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

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

#[derive(Display)]
pub enum CheckerError {
    #[display("Unbound: {_0}")]
    Unbound(String),
    #[display("Rebinding: {_0}")]
    Rebinding(String),
    #[display("The following functionality is unimplemented: {_0}")]
    Unimplemented(String),
    #[display("Element {value} claimed to belong to {claimed} but actually belongs to {real}")]
    WrongSetForElement {
        value: ElementValue,
        claimed: Set,
        real: Set,
    },
    #[display("Element {element} does belong to set {claimed}: {reason}")]
    ElementDoesNotBelong {
        element: Element,
        claimed: Set,
        reason: String,
    },
    #[display("Case analysis {_0} does not have consistent set for scrutinee")]
    IncosistentCaseScrutineeSet(Element),
    #[display("Incomplete case analysis: covered [{}] but required [{}]",
        found.join(", "),
        required.join(", ")
    )]
    IncompleteCaseAnalysis {
        found: Vec<String>,
        required: Vec<String>,
    },
}

#[derive(Display, Clone)]
#[display("{field} @ {owner}")]
pub struct Field<T: std::fmt::Display> {
    pub field: T,
    pub owner: T,
}

#[derive(Display, Clone)]
pub enum ElementValue {
    Concrete(Element),
    #[display("_ : {_0}")]
    Hypothetical(Set),
}

impl From<Element> for ElementValue {
    fn from(e: Element) -> ElementValue {
        ElementValue::Concrete(e)
    }
}

#[derive(Display, Clone)]
#[display("{value} : {set}")]
pub struct CheckedElement {
    pub value: ElementValue,
    pub set: Set,
}

#[derive(Default, Clone)]
pub struct CheckerState {
    wf_sets: HashMap<String, Set>,
    wf_elements: HashMap<String, CheckedElement>,
    wf_signatures: HashMap<String, Signature>,
    wf_instances: HashMap<String, Instance>,
    record_fields: HashMap<String, Field<Set>>,
    variant_fields: HashMap<String, Field<Set>>,
    signature_fields: HashMap<String, Field<Signature>>,
}

impl fmt::Display for CheckerState {
    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,
        {
            writeln!(f, "  {name} = {{")?;
            let mut entries: Vec<_> = map.iter().collect();
            entries.sort_by(|a, b| a.0.cmp(b.0));
            for (k, v) in entries {
                writeln!(f, "    {k} ~> {v},")?;
            }
            writeln!(f, "  }},")
        }

        writeln!(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)?;
        writeln!(f, " }}")?;
        Ok(())
    }
}

// -----------------------------------------------------------------------------
// Equality

// The invariant we're maintaining is that everything is fully evaluated before
// we commit it to be stored in the state. Because of our invariant we don't
// actually need to do anything non-trivial here.

pub trait DecideEquality<T> {
    fn equal(&self, thing_a: &T, thing_b: &T) -> bool;
}

impl DecideEquality<Set> for CheckerState {
    #[instrument(skip(self), level = "debug", fields(%set_a, %set_b))]
    fn equal(&self, set_a: &Set, set_b: &Set) -> bool {
        set_a == set_b
    }
}

impl DecideEquality<Signature> for CheckerState {
    #[instrument(skip(self), level = "debug", fields(%signature_a, %signature_b))]
    fn equal(&self, signature_a: &Signature, signature_b: &Signature) -> bool {
        signature_a == signature_b
    }
}

impl CheckerState {
    #[instrument(skip(self), level = "debug", fields(%name, %field, %belongs_to))]
    fn assert_correct_owner<T>(
        &self,
        name: &String,
        field: &Field<T>,
        belongs_to: &T,
    ) -> Result<(), CheckerError>
    where
        Self: DecideEquality<T>,
        T: std::fmt::Display,
    {
        if !self.equal(&field.owner, belongs_to) {
            Err(CheckerError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }
}

// -----------------------------------------------------------------------------
// Sets
impl CheckerState {
    fn assert_unbound_set(&self, name: &String) -> Result<(), CheckerError> {
        if self.wf_sets.contains_key(name) {
            Err(CheckerError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    fn assert_unbound_element(&self, name: &String) -> Result<(), CheckerError> {
        if self.wf_elements.contains_key(name) {
            Err(CheckerError::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<(), CheckerError> {
        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(),
            Field {
                field: field_set.clone(),
                owner: 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<(), CheckerError> {
        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(),
            Field {
                field: field_set.clone(),
                owner: owner_set.clone(),
            },
        );
        Ok(())
    }

    #[instrument(skip(self), level = "debug", fields(%name, %set))]
    pub fn add_set(&mut self, name: &String, set: Set) -> Result<(), CheckerError> {
        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(())
    }

    #[instrument(skip(self), level = "debug", fields(%name, %element, %set))]
    pub fn add_element(
        &mut self,
        name: String,
        element: ElementValue,
        set: Set,
    ) -> Result<(), CheckerError> {
        self.assert_unbound_element(&name)?;
        self.wf_elements.insert(
            name,
            CheckedElement {
                value: element,
                set,
            },
        );
        Ok(())
    }

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

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

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

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

// -----------------------------------------------------------------------------
// Signatures
impl CheckerState {
    fn assert_unbound_signature(&self, name: &String) -> Result<(), CheckerError> {
        if self.wf_signatures.contains_key(name) {
            Err(CheckerError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    // fn assert_unbound_instance(&self, name: &String) -> Result<(), CheckerError> {
    //     if self.wf_instances.contains_key(name) {
    //         Err(CheckerError::Rebinding(name.clone()))
    //     } else {
    //         Ok(())
    //     }
    // }

    #[instrument(skip(self), level = "debug", fields(%name, %field_signature, %owner_signature))]
    fn add_signature_field(
        &mut self,
        name: &String,
        field_signature: &Signature,
        owner_signature: &Signature,
    ) -> Result<(), CheckerError> {
        if let Some(signature_ref) = self.signature_fields.get(name) {
            self.assert_correct_owner(name, signature_ref, owner_signature)?;
        };
        self.signature_fields.insert(
            name.clone(),
            Field {
                field: field_signature.clone(),
                owner: owner_signature.clone(),
            },
        );
        Ok(())
    }

    #[instrument(skip(self), level = "debug", fields(%name, %signature))]
    pub fn add_signature(
        &mut self,
        name: &String,
        signature: Signature,
    ) -> Result<(), CheckerError> {
        self.assert_unbound_signature(name)?;

        match &signature {
            Signature::Theory(fields) => {
                for SigField {
                    name: field_name,
                    signature: field_sig,
                } in fields
                {
                    self.add_signature_field(field_name, field_sig, &signature)?;
                }
            }
            // TODO: is there more?
            _ => (),
        };

        self.wf_signatures.insert(name.clone(), signature);

        Ok(())
    }

    pub fn lookup_signature(&self, name: &String) -> Result<&Signature, CheckerError> {
        self.wf_signatures
            .get(name)
            .map_or(Err(CheckerError::Unbound(name.clone())), Ok)
    }
}