// Time-stamp: <2026-04-30 17h19 BST (9561bc0c)>

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::Case { scrutinee, arms } => {
                if arms.is_empty() {
                    return Err(CheckerError::Unimplemented(
                        "mapping out of bottom types".to_string(),
                    ));
                }

                let arm_owners = arms
                    .iter()
                    .map(|ca| self.lookup_variant_field(&ca.tag).map(|sf| &sf.owner))
                    .collect::<Result<Vec<_>, _>>()?;
                let owner = arm_owners[0];

                if !arm_owners.into_iter().all(|o| self.equal(owner, o)) {
                    todo!("inconsistent case scrutinee set");
                }

                let Set::Variant(fields) = owner else {
                    panic!(
                        "invariant violation: looking up the owner of a variant field resulted in a non-variant set",
                    )
                };
                let mut required_field_names_sorted: Vec<String> =
                    fields.iter().map(|vf| vf.name.clone()).collect();
                required_field_names_sorted.sort();
                let mut covered_field_names_sorted: Vec<String> =
                    arms.iter().map(|ca| ca.tag.clone()).collect();
                covered_field_names_sorted.sort();
                if required_field_names_sorted != covered_field_names_sorted {
                    return Err(CheckerError::IncompleteCaseAnalysis {
                        found: covered_field_names_sorted,
                        required: required_field_names_sorted,
                    });
                }
                let scrutinee = self.check_element(scrutinee, owner)?;

                let matching: Option<(String, Element)> = match scrutinee {
                    Element::Inject {
                        ref field,
                        element: ref inner,
                    } => Some((field.clone(), *inner.clone())),
                    Element::Var(_) | Element::Project { .. } | Element::Case { .. } => None,
                    Element::Literal(_) | Element::Record(_) => panic!(
                        "invariant violation: scrutinee is a non-variant value at variant set"
                    ),
                };

                let mut computed_output = None;
                let mut processed_arms = Vec::new();
                for arm in arms {
                    let Field {
                        field: field_set, ..
                    } = self.lookup_variant_field(&arm.tag)?;

                    let mut ctx = self.clone();
                    let binding_name = arm.bound.clone();
                    let binding_set = field_set.clone();

                    let case_arm = if let Some((tag, inner)) = &matching
                        && *tag == arm.tag
                    {
                        let canonical =
                            ctx.make_element_definition(binding_name, inner.clone(), binding_set)?;
                        let output = ctx.check_instance((&arm.body).into(), signature)?;
                        if matches!(computed_output, Some(_)) {
                            panic!(
                                "invariant violation: we somehow matched multiple arms in case analysis"
                            )
                        }
                        computed_output = Some(output.clone());
                        InstCaseArm {
                            tag: arm.tag.clone(),
                            bound: canonical,
                            body: output,
                        }
                    } else {
                        let canonical = ctx.make_element_binding(binding_name, binding_set)?;
                        let body = ctx.check_instance((&arm.body).into(), signature)?;
                        InstCaseArm {
                            tag: arm.tag.clone(),
                            bound: canonical,
                            body,
                        }
                    };

                    processed_arms.push(case_arm);
                }
                Ok(computed_output.unwrap_or(Instance::Case {
                    scrutinee: Box::new(scrutinee),
                    arms: processed_arms,
                }))
            }
            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 {
                    let mut ctx = self.clone();
                    let inst_params = inst_params
                        .iter()
                        .map(|Param { name, set }| {
                            let set = ctx.check_set(set)?;
                            let canon = ctx.make_element_binding(name.clone(), set.clone())?;
                            Ok(Param {
                                name: canon,
                                set: set,
                            })
                        })
                        .collect::<Result<Vec<_>, _>>()?;
                    let body = ctx.check_instance(body, None)?;
                    Ok(Instance::For {
                        params: inst_params,
                        body: Box::new(body),
                    })
                }
            }
            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 parameters, but here we are"
                            );
                        }
                        let mut ctx = self.clone();
                        let first_set = ctx.check_set(&params[0].set)?;
                        let element_checked = self.check_element(element, &first_set)?;
                        ctx.make_element_definition(
                            params[0].name.clone(),
                            element_checked,
                            first_set,
                        )?;

                        if params.len() == 1 {
                            ctx.check_instance(body, signature)
                        } else {
                            let residual = Instance::For {
                                params: params[1..].to_vec(),
                                body: body.clone(),
                            };
                            ctx.check_instance(&residual, signature)
                        }
                    }
                    Instance::Record(_) | Instance::SetCoerce(_) => {
                        Err(CheckerError::NonFunctionalInstance {
                            instance: inner,
                            element: *element.clone(),
                        })
                    }
                    Instance::Project { .. } | Instance::App(_, _) | Instance::Case { .. } => {
                        // 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()))
                    }
                }
            }
        }
    }
}