use crate::ast::*;
use crate::check_state::*;

use std::iter::zip;
use tracing::instrument;

impl CheckState {
    #[instrument(skip(self), level = "debug", fields(%set))]
    pub fn check_set(&self, set: &Set) -> Result<Set, CheckError> {
        match set {
            Set::BuiltIn(_) => Ok(set.clone()),
            Set::Record(fields) => {
                let fields = fields
                    .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
                    .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(CheckError::Unimplemented("instances as sets".to_string())),
            Set::Var(v) => {
                let deref = self.lookup_set(v)?;
                Ok(deref.clone())
            }
        }
    }

    fn _check_literal_set_helper(&self, claimed: &Set, should_be: Set) -> Result<(), CheckError> {
        if !self.set_equal(claimed, &should_be) {
            Err(CheckError::WrongSetForElement(claimed.clone(), should_be))
        } else {
            Ok(())
        }
    }

    #[instrument(skip(self), level = "debug", fields(%element, %set))]
    pub fn check_element(&self, element: &Element, set: &Set) -> Result<Element, CheckError> {
        match element {
            Element::Literal(lit) => {
                // we may infer the type from the element
                match lit {
                    Literal::Int(_) => {
                        self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Int))?;
                    }
                    Literal::Nat(_) => {
                        self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Nat))?;
                    }
                    Literal::Str(_) => {
                        self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Str))?;
                    }
                    Literal::Bool(_) => {
                        self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Bool))?;
                    }
                    Literal::Float(_) => {
                        self._check_literal_set_helper(set, Set::BuiltIn(BuiltIn::Float))?;
                    }
                }
                Ok(element.clone())
            }
            Element::Var(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.set_equal(set, &lookup.set) {
                    return Err(CheckError::WrongSetForElement(
                        set.clone(),
                        lookup.set.clone(),
                    ));
                }
                Ok(lookup.element.clone())
            }
            Element::Record(assignations) => {
                let rej = |reason| CheckError::ElementDoesNotBelong {
                    element: element.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, &e_s))
                    .collect::<Result<Vec<_>, _>>()?;
                // rebuild
                let assignations = zip(element_fnames, sub_els)
                    .map(|(name, element)| ElemAssign { name, element })
                    .collect();
                // resign?
                Ok(Element::Record(assignations))
            }
            Element::Project {
                element: inner,
                field,
            } => {
                // globally unique projections mean we know what the sets going
                // in and out must be
                let SetField {
                    field_set,
                    owner_set,
                } = self.lookup_record_field(field)?;

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

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

                // 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)
            }
            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 SetField {
                    field_set,
                    owner_set,
                } = self.lookup_variant_field(field)?;

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

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

                Ok(Element::Inject {
                    element: Box::new(element),
                    field: field.clone(),
                })
            }
            Element::Case { .. } => Err(CheckError::Unimplemented("element case".to_string())),
        }
    }
}