use crate::ast::*;
use crate::checker_state::*;
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
                    .into_iter()
                    .map(|p| {
                        let set = ctx.check_set(p.set.clone())?;
                        ctx.make_element_binding(p.name.clone(), set.clone())?;
                        Ok(Param { set, name: p.name })
                    })
                    .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 fields = fields
                    .into_iter()
                    .map(|SigField { signature, name }| {
                        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(),
                            )?;
                        }
                        Ok(SigField { name, signature })
                    })
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(Signature::Theory(fields))
            }
        }
    }

    #[instrument(skip(self), level = "debug", fields(%instance, ?signature))]
    pub fn check_instance(
        &self,
        instance: Instance,
        signature: &Signature,
    ) -> Result<Instance, CheckerError> {
        match instance {
            Instance::SetCoerce(ref set) => {
                if *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.clone())?);
                if let Set::ClaimedSet(inner) = *set {
                    Ok(inner)
                } else {
                    Ok(Instance::SetCoerce(set))
                }
            }
            Instance::Var(ref v) => {
                // Exactly the same discipline as for Element::Var, see there
                // for some sparse comments
                let lookup = self.lookup_instance(&v)?;
                if !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(ref assignations) => {
                // once again, mutatis mutandis from elements
                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("instance is not a record instance".to_string()))
                }?;

                let (signature_fnames, signature_fsigs): (Vec<String>, Vec<Signature>) = fields
                    .iter()
                    .map(|SigField { name, signature }| (name.clone(), signature.clone()))
                    .unzip();
                let mut signature_fnames_sorted = signature_fnames.clone();
                signature_fnames_sorted.sort();

                let (instance_fnames, instance_finstances): (Vec<String>, Vec<&Instance>) =
                    assignations
                        .iter()
                        .map(|InstAssign { name, instance }| (name.clone(), instance))
                        .unzip();

                let mut instance_fnames_sorted = instance_fnames.clone();
                instance_fnames_sorted.sort();

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

                let sub_els = zip(instance_finstances, signature_fsigs)
                    .map(|(e_f, e_s)| self.check_instance(e_f.clone(), &e_s))
                    .collect::<Result<Vec<_>, _>>()?;
                let assignations = zip(instance_fnames, sub_els)
                    .map(|(name, instance)| InstAssign { name, instance })
                    .collect();
                Ok(Instance::Record(assignations))
            }
            Instance::Project {
                ref instance,
                ref field,
            } => {
                let Field {
                    field: field_signature,
                    owner: owner_signature,
                } = self.lookup_signature_field(&field)?;

                if !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.clone(), owner_signature)?;

                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, body } => {
                Err(CheckerError::Unimplemented("instance for".to_string()))
            }
            Instance::App(inst, elem) => {
                println!("{}", self);
                Err(CheckerError::Unimplemented("instance app".to_string()))
            }
        }
    }
}