use crate::ast::*;
use crate::checker_state::*;
use std::collections::HashMap;
use std::iter::zip;

use tracing::instrument;

impl CheckerState {
    #[instrument(skip(self), level = "debug", fields(%signature))]
    pub fn check_signature(&self, signature: &Signature) -> Result<Signature, CheckerError> {
        match signature {
            Signature::Set => Ok(Signature::Set),
            Signature::Var(v) => {
                let deref = self.lookup_signature(&v)?;
                Ok(deref.clone())
            }
            Signature::Ext { params, codomain } => {
                let mut ctx = self.clone();
                let params = params
                    .iter()
                    .map(|p| {
                        let set = ctx.check_set(&p.set)?;
                        let canon = ctx.make_element_binding(p.name.clone(), set.clone())?;
                        Ok(Param { set, name: canon })
                    })
                    .collect::<Result<Vec<_>, _>>()?;
                let codomain = Box::new(ctx.check_signature(codomain)?);
                Ok(Signature::Ext { params, codomain })
            }
            Signature::Theory(fields) => {
                let mut ctx = self.clone();
                let temp_name = ctx.make_unique_name();
                let mut new_fields = Vec::new();
                for SigField { signature, name } in fields {
                    let signature = ctx.check_signature(signature)?;
                    ctx.add_instance(name.clone(), InstanceValue::Hypothetical, signature.clone())?;
                    // And lo, the special case, our chosen canonical form
                    if signature == Signature::Set {
                        ctx.add_set(
                            name.clone(),
                            Set::ClaimedSet(Instance::Var(name.clone())).into(),
                        )?;
                    }
                    new_fields.push(SigField {
                        name: name.clone(),
                        signature,
                    });
                    // We must iteratively add the entire signature so that
                    // field lookup does something, as we rely on that for type
                    // checking. We could hack together a signature i suppose,
                    // but the cleanest thing is to add the truncations of this
                    // signature. In any event the context is discarded
                    // afterward.
                    ctx.add_signature(&temp_name, Signature::Theory(new_fields.clone()), true)?;
                }
                Ok(Signature::Theory(new_fields))
            }
        }
    }

    #[instrument(skip(self), level = "debug", fields(%instance, ?signature))]
    pub fn check_instance(
        &self,
        instance: &Instance,
        signature: Option<&Signature>,
    ) -> Result<Instance, CheckerError> {
        match instance {
            Instance::SetCoerce(set) => {
                if let Some(signature) = signature
                    && *signature != Signature::Set
                {
                    return Err(CheckerError::WrongSignatureForInstance {
                        value: instance.clone().into(),
                        real: Signature::Set,
                        claimed: signature.clone(),
                    });
                };
                let set = Box::new(self.check_set(set)?);
                if let Set::ClaimedSet(inner) = *set {
                    Ok(inner)
                } else {
                    Ok(Instance::SetCoerce(set))
                }
            }
            Instance::Var(v) => {
                // Exactly the same discipline as for Element::Var, see there
                // for some sparse comments
                let lookup = self.lookup_instance(&v)?;
                if let Some(signature) = signature
                    && !self.equal(signature, &lookup.container)
                {
                    return Err(CheckerError::WrongSignatureForInstance {
                        value: instance.clone().into(),
                        claimed: signature.clone(),
                        real: lookup.container.clone(),
                    });
                }
                if let InstanceValue::Concrete(ref deref) = lookup.value {
                    Ok(deref.clone())
                } else {
                    Ok(Instance::Var(v.clone()))
                }
            }

            Instance::Record(assignations) => {
                if let Some(signature) = signature {
                    let rej = |reason| CheckerError::InstanceDoesNotBelong {
                        instance: instance.clone(),
                        claimed: signature.clone(),
                        reason,
                    };

                    let fields = if let Signature::Theory(fields) = signature {
                        Ok(fields)
                    } else {
                        Err(rej("signature has no fields".to_string()))
                    }?;

                    let mut signature_fnames_sorted: Vec<String> =
                        fields.iter().map(|x| x.name.clone()).collect();
                    signature_fnames_sorted.sort();

                    let mut instance_fnames_sorted: Vec<String> =
                        assignations.iter().map(|x| x.name.clone()).collect();
                    instance_fnames_sorted.sort();

                    if signature_fnames_sorted != instance_fnames_sorted {
                        return Err(rej(format!(
                            "expected [{}] but found [{}]",
                            signature_fnames_sorted.join(", "),
                            instance_fnames_sorted.join(", "),
                        )));
                    }

                    let assignations = assignations
                        .iter()
                        .map(|x| (&x.name, &x.instance))
                        .collect::<HashMap<_, _>>();

                    let mut ctx = self.clone();
                    let sub_instances = fields
                        .iter()
                        .map(
                            |SigField {
                                 name: f_n,
                                 signature: f_s,
                             }| {
                                let f_i = assignations
                                    .get(f_n)
                                    .expect("we have already checked that all fields are present");
                                let f_s = ctx.check_signature(f_s)?;
                                let f_i = ctx.check_instance(f_i, Some(&f_s))?;

                                if f_s == Signature::Set {
                                    ctx.add_set(
                                        f_n.clone(),
                                        Set::ClaimedSet(Instance::Var(f_n.clone())).into(),
                                    )?;
                                }

                                ctx.add_instance(f_n.clone(), f_i.clone().into(), f_s.clone())?;
                                Ok(InstAssign {
                                    name: f_n.clone(),
                                    instance: f_i,
                                })
                            },
                        )
                        .collect::<Result<Vec<_>, _>>()?;

                    Ok(Instance::Record(sub_instances))
                } else {
                    let sub_insts = assignations
                        .iter()
                        .map(|InstAssign { name, instance }| {
                            Ok(InstAssign {
                                name: name.clone(),
                                instance: self.check_instance(instance, None)?,
                            })
                        })
                        .collect::<Result<Vec<_>, _>>()?;
                    Ok(Instance::Record(sub_insts))
                }
            }
            Instance::Project { instance, field } => {
                let Field {
                    field: field_signature,
                    owner: owner_signature,
                } = self.lookup_signature_field(&field)?;

                if let Some(signature) = signature
                    && !self.equal(signature, field_signature)
                {
                    return Err(CheckerError::WrongSignatureForInstance {
                        value: (*instance.clone()).into(),
                        claimed: signature.clone(),
                        real: field_signature.clone(),
                    });
                }

                let inner = self.check_instance(instance, owner_signature.into())?;

                match inner {
                    Instance::Var(_) | Instance::Project { .. } => Ok(Instance::Project {
                        instance: Box::new(inner),
                        field: field.clone(),
                    }),
                    Instance::Record(assignations) => {
                        let sub_element = assignations
                            .into_iter()
                            .find(|a| a.name == *field)
                            .expect(
                                "invariant violation: record missing field that was type-checked",
                            )
                            .instance;
                        Ok(sub_element)
                    }
                    _ => panic!(
                        "invariant violation: check_instance returned neither a record or stuck computation for record set"
                    ),
                }
            }
            Instance::For {
                params: inst_params,
                body,
            } => {
                if let Some(signature) = signature {
                    if inst_params.is_empty() {
                        todo!("should be impossible");
                    };
                    let Signature::Ext {
                        params: sig_params,
                        codomain,
                    } = signature
                    else {
                        todo!("need to raise error");
                    };
                    if sig_params.is_empty() {
                        todo!("should be impossible")
                    }
                    if sig_params.len() != inst_params.len() {
                        todo!("this is a type error")
                    }
                    let mut ctx = self.clone();
                    let inst_params = zip(inst_params, sig_params)
                        .map(
                            |(
                                Param {
                                    name: inst_n,
                                    set: inst_s,
                                },
                                Param {
                                    name: set_n,
                                    set: set_s,
                                },
                            )| {
                                let inst_s = ctx.check_set(inst_s)?;
                                let set_s = ctx.check_set(set_s)?;
                                if !ctx.equal(&inst_s, &set_s) {
                                    todo!("type error")
                                }
                                ctx.add_element(
                                    inst_n.clone(),
                                    Element::Var(set_n.clone()).into(),
                                    set_s.clone(),
                                )?;
                                Ok(Param {
                                    name: set_n.clone(),
                                    set: set_s,
                                })
                            },
                        )
                        .collect::<Result<Vec<_>, _>>()?;
                    let body = ctx.check_instance(body, Some(&*codomain))?;
                    Ok(Instance::For {
                        body: Box::new(body),
                        params: inst_params,
                    })
                } else {
                    todo!("expand only mode???")
                }
            }
            Instance::App(inner, element) => {
                // this is the only time that we ever call check_instance with
                // signature = None, and in this mode all we want is to put
                // inner into a canonical form pushing stuck terms to the leaves
                // and simplifying everything else.
                let inner = self.check_instance(inner, None)?;
                match &inner {
                    Instance::Var(v) => {
                        let field = self.lookup_signature_field(&v)?;
                        let Signature::Ext {
                            ref params,
                            ref codomain,
                        } = field.field
                        else {
                            return Err(CheckerError::WrongSignatureForInstance {
                                value: inner.clone().into(),
                                claimed: Signature::Ext {
                                    params: vec![Param {
                                        name: "...".to_string(),
                                        set: Set::Var("...".to_string()),
                                    }],
                                    codomain: Box::new(Signature::Var("...".to_string())),
                                },
                                real: field.owner.clone(),
                            });
                        };
                        if params.is_empty() {
                            panic!(
                                "It should have been impossible to construct an Ext with no params, but here we are"
                            );
                        }
                        if let Some(signature) = signature {
                            if !self.equal(&**codomain, signature) {
                                return Err(CheckerError::WrongSignatureForInstance {
                                    value: instance.clone().into(),
                                    claimed: signature.clone(),
                                    real: *codomain.clone(),
                                });
                            }
                        }
                        let element = self.check_element(element, &params[0].set)?;
                        Ok(Instance::App(
                            Box::new(Instance::Var(v.clone())),
                            Box::new(element),
                        ))
                    }
                    Instance::For { params, body } => {
                        // if params.is_empty() {
                        //     panic!(
                        //         "It should have been impossible to construct a For with no params, but here we are"
                        //     );
                        // }
                        // if let Some(signature) = signature {
                        //     let Signature::Ext {
                        //         params: sig_params,
                        //         codomain,
                        //     } = signature
                        //     else {
                        //         return Err(CheckerError::WrongSignatureForInstance {
                        //             value: inner.clone().into(),
                        //             claimed: signature.clone(),
                        //             real: Signature::Ext {
                        //                 params: vec![Param {
                        //                     name: "...".to_string(),
                        //                     set: Set::Var("...".to_string()),
                        //                 }],
                        //                 codomain: Box::new(Signature::Var("...".to_string())),
                        //             },
                        //         });
                        //     };
                        //     if params.len() != sig_params.len() {
                        //         todo!("raise an error about incorrect signature")
                        //     };
                        //     let mut ctx = self.clone();
                        //     let (params, sig_params): (Vec<Param>, Vec<Param>) =
                        //         zip(params, sig_params)
                        //             .enumerate()
                        //             .map(|(idx, (a_p, s_p))| {
                        //                 let a_set = ctx.check_set(&a_p.set)?;
                        //                 let p_set = ctx.check_set(&s_p.set)?;
                        //                 if !ctx.equal(&a_set, &p_set) {
                        //                     todo!(
                        //                         "raise an error about the for/Ext being incorrect"
                        //                     );
                        //                 }
                        //                 if idx == 0 {
                        //                     let element = ctx.check_element(&*element, &p_set)?;
                        //                     ctx.make_element_definition(
                        //                         a_p.name.clone(),
                        //                         element.clone(),
                        //                         a_set.clone(),
                        //                     )?;
                        //                     ctx.make_element_definition(
                        //                         s_p.name.clone(),
                        //                         element,
                        //                         p_set.clone(),
                        //                     )?;
                        //                 } else {
                        //                     let canon_a = ctx.make_element_binding(
                        //                         a_p.name.clone(),
                        //                         a_set.clone(),
                        //                     )?;
                        //                     ctx.add_element(
                        //                         s_p.name.clone(),
                        //                         Element::Var(canon_a).into(),
                        //                         p_set.clone(),
                        //                     )?;
                        //                 }
                        //                 Ok((
                        //                     Param {
                        //                         name: a_p.name.clone(),
                        //                         set: a_set,
                        //                     },
                        //                     Param {
                        //                         name: s_p.name.clone(),
                        //                         set: p_set,
                        //                     },
                        //                 ))
                        //             })
                        //             .collect::<Result<_, _>>()?;
                        //     // TODO: now we need to build a new For and a new
                        //     // Ext, and check the latter then check that the
                        //     // former fits the latter. The base case would be
                        //     // checking that body's expected signature is
                        //     // codomain.
                        // }
                        todo!("finish this");
                    }
                    Instance::Record(_) | Instance::SetCoerce(_) => {
                        Err(CheckerError::NonFunctionalInstance {
                            instance: inner,
                            element: *element.clone(),
                        })
                    }
                    Instance::Project { .. } | Instance::App(_, _) => {
                        // it would appear that we are stuck here, so our only
                        // choice is to continue to be so
                        Ok(Instance::App(Box::new(inner), element.clone()))
                    }
                }
            }
        }
    }
}