use crate::ast::*;
use crate::checker_state::*;

use std::collections::{HashMap, HashSet};
use tracing::instrument;

impl CheckerState {
    #[instrument(skip(self), level = "debug", fields(%set))]
    pub fn check_set(&self, set: &Set) -> Result<Set, CheckerError> {
        match set {
            Set::BuiltIn(_) => Ok(set.clone()),
            Set::Record(fields) => {
                let field_set = fields.iter().map(|f| &f.name).collect::<HashSet<&String>>();
                if field_set.len() != fields.len() {
                    return Err(CheckerError::DuplicateFieldsSet(set.clone()));
                }

                let mut ctx = self.clone();
                let temp_name = ctx.make_unique_name();
                let mut new_fields = Vec::new();
                for Field { name, carries } in fields {
                    let set = ctx.check_set(carries)?;
                    ctx._recursively_add_hypothetical_element(name.clone(), set.clone(), None)?;
                    new_fields.push(Field {
                        name: name.clone(),
                        carries: set,
                    });
                    // 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_set(temp_name.clone(), Set::Record(new_fields.clone()), true)?;
                }
                Ok(Set::Record(new_fields))
            }
            Set::Variant(fields) => {
                let field_set = fields.iter().map(|f| &f.name).collect::<HashSet<&String>>();
                if field_set.len() != fields.len() {
                    return Err(CheckerError::DuplicateFieldsSet(set.clone()));
                }

                let fields = fields
                    .into_iter()
                    .map(|Field { name, carries }| {
                        let set = self.check_set(carries)?;
                        Ok(Field {
                            name: name.clone(),
                            carries: set,
                        })
                    })
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(Set::Variant(fields))
            }
            Set::ClaimedSet(instance) => {
                let instance = self.check_instance(instance, Some(&Signature::Set))?;
                if let Instance::SetCoerce(set) = instance {
                    Ok(*set)
                } else {
                    Ok(Set::ClaimedSet(instance))
                }
            }
            Set::Var(v) => {
                let deref = self.lookup_set(&v)?;
                match deref {
                    SetValue::Hypothetical => Ok(Set::Var(v.clone())),
                    SetValue::Concrete(deref) => Ok(deref.clone()),
                }
            }
        }
    }

    // the goal here is to spread the love: if we are adding a hypothetical of
    // some set _ : record { ... } then we must recurse into all of those
    // fields and add hypotheticals for them---but, we need to build the tree as
    // we go, giving them the concrete value of their path from the root (our
    // canonical form).
    fn _recursively_add_hypothetical_element(
        &mut self,
        name: String,
        set: Set,
        head: Option<&Element>,
    ) -> Result<(), CheckerError> {
        let value = match head {
            Some(h) => ElementValue::Concrete(Element::Project {
                element: Box::new(h.clone()),
                field: name.clone(),
            }),
            None => ElementValue::Hypothetical,
        };
        let self_element = match head {
            Some(h) => Element::Project {
                element: Box::new(h.clone()),
                field: name.clone(),
            },
            None => Element::Var(name.clone()),
        };
        self.add_element(name.clone(), value, set.clone())?;
        if let Set::Record(fields) = set {
            for f in fields {
                self._recursively_add_hypothetical_element(f.name, f.carries, Some(&self_element))?;
            }
        }
        Ok(())
    }

    fn _check_literal_set_helper(
        &self,
        value: Element,
        claimed: &Set,
        should_be: Set,
    ) -> Result<(), CheckerError> {
        if !self.equal(claimed, &should_be) {
            Err(CheckerError::WrongSetForElement {
                value: value.clone().into(),
                claimed: claimed.clone(),
                real: should_be,
            })
        } else {
            Ok(())
        }
    }

    #[instrument(skip(self), level = "debug", fields(%element, set=%set.map(|s| s.to_string()).unwrap_or_default()))]
    pub fn check_element(
        &self,
        element: &Element,
        set: Option<&Set>,
    ) -> Result<Element, CheckerError> {
        match element {
            Element::Literal(lit) => {
                let value = element.clone();
                // we may infer the type from the element
                if let Some(set) = set {
                    match lit {
                        Literal::Int(_) => {
                            self._check_literal_set_helper(value, set, Set::BuiltIn(BuiltIn::Int))?;
                        }
                        Literal::Nat(_) => {
                            self._check_literal_set_helper(value, set, Set::BuiltIn(BuiltIn::Nat))?;
                        }
                        Literal::Str(_) => {
                            self._check_literal_set_helper(value, set, Set::BuiltIn(BuiltIn::Str))?;
                        }
                        Literal::Bool(_) => {
                            self._check_literal_set_helper(
                                value,
                                set,
                                Set::BuiltIn(BuiltIn::Bool),
                            )?;
                        }
                        Literal::Float(_) => {
                            self._check_literal_set_helper(
                                value,
                                set,
                                Set::BuiltIn(BuiltIn::Float),
                            )?;
                        }
                    }
                }
                Ok(element.clone().into())
            }
            Element::ClaimedElement(instance) => {
                let instance = self.check_instance(
                    instance.as_ref(),
                    set.map(|s| Signature::FromSet(s.clone())).as_ref(),
                )?;
                if let Instance::ElementCoerce(element) = instance {
                    Ok(element)
                } else {
                    Ok(Element::ClaimedElement(Box::new(instance)))
                }
            }
            Element::Var(v) => {
                let lookup = self.lookup_element(&v)?;

                if let Some(set) = set {
                    let container = self.check_set(&lookup.container)?;
                    // we have previously done the work to discover the type of
                    // this element, so what we're claiming now must match!
                    if !self.equal(set, &container) {
                        return Err(CheckerError::WrongSetForElement {
                            value: element.clone().into(),
                            claimed: set.clone(),
                            real: container,
                        });
                    }
                }
                // If we found a formal binding, we have no value to report.
                // This is the ONLY source of Var as a return value for
                // check_element, so in other branches we condition our logic
                // for formal bindings on finding Var after recursing.
                if let ElementValue::Concrete(ref deref) = lookup.value {
                    Ok(deref.clone())
                } else {
                    Ok(Element::Var(v.clone()))
                }
            }
            Element::Record(assignations) => {
                // there's a very short path here for constructing {} : record {}
                if assignations.is_empty() {
                    let element = Element::Record(Vec::new());
                    if let Some(set) = set {
                        if !self.equal(set, &Set::Record(Vec::new())) {
                            return Err(CheckerError::ElementDoesNotBelong {
                                element,
                                claimed: set.clone(),
                                reason: "the set is not the empty record".to_string(),
                            });
                        };
                    };
                    return Ok(element);
                }
                // from here on assignations is non-empty

                // look up the owner for each tag
                let owners = assignations
                    .iter()
                    .map(|ea| self.lookup_record_field(&ea.name).map(|x| &x.owner))
                    .collect::<Result<Vec<_>, _>>()?;

                // make sure they're all the same
                let owner = owners[0];
                if !owners[1..].into_iter().all(|x| self.equal(*x, owner)) {
                    return Err(CheckerError::ElementInconsistentFieldChoice(
                        element.clone(),
                    ));
                };

                // if in addition we know the set, make sure it agrees
                if let Some(set) = set
                    && !self.equal(set, owner)
                {
                    return Err(CheckerError::WrongSetForElement {
                        value: element.clone().into(),
                        claimed: set.clone(),
                        real: owner.clone(),
                    });
                }

                let Set::Record(fields) = owner else {
                    panic!("invariant violation: looking up field owners did not retrieve a record")
                };

                let mut set_fnames_sorted: Vec<String> =
                    fields.iter().map(|f| f.name.clone()).collect();
                set_fnames_sorted.sort();

                let mut element_fnames_sorted: Vec<String> =
                    assignations.iter().map(|x| x.name.clone()).collect();
                element_fnames_sorted.sort();

                // make sure that we are correctly filling the record
                if set_fnames_sorted != element_fnames_sorted {
                    return Err(CheckerError::ElementDoesNotBelong {
                        element: element.clone(),
                        claimed: owner.clone(),
                        reason: format!(
                            "expected [{}] but found [{}]",
                            set_fnames_sorted.join(", "),
                            element_fnames_sorted.join(", "),
                        ),
                    });
                }

                let assignations = assignations
                    .into_iter()
                    .map(|x| (&x.name, &x.element))
                    .collect::<HashMap<_, _>>();

                let mut ctx = self.clone();
                // the basic pattern here is that we use ctx.check_* to perform
                // substitutions for us, as we steadily march through the users
                // definitions
                let sub_elements = fields
                    .iter()
                    .map(
                        |Field {
                             name: f_n,
                             carries: f_s,
                         }| {
                            let f_e = assignations
                                .get(f_n)
                                .expect("we have already checked that all fields are present");

                            let f_s = ctx.check_set(f_s)?;
                            let f_e = ctx.check_element(f_e, Some(&f_s))?;
                            ctx.add_element(f_n.clone(), f_e.clone().into(), f_s)?;

                            Ok(ElemAssign {
                                name: f_n.clone(),
                                element: f_e,
                            })
                        },
                    )
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(Element::Record(sub_elements))
            }
            Element::Project {
                element: inner,
                field,
            } => {
                // globally unique projections mean we know what the sets going
                // in and out must be
                let OwnedField {
                    field: field_set,
                    owner: owner_set,
                } = self.lookup_record_field(&field)?;

                // enforce the correct typing of the claimed result
                if let Some(set) = set
                    && !self.equal(set, field_set)
                {
                    return Err(CheckerError::WrongSetForElement {
                        value: element.clone().into(),
                        claimed: set.clone(),
                        real: field_set.clone(),
                    });
                }

                // enforce the correct typing of the element
                let inner = self.check_element(inner, Some(owner_set))?;

                // Unfortunately we still have to do something nasty here to
                // obtain the data
                match inner {
                    // We're stuck on something that bottoms out in a binding
                    // blocking computation, nothing to be done here
                    Element::Var(_)
                    | Element::Project { .. }
                    | Element::Case { .. }
                    | Element::ClaimedElement(_) => Ok(Element::Project {
                        element: Box::new(inner),
                        field: field.clone(),
                    }),
                    Element::Record(assignations) => {
                        let sub_element = assignations
                            .into_iter()
                            .find(|a| a.name == *field)
                            .expect(
                                "invariant violation: record missing field that was type-checked",
                            )
                            .element;
                        Ok(sub_element)
                    }
                    Element::Inject { .. } | Element::Literal(_) => panic!(
                        "invariant violation: check_element returned neither a record or stuck computation for record set"
                    ),
                }
            }
            Element::Inject {
                element: inner,
                field,
            } => {
                // globally unique injections mean that we know what the sets
                // going in and out must be, but compared to projections their
                // roles are here interchanged
                let OwnedField {
                    field: field_set,
                    owner: owner_set,
                } = self.lookup_variant_field(&field)?;

                // enforce the correct typing of the claimed result
                if let Some(set) = set
                    && !self.equal(set, owner_set)
                {
                    return Err(CheckerError::WrongSetForElement {
                        value: element.clone().into(),
                        claimed: set.clone(),
                        real: owner_set.clone(),
                    });
                }

                // enforce the correct typing of the element
                let element = self.check_element(inner, Some(field_set))?;
                Ok(Element::Inject {
                    element: Box::new(element),
                    field: field.clone(),
                })
            }
            Element::Case { arms, scrutinee } => {
                if arms.is_empty() {
                    return Err(CheckerError::Unimplemented(
                        "mapping out of bottom types".to_string(),
                    ));
                }

                // 1. Syntactic checks
                // -------------------
                // arms agree on the set to which the scrutinee should belong
                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]; // safe because of the above decision about bottom
                if !arm_owners.into_iter().all(|o| self.equal(owner, o)) {
                    return Err(CheckerError::ElementInconsistentCaseScrutineeSet(
                        element.clone(),
                    ));
                }

                // all cases are handled
                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,
                    });
                }

                // 2. semantic checks
                // ------------------

                // scrutinee must be of the same set that all the arms are
                // implying, in particular this implies that the following holds
                // `inner : self.lookup_variant_field(field).field_set`
                let scrutinee = self.check_element(scrutinee, Some(owner))?;

                // which variant are we, if any
                let matching: Option<(String, Element)> = match scrutinee {
                    Element::Inject {
                        ref field,
                        element: ref inner,
                    } => Some((field.clone(), *inner.clone())),
                    // These are all the cases which could become stuck on a
                    // formal binding
                    Element::Var(_)
                    | Element::Project { .. }
                    | Element::Case { .. }
                    | Element::ClaimedElement(_) => None,
                    Element::Literal(_) | Element::Record(_) => panic!(
                        "invariant violation: scrutinee is a non-variant value at variant set"
                    ),
                };

                // are we allowed to posit the equality of elements scrutinee =
                // (arm.tag). (arm.bound) when looking at necessarily
                // non-matching arms?
                let posit_equality_with = match &scrutinee {
                    Element::Var(v) => Some(v.clone()),
                    Element::Inject { .. }
                    | Element::ClaimedElement(_)
                    | Element::Project { .. }
                    | Element::Case { .. }
                    | Element::Literal(_)
                    | Element::Record(_) => None,
                };

                // for each arm, recurse with a concrete value (if we have one)
                // otherwise fall back to hypothetical elements; in the former
                // case record the end result
                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_set = set.map(|set| ctx.check_set(set)).transpose()?;

                        let output = ctx.check_element((&arm.body).into(), this_set.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.make_element_definition(
                                scrutinee_var.clone(),
                                Element::Inject {
                                    field: arm.tag.clone(),
                                    element: Box::new(Element::Var(canonical.clone())),
                                },
                                owner.clone(),
                            )?;
                        };
                        let this_set = set.map(|set| ctx.check_set(set)).transpose()?;
                        let body = ctx.check_element((&arm.body).into(), this_set.as_ref())?;

                        CaseArm {
                            tag: arm.tag.clone(),
                            bound: canonical,
                            body,
                        }
                    };

                    processed_arms.push(case_arm);
                }
                Ok(computed_output.unwrap_or(Element::Case {
                    scrutinee: Box::new(scrutinee),
                    arms: processed_arms,
                }))
            }
        }
    }
}