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

use std::iter::zip;
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 fields = fields
                    .into_iter()
                    .map(|RecordField { name, set }| {
                        let set = self.check_set(set)?;
                        Ok(RecordField {
                            name: name.clone(),
                            set,
                        })
                    })
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(Set::Record(fields))
            }
            Set::Variant(fields) => {
                let fields = fields
                    .into_iter()
                    .map(|VariantField { name, set }| {
                        let set = self.check_set(set)?;
                        Ok(VariantField {
                            name: name.clone(),
                            set,
                        })
                    })
                    .collect::<Result<Vec<_>, _>>()?;
                Ok(Set::Variant(fields))
            }
            Set::ClaimedSet(_) => Err(CheckerError::Unimplemented("instances as sets".to_string())),
            Set::Var(v) => {
                let deref = self.lookup_set(&v)?;
                Ok(deref.clone())
            }
        }
    }

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

    #[instrument(skip(self), level = "debug", fields(%value, %set))]
    pub fn check_element(
        &self,
        value: ElementValue,
        set: &Set,
    ) -> Result<ElementValue, CheckerError> {
        match value {
            // This arm ensures that if we ever in the position of obtaining a
            // hypothetical from a call to check_element, in the context of
            // check_element, we can safely ignore its payload. I'll point this
            // out later as (*)
            ElementValue::Hypothetical(ref h_set) => {
                if !self.equal(set, &h_set) {
                    Err(CheckerError::WrongSetForElement {
                        value: value.clone(),
                        claimed: set.clone(),
                        real: h_set.clone(),
                    })
                } else {
                    Ok(value)
                }
            }
            ElementValue::Concrete(ref element @ Element::Literal(ref lit)) => {
                let value = value.clone();
                // we may infer the type from the element
                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())
            }
            ElementValue::Concrete(Element::Var(ref v)) => {
                let lookup = self.lookup_element(&v)?;
                // 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, &lookup.set) {
                    return Err(CheckerError::WrongSetForElement {
                        value: value.clone(),
                        claimed: set.clone(),
                        real: lookup.set.clone(),
                    });
                }
                Ok(lookup.value.clone())
            }
            ElementValue::Concrete(ref concrete @ Element::Record(ref assignations)) => {
                let rej = |reason| CheckerError::ElementDoesNotBelong {
                    element: concrete.clone(),
                    claimed: set.clone(),
                    reason,
                };

                // make sure we are filling a record
                let fields = if let Set::Record(fields) = set {
                    Ok(fields)
                } else {
                    Err(rej("element is a record instance".to_string()))
                }?;

                let (set_fnames, set_fsets): (Vec<String>, Vec<Set>) = fields
                    .iter()
                    .map(|RecordField { name, set }| (name.clone(), set.clone()))
                    .unzip();
                let mut set_fnames_sorted = set_fnames.clone();
                set_fnames_sorted.sort();

                let (element_fnames, element_felements): (Vec<String>, Vec<&Element>) =
                    assignations
                        .iter()
                        .map(|ElemAssign { name, element }| (name.clone(), element))
                        .unzip();

                let mut element_fnames_sorted = element_fnames.clone();
                element_fnames_sorted.sort();

                if set_fnames_sorted != element_fnames_sorted {
                    return Err(rej(format!(
                        "expected [{}] but found [{}]",
                        set_fnames.join(", "),
                        element_fnames.join(", "),
                    )));
                }

                // recurse, sets have already been completely expanded
                let sub_els = zip(element_felements, set_fsets)
                    .map(|(e_f, e_s)| self.check_element(e_f.clone().into(), &e_s))
                    .collect::<Result<Vec<_>, _>>()?;
                // rebuild, hypotheticals are contagious
                let assignations = zip(element_fnames, sub_els)
                    .map(|(name, element)| match element {
                        ElementValue::Concrete(element) => Some(ElemAssign { name, element }),
                        ElementValue::Hypothetical(_) => None,
                    })
                    .collect();
                // resign?
                Ok(if let Some(assignations) = assignations {
                    Element::Record(assignations).into()
                } else {
                    ElementValue::Hypothetical(set.clone())
                })
            }
            ElementValue::Concrete(Element::Project {
                element: ref inner,
                ref field,
            }) => {
                // globally unique projections mean we know what the sets going
                // in and out must be
                let Field {
                    field: field_set,
                    owner: owner_set,
                } = self.lookup_record_field(&field)?;

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

                // enforce the correct typing of the element
                let inner = self.check_element((*inner.clone()).into(), owner_set)?;
                match inner {
                    ElementValue::Concrete(inner) => {
                        // Unfortunately we still have to do something nasty here to obtain the data
                        let Element::Record(assignations) = inner else {
                            panic!(
                                "invariant violation: check_element returned non-record for record set"
                            );
                        };
                        let sub_element = assignations
                            .into_iter()
                            .find(|a| a.name == *field)
                            .expect(
                                "invariant violation: record missing field that was type-checked",
                            )
                            .element
                            .clone();

                        Ok(sub_element.into())
                    }
                    // correct by (*)
                    ElementValue::Hypothetical(_) => Ok(ElementValue::Hypothetical(set.clone())),
                }
            }
            ElementValue::Concrete(Element::Inject {
                element: ref inner,
                ref 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 Field {
                    field: field_set,
                    owner: owner_set,
                } = self.lookup_variant_field(&field)?;

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

                // enforce the correct typing of the element
                let element = self.check_element((*inner.clone()).into(), field_set)?;

                match element {
                    ElementValue::Concrete(element) => Ok(Element::Inject {
                        element: Box::new(element),
                        field: field.clone(),
                    }
                    .into()),
                    // correct by (*)
                    ElementValue::Hypothetical(_) => Ok(ElementValue::Hypothetical(set.clone())),
                }
            }
            ElementValue::Concrete(
                ref element @ Element::Case {
                    ref arms,
                    ref scrutinee,
                },
            ) => {
                // TODO: do we allow mapping out of bottom?
                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.iter().all(|o| self.equal(owner, o)) {
                    return Err(CheckerError::IncosistentCaseScrutineeSet(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.clone()).into(), owner)?;

                // which variant are we, if any
                let matching: Option<(String, Element)> = match scrutinee {
                    ElementValue::Hypothetical(_) => None,
                    ElementValue::Concrete(Element::Inject {
                        field,
                        element: inner,
                    }) => Some((field, *inner)),
                    _ => {
                        panic!(
                            "invariant violation: we believe element is of a variant set but it's not an injection"
                        );
                    }
                };

                // 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;
                for arm in arms {
                    let Field {
                        field: field_set, ..
                    } = self.lookup_variant_field(&arm.tag)?;

                    // TODO: if we were worried about overhead we'd have a separate
                    // locals stack, though truly if we were worried about overhead
                    // we'd not have NNN instances of clone elsewhere in the
                    // codebase and we wouldn't be eagerly evaluating all
                    // expressions fully.

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

                    if let Some((tag, inner)) = &matching
                        && *tag == arm.tag
                    {
                        new_context.add_element(binding_name, inner.clone().into(), binding_set)?;
                        let output = new_context.check_element(arm.body.clone().into(), set)?;
                        if matches!(computed_output, Some(_)) {
                            panic!(
                                "invariant violation: we somehow matched multiple arms in case analysis"
                            )
                        }
                        computed_output = Some(output);
                    } else {
                        new_context.add_element(
                            binding_name,
                            ElementValue::Hypothetical(field_set.clone()),
                            binding_set,
                        )?;
                        new_context.check_element(arm.body.clone().into(), set)?;
                    };
                }
                Ok(computed_output.unwrap_or(ElementValue::Hypothetical(set.clone())))
            }
        }
    }
}