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 = ctx.check_signature(codomain)?;

                let result = if let Signature::Ext {
                    params: inner,
                    codomain: deep,
                } = codomain
                {
                    let mut merged = params.clone();
                    merged.extend(inner.iter().cloned());
                    Signature::Ext {
                        params: merged,
                        codomain: deep,
                    }
                } else {
                    Signature::Ext {
                        params,
                        codomain: Box::new(codomain),
                    }
                };

                Ok(result)
            }
            Signature::Theory(fields) => {
                let mut ctx = self.clone();
                let temp_name = ctx.make_unique_name();
                let mut new_fields = Vec::new();
                for Field { carries, name } in fields {
                    let signature = ctx.check_signature(carries)?;
                    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())))?;
                    }
                    new_fields.push(Field {
                        name: name.clone(),
                        carries: 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(
                            |Field {
                                 name: f_n,
                                 carries: 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())),
                                    )?;
                                }

                                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 OwnedField {
                    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)) {
                    return Err(CheckerError::InstanceInconsistentCaseScrutineeSet(
                        instance.clone(),
                    ));
                }

                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 OwnedField {
                        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());
                        CaseArm {
                            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)?;
                        CaseArm {
                            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 inst_params.is_empty() {
                    return Err(CheckerError::Unimplemented(
                        "for instance with empty params".to_string(),
                    ));
                }

                // deal with left-nesting
                let mut ctx = self.clone();
                for Param { name, set } in inst_params {
                    ctx.make_element_binding(name.clone(), set.clone())?;
                }
                let body = ctx.check_instance(body, None)?;
                let (inst_params, body) = if let Instance::For {
                    params: inner_params,
                    body: inner_body,
                } = body
                {
                    let mut merged = inst_params.clone();
                    merged.extend(inner_params.clone());
                    (merged, *inner_body)
                } else {
                    (inst_params.clone(), body)
                };

                if let Some(signature) = signature {
                    if inst_params.is_empty() {
                        return Err(CheckerError::Unimplemented(
                            "instance for with empty params".to_string(),
                        ));
                    };

                    let Signature::Ext {
                        params: sig_params,
                        codomain,
                    } = signature
                    else {
                        return Err(CheckerError::WrongSignatureForInstance {
                            real: Signature::Ext {
                                params: vec![Param {
                                    name: "_".to_string(),
                                    set: Set::Var("_".to_string()),
                                }],
                                codomain: Box::new(Signature::Var("_".to_string())),
                            },
                            value: instance.clone().into(),
                            claimed: signature.clone(),
                        });
                    };

                    if sig_params.is_empty() {
                        return Err(CheckerError::Unimplemented(
                            "extension signature with empty params".to_string(),
                        ));
                    }

                    if sig_params.len() != inst_params.len() {
                        return Err(CheckerError::InstanceDoesNotBelong {
                            instance: instance.clone(),
                            reason: "instance and signature have differing number of parameters"
                                .to_string(),
                            claimed: signature.clone(),
                        });
                    }
                    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) {
                                    return Err(CheckerError::InstanceDoesNotBelong{
                                        instance: instance.clone(),
                                        claimed: signature.clone(),
                                        reason: format!("the signature specifies set {} but in this position the instance has set {}",
                                            set_s, inst_s),
                                    });
                                }
                                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 {
                instance: inner,
                args,
            } => {
                let subject = self.check_instance(inner, None)?;
                // the whole game here is to make sure that we have no left
                // nesting, and that we're fully evaluated. If that's true then
                // structural equality is much more powerful, and partial
                // application is simpler.
                let (subject, args) = match subject {
                    Instance::App {
                        instance: inner_inner,
                        args: inner_args,
                    } => {
                        let mut merged = inner_args;
                        merged.extend(args.iter().cloned());
                        return self.check_instance(
                            &Instance::App {
                                instance: inner_inner,
                                args: merged,
                            },
                            signature,
                        );
                    }
                    other => (other, args.clone()),
                };

                // nevertheless we need the tiniest amount of bidirectionality
                // here to deal with case, project, and var recursively
                let subject_sig: Option<Signature> = self.stuck_subject_signature(&subject)?;

                if let Some(subject_sig) = subject_sig {
                    let Signature::Ext { params, codomain } = subject_sig else {
                        return Err(CheckerError::NonFunctionalInstance {
                            instance: subject,
                            elements: args,
                        });
                    };
                    let (ctx, checked_args) = self._bind_args(instance, &params, &args)?;

                    if let Some(expected) = signature {
                        let result_sig = if args.len() == params.len() {
                            ctx.check_signature(&codomain)?
                        } else {
                            let remaining: Vec<Param> = params[args.len()..]
                                .iter()
                                .map(|p| {
                                    let s = ctx.check_set(&p.set)?;
                                    Ok(Param {
                                        name: p.name.clone(),
                                        set: s,
                                    })
                                })
                                .collect::<Result<_, CheckerError>>()?;
                            let cod = ctx.check_signature(&codomain)?;
                            Signature::Ext {
                                params: remaining,
                                codomain: Box::new(cod),
                            }
                        };
                        if !self.equal(expected, &result_sig) {
                            return Err(CheckerError::WrongSignatureForInstance {
                                value: subject.clone().into(),
                                claimed: expected.clone(),
                                real: result_sig,
                            });
                        }
                    }

                    return Ok(Instance::App {
                        instance: Box::new(subject),
                        args: checked_args,
                    });
                }

                match subject {
                    Instance::For { params, body } => {
                        let (ctx, _checked) = self._bind_args(instance, &params, &args)?;
                        if args.len() == params.len() {
                            ctx.check_instance(&body, signature)
                        } else {
                            let residual = Instance::For {
                                params: params[args.len()..].to_vec(),
                                body,
                            };
                            ctx.check_instance(&residual, signature)
                        }
                    }
                    Instance::Record(_) | Instance::SetCoerce(_) => {
                        Err(CheckerError::NonFunctionalInstance {
                            instance: subject,
                            elements: args,
                        })
                    }
                    Instance::Case { .. } => {
                        unreachable!("_stuck_subject_signature should have dealt with this")
                    }
                    Instance::Var(_) | Instance::Project { .. } => {
                        unreachable!("handled in the stuck-head branch above")
                    }
                    Instance::App { .. } => {
                        unreachable!("handled by the merge-and-recurse branch above")
                    }
                }
            }
        }
    }

    fn _bind_args(
        &self,
        instance: &Instance,
        params: &[Param],
        args: &[Element],
    ) -> Result<(CheckerState, Vec<Element>), CheckerError> {
        if args.len() > params.len() {
            return Err(CheckerError::OverApplication {
                instance: instance.clone(),
                applied_to: args.len(),
            });
        }
        let mut ctx = self.clone();
        let checked = zip(params.iter(), args.iter())
            .map(|(p, a)| {
                let p_set = ctx.check_set(&p.set)?;
                let a = ctx.check_element(a, &p_set)?;
                ctx.add_element(p.name.clone(), a.clone().into(), p_set)?;
                Ok(a)
            })
            .collect::<Result<Vec<_>, _>>()?;
        Ok((ctx, checked))
    }

    fn stuck_subject_signature(&self, inst: &Instance) -> Result<Option<Signature>, CheckerError> {
        match inst {
            Instance::Var(v) => Ok(Some(self.lookup_instance(v)?.container.clone())),
            Instance::Project { field, .. } => {
                Ok(Some(self.lookup_signature_field(field)?.field.clone()))
            }
            // All arms of a stuck Case share a signature by the case
            // elimination typing rule, and we've already expanded the body, so
            // we can pick any arm.
            Instance::Case { arms, .. } => self
                .stuck_subject_signature(&arms.first().expect("we don't allow bottom type").body),

            Instance::For { .. } | Instance::Record(_) | Instance::SetCoerce(_) => Ok(None),

            Instance::App { .. } => unreachable!(
                "invariant violation: _stuck_head_signature called on a left-nested App"
            ),
        }
    }
}