use crate::ast::*;
use crate::checker_state::*;
use std::collections::{HashMap, HashSet};
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::FromSet(set) => Ok(Signature::FromSet(self.check_set(set)?)),
            Signature::Var(v) => {
                let deref = self.lookup_signature(&v)?;
                Ok(deref.clone())
            }
            Signature::Ext { params, codomain } => {
                if params.is_empty() {
                    return Err(CheckerError::Unimplemented(
                        "extension signature with empty params".to_string(),
                    ));
                }
                let mut ctx = self.clone();

                // a little juggling here to ensure that we merge names without
                // collision, the pattern is: bind to something unique, recurse,
                // fix the names which involves in particular messing about with
                // counters
                let params = params
                    .iter()
                    .map(|p| {
                        let set = ctx.check_set(&p.set)?;
                        let canon = ctx.make_unique_name();
                        ctx.add_element(
                            p.name.clone(),
                            Element::Var(canon.clone()).into(),
                            set.clone(),
                        )?;
                        ctx.add_element(canon.clone(), ElementValue::Hypothetical, set.clone())?;
                        Ok(Param { set, name: canon })
                    })
                    .collect::<Result<Vec<_>, _>>()?;

                // use tactic "trust_me"
                let snapshot = ctx.get_binder_counter();
                let codomain = ctx.check_signature(codomain)?;
                ctx.set_binder_counter(snapshot);

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

                // normalise
                let mut ctx = self.clone();
                let params = params
                    .into_iter()
                    .map(|Param { name, set }| {
                        let canonical = ctx.make_element_binding(name, set.clone())?;
                        let set = ctx.check_set(&set)?;
                        Ok(Param {
                            name: canonical,
                            set,
                        })
                    })
                    .collect::<Result<Vec<_>, _>>()?;
                let codomain = ctx.check_signature(&codomain)?;

                Ok(Signature::Ext {
                    params,
                    codomain: Box::new(codomain),
                })
            }
            Signature::Theory(fields) => {
                let field_set = fields.iter().map(|f| &f.name).collect::<HashSet<&String>>();
                if field_set.len() != fields.len() {
                    return Err(CheckerError::DuplicateFieldsSignature(signature.clone()));
                }

                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.recursively_add_hypothetical_instance(
                        name.clone(),
                        signature.clone(),
                        None,
                    )?;
                    new_fields.push(Field {
                        name: name.clone(),
                        carries: signature,
                    });
                    ctx.add_signature(&temp_name, Signature::Theory(new_fields.clone()), true)?;
                }
                Ok(Signature::Theory(new_fields))
            }
        }
    }

    #[instrument(skip(self), level = "debug", fields(%name, %signature, head=%head.map(|s| s.to_string()).unwrap_or_default()) )]
    fn recursively_add_hypothetical_instance(
        &mut self,
        name: String,
        signature: Signature,
        head: Option<&Instance>,
    ) -> Result<(), CheckerError> {
        let value = match head {
            Some(h) => InstanceValue::Concrete(Instance::Project {
                instance: Box::new(h.clone()),
                field: name.clone(),
            }),
            None => InstanceValue::Hypothetical,
        };
        let self_instance = match head {
            Some(h) => Instance::Project {
                instance: Box::new(h.clone()),
                field: name.clone(),
            },
            None => Instance::Var(name.clone()),
        };
        // handle the special case canonical form for _ :: Set
        if signature == Signature::Set {
            self.add_set(name.clone(), Set::ClaimedSet(self_instance.clone()), false)?;
        }
        self.add_instance(name.clone(), value, signature.clone())?;
        if let Signature::Theory(fields) = signature {
            for f in fields {
                self.recursively_add_hypothetical_instance(
                    f.name,
                    f.carries,
                    Some(&self_instance),
                )?;
            }
        }
        Ok(())
    }

    #[instrument(skip(self), level = "debug", fields(%instance, signature=%signature.map(|s| s.to_string()).unwrap_or_default()) )]
    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::ElementCoerce(element) => {
                let set = if let Some(signature) = signature {
                    let Signature::FromSet(set) = signature else {
                        return Err(CheckerError::WrongSignatureForInstance {
                            value: instance.clone().into(),
                            real: Signature::FromSet(Set::Var("_".to_string())),
                            claimed: signature.clone(),
                        });
                    };
                    Some(set)
                } else {
                    None
                };

                let element = self.check_element(element, set)?;
                Ok(Instance::ElementCoerce(element))
            }
            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())),
                                        false,
                                    )?;
                                }

                                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)?;

                let field_signature = self.check_signature(field_signature)?;

                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)
                    }
                    Instance::For { .. }
                    | Instance::ElementCoerce(_)
                    | Instance::SetCoerce(_)
                    | Instance::App { .. }
                    | Instance::Case { .. } => 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, Some(owner))?;

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

                let posit_equality_with = match &scrutinee {
                    Element::Var(v) => Some(v.clone()),
                    Element::ClaimedElement(_)
                    | Element::Inject { .. }
                    | Element::Project { .. }
                    | Element::Case { .. }
                    | Element::Literal(_)
                    | Element::Record(_) => None,
                };

                let mut computed_output = None;
                let mut processed_arms = Vec::new();
                for arm in arms {
                    let OwnedField {
                        field: field_set,
                        owner,
                    } = 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 this_signature =
                            signature.map(|s| ctx.check_signature(s)).transpose()?;

                        let output =
                            ctx.check_instance((&arm.body).into(), this_signature.as_ref())?;
                        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)?;
                        if let Some(ref scrutinee_var) = posit_equality_with {
                            ctx.add_element(
                                scrutinee_var.clone(),
                                Element::Inject {
                                    field: arm.tag.clone(),
                                    element: Box::new(Element::Var(canonical.clone())),
                                }
                                .into(),
                                owner.clone(),
                            )?;
                        };
                        let this_signature =
                            signature.map(|s| ctx.check_signature(s)).transpose()?;
                        let body =
                            ctx.check_instance((&arm.body).into(), this_signature.as_ref())?;
                        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();
                let inst_params: Vec<Param> = 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 })
                    })
                    .collect::<Result<_, _>>()?;

                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(),
                                )?;
                                // TODO should this be unconditional?
                                ctx.add_element(
                                    set_n.clone(),
                                    ElementValue::Hypothetical,
                                    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 {
                    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(_) | Instance::ElementCoerce(_) => {
                        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, Some(&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()))
            }
            // until we properly support motives there's nothing we can really do here
            Instance::Case { .. } => {
                let msg = format!(
                    "this type checker has no motives yet, and was called upon to infer the signature of {inst}, which leads with a `case`, and so has no option but to fail"
                );
                Err(CheckerError::Unimplemented(msg))
            }
            Instance::ElementCoerce(_)
            | Instance::For { .. }
            | Instance::Record(_)
            | Instance::SetCoerce(_) => Ok(None),

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