use crate::ast::*;
use tracing::{debug, instrument};

use derive_more::Display;
use std::collections::HashMap;
use std::fmt;
use std::iter::zip;

impl Programme {
    pub fn check(&self) -> Result<(), CheckError> {
        let mut state = CheckState::default();
        state.check(self)
    }
}

#[derive(Display)]
pub enum CheckError {
    #[display("Unbound: {_0}")]
    Unbound(String),
    #[display("Rebinding: {_0}")]
    Rebinding(String),
    #[display("The following functionality is unimplemented: {_0}")]
    Unimplemented(String),
    #[display("Element claimed to belong to {_0} but actually belongs to {_1}")]
    WrongSetForElement(Set, Set),
    #[display("Element {element} does belong to set {claimed}: {reason}")]
    ElementDoesNotBelong {
        element: Element,
        claimed: Set,
        reason: String,
    },
}

#[derive(Display)]
#[display("{field_set} @ {owner_set}")]
struct SetField {
    field_set: Set,
    owner_set: Set,
}

#[derive(Display)]
#[display("{element} : {set}")]
struct CheckedElement {
    element: Element,
    set: Set,
}

#[derive(Default)]
struct CheckState {
    wf_sets: HashMap<String, Set>,
    wf_elements: HashMap<String, CheckedElement>,
    wf_signatures: HashMap<String, Signature>,
    wf_instances: HashMap<String, Instance>,
    record_fields: HashMap<String, SetField>,
    variant_fields: HashMap<String, SetField>,
}

impl fmt::Display for CheckState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fn section<K, V>(f: &mut fmt::Formatter<'_>, name: &str, map: &HashMap<K, V>) -> fmt::Result
        where
            K: fmt::Display + Ord,
            V: fmt::Display,
        {
            write!(f, " {name} = {{")?;
            let mut entries: Vec<_> = map.iter().collect();
            entries.sort_by(|a, b| a.0.cmp(b.0));
            for (k, v) in entries {
                write!(f, "{k} ~> {v}, ")?;
            }
            write!(f, "}},")
        }

        write!(f, "CheckState {{")?;
        section(f, "sets", &self.wf_sets)?;
        section(f, "elements", &self.wf_elements)?;
        section(f, "record_fields", &self.record_fields)?;
        section(f, "variant_fields", &self.variant_fields)?;
        section(f, "signatures", &self.wf_signatures)?;
        section(f, "instances", &self.wf_instances)?;
        write!(f, " }}")?;
        Ok(())
    }
}

impl CheckState {
    #[instrument(skip(self), level = "debug")]
    fn assert_unbound_set(&self, name: &String) -> Result<(), CheckError> {
        if self.wf_sets.contains_key(name) {
            Err(CheckError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    #[instrument(skip(self), level = "debug")]
    fn assert_unbound_element(&self, name: &String) -> Result<(), CheckError> {
        if self.wf_elements.contains_key(name) {
            Err(CheckError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    #[instrument(skip(self), level = "debug", fields(%name, %set_ref, %belongs_to))]
    fn assert_correct_owner(
        &self,
        name: &String,
        set_ref: &SetField,
        belongs_to: &Set,
    ) -> Result<(), CheckError> {
        if !self.set_equal(&set_ref.owner_set, belongs_to) {
            Err(CheckError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    #[instrument(skip(self), level = "debug", fields(%name, %field_set, %owner_set))]
    fn add_record_field(
        &mut self,
        name: &String,
        field_set: &Set,
        owner_set: &Set,
    ) -> Result<(), CheckError> {
        if let Some(set_ref) = self.record_fields.get(name) {
            self.assert_correct_owner(name, set_ref, owner_set)?;
        };
        self.record_fields.insert(
            name.clone(),
            SetField {
                field_set: field_set.clone(),
                owner_set: owner_set.clone(),
            },
        );
        Ok(())
    }

    #[instrument(skip(self), level = "debug", fields(%name, %field_set, %owner_set))]
    fn add_variant_field(
        &mut self,
        name: &String,
        field_set: &Set,
        owner_set: &Set,
    ) -> Result<(), CheckError> {
        if let Some(set_ref) = self.variant_fields.get(name) {
            self.assert_correct_owner(name, set_ref, owner_set)?;
        };
        self.variant_fields.insert(
            name.clone(),
            SetField {
                field_set: field_set.clone(),
                owner_set: owner_set.clone(),
            },
        );
        Ok(())
    }

    #[instrument(skip(self), level = "debug", fields(%name, %set))]
    fn add_set(&mut self, name: &String, set: Set) -> Result<(), CheckError> {
        self.assert_unbound_set(name)?;
        match &set {
            Set::Record(fields) => {
                for RecordField {
                    name: rfn,
                    set: field_set,
                } in fields
                {
                    self.add_record_field(rfn, field_set, &set)?;
                }
            }
            Set::Variant(fields) => {
                for VariantField {
                    name: vfn,
                    set: field_set,
                } in fields
                {
                    self.add_variant_field(vfn, field_set, &set)?;
                }
            }
            _ => (),
        };
        self.wf_sets.insert(name.clone(), set);
        Ok(())
    }

    fn add_element(&mut self, name: &String, element: Element, set: Set) -> Result<(), CheckError> {
        self.assert_unbound_element(name)?;
        self.wf_elements
            .insert(name.clone(), CheckedElement { element, set });
        Ok(())
    }
}

// The invariant we're maintaining is that everything is fully evaluated before
// we commit it to be stored in the state.
impl CheckState {
    // Because of our invariant we don't actually need to do anything
    // non-trivial here.
    #[instrument(skip(self), level = "debug", fields(%set_a, %set_b))]
    fn set_equal(&self, set_a: &Set, set_b: &Set) -> bool {
        set_a == set_b
    }

    #[instrument(skip(self, prog), level = "debug")]
    pub fn check(&mut self, prog: &Programme) -> Result<(), CheckError> {
        let Programme(decls) = prog;

        for decl in decls {
            debug!(%self);
            debug!(%decl);
            match decl {
                Decl::Set { name, set } => {
                    let set = self.check_set(set)?;
                    self.add_set(name, set)
                }

                Decl::Element { name, element, set } => {
                    let set = self.check_set(set)?;
                    let element = self.check_element(element, &set)?;
                    self.add_element(name, element, set)
                }
                Decl::Signature { .. } => {
                    return Err(CheckError::Unimplemented("signatures".to_string()));
                }
                Decl::Instance { .. } => {
                    return Err(CheckError::Unimplemented("instances".to_string()));
                }
            }?;
        }
        debug!(%self, "END");
        Ok(())
    }

    #[instrument(skip(self), level = "debug", fields(%set))]
    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) => {
                if let Some(deref) = self.wf_sets.get(v) {
                    Ok(deref.clone())
                } else {
                    Err(CheckError::Unbound(v.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))]
    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) => {
                if let Some(CheckedElement {
                    element: found_element,
                    set: found_set,
                }) = self.wf_elements.get(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, found_set) {
                        return Err(CheckError::WrongSetForElement(
                            set.clone(),
                            found_set.clone(),
                        ));
                    }
                    Ok(found_element.clone())
                } else {
                    Err(CheckError::Unbound(v.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 Some(SetField {
                    field_set,
                    owner_set,
                }) = self.record_fields.get(field)
                else {
                    return Err(CheckError::Unbound(field.clone()));
                };

                // 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 Some(SetField {
                    field_set,
                    owner_set,
                }) = self.variant_fields.get(field)
                else {
                    return Err(CheckError::Unbound(field.clone()));
                };

                // 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())),
        }
    }
}