use crate::ast::*;
use tracing::instrument;

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

#[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}")]
pub struct SetField {
    pub field_set: Set,
    pub owner_set: Set,
}

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

#[derive(Default)]
pub 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(())
    }
}

// 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))]
    pub fn set_equal(&self, set_a: &Set, set_b: &Set) -> bool {
        set_a == set_b
    }

    #[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))]
    pub 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(())
    }

    pub 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(())
    }

    pub fn lookup_set(&self, name: &String) -> Result<&Set, CheckError> {
        self.wf_sets
            .get(name)
            .map_or(Err(CheckError::Unbound(name.clone())), Ok)
    }

    pub fn lookup_element(&self, name: &String) -> Result<&CheckedElement, CheckError> {
        self.wf_elements
            .get(name)
            .map_or(Err(CheckError::Unbound(name.clone())), Ok)
    }

    pub fn lookup_record_field(&self, name: &String) -> Result<&SetField, CheckError> {
        self.record_fields
            .get(name)
            .map_or(Err(CheckError::Unbound(name.clone())), Ok)
    }

    pub fn lookup_variant_field(&self, name: &String) -> Result<&SetField, CheckError> {
        self.variant_fields
            .get(name)
            .map_or(Err(CheckError::Unbound(name.clone())), Ok)
    }
}