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
                    .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) => {
                // once again, mutatis mutandis from elements
                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();

                let signature_fsigs: Vec<Option<Signature>>;
                if let Some(signature) = signature {
                    let rej = |reason| CheckerError::InstanceDoesNotBelong {
                        instance: instance.clone(),
                        claimed: signature.clone(),
                        reason: reason,
                    };

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

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

                    if signature_fnames_sorted != instance_fnames_sorted {
                        return Err(rej(format!(
                            "expected [{}] but found [{}]",
                            signature_fnames.join(", "),
                            instance_fnames.join(", "),
                        )));
                    }
                } else {
                    signature_fsigs = std::iter::repeat(None)
                        .take(instance_finstances.len())
                        .collect();
                }

                let sub_els = zip(instance_finstances, signature_fsigs)
                    .map(|(e_f, e_s)| self.check_instance(e_f, (&e_s).into()))
                    .collect::<Result<Vec<_>, _>>()?;
                let assignations = zip(instance_fnames, sub_els)
                    .map(|(name, instance)| InstAssign { name, instance })
                    .collect();
                Ok(Instance::Record(assignations))
            }
            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 { .. } => {
                todo!("instance for")
            }
            Instance::App(inner, element) => {
                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"
                            );
                        }
                        let mut ctx = self.clone();
                        for (idx, param) in params.iter().enumerate() {
                            let Param {
                                name: param_name,
                                set: param_set,
                            } = param.clone();
                            if idx == 0 {
                                let element = self.check_element(&*element, &param_set)?;
                                ctx.make_element_definition(param_name, element, param_set)?;
                            } else {
                                ctx.make_element_binding(param_name, param_set)?;
                            }
                        }
                        let instance = ctx.check_instance(body, signature)?;
                        if params.len() == 1 {
                            Ok(instance)
                        } else {
                            Ok(Instance::For {
                                params: params[1..].iter().map(|p| p.clone()).collect(),
                                body: Box::new(instance),
                            })
                        }
                    }
                    Instance::Record(_) | Instance::SetCoerce(_) => {
                        Err(CheckerError::NonFunctionalInstance {
                            instance: inner,
                            element: *element.clone(),
                        })
                    }
                    Instance::Project { .. } | Instance::App(_, _) => panic!(
                        "invariant violation: we did not completely expand the inner instance in our app"
                    ),
                }
            }
        }
    }
}