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>,
    },
    #[display("Instance {value} claimed to belong to {claimed} but actually belongs to {real}")]
    WrongSignatureForInstance {
        value: InstanceValue,
        claimed: Signature,
        real: Signature,
    },
    #[display("Instance {instance} is not of signature {claimed}: {reason}")]
    InstanceDoesNotBelong {
        instance: Instance,
        claimed: Signature,
        reason: String,
    },
    #[display("Non-functional instance {instance} found in application to element {element}")]
    NonFunctionalInstance {
        instance: Instance,
        element: Element,
    },
}

// -----------------------------------------------------------------------------
// Generics for wrapping fields, values, and coercing them
#[derive(Display, Clone)]
#[display("{field} @ {owner}")]
pub struct Field<T: std::fmt::Display> {
    pub field: T,
    pub owner: T,
}

#[derive(Display, Clone)]
pub enum Value<Term: std::fmt::Display> {
    /// The storage format for concrete terms.
    Concrete(Term),
    #[display("_")]
    /// The storage format for formal bindings.
    Hypothetical,
}

pub type ElementValue = Value<Element>;
pub type InstanceValue = Value<Instance>;
pub type SetValue = Value<Set>;

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

impl From<Instance> for InstanceValue {
    fn from(i: Instance) -> InstanceValue {
        Value::Concrete(i)
    }
}

impl From<Set> for SetValue {
    fn from(s: Set) -> SetValue {
        Value::Concrete(s)
    }
}

#[derive(Display, Clone)]
#[display("{value} : {container}")]
pub struct Checked<Term: std::fmt::Display, Type: std::fmt::Display> {
    pub value: Value<Term>,
    pub container: Type,
}

pub type CheckedElement = Checked<Element, Set>;
pub type CheckedInstance = Checked<Instance, Signature>;

// -----------------------------------------------------------------------------
// The checker state
#[derive(Default, Clone)]
pub struct CheckerState {
    wf_sets: HashMap<String, SetValue>,
    wf_elements: HashMap<String, CheckedElement>,
    wf_signatures: HashMap<String, Signature>,
    wf_instances: HashMap<String, CheckedInstance>,
    record_fields: HashMap<String, Field<Set>>,
    variant_fields: HashMap<String, Field<Set>>,
    signature_fields: HashMap<String, Field<Signature>>,
    // TODO: do we need to make these strictly monotonic somewhere somehow?
    binder_element: usize,
    unique_name: usize,
}

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)?;
        section(f, "signature_fields", &self.signature_fields)?;
        writeln!(f, " }}")?;
        Ok(())
    }
}

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

// we work very hard to store canonical forms so that equality is purely
// structural. This approach may or may not survive contact with reality.
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 {
    pub fn assert_unbound_set(&self, name: &String) -> Result<(), CheckerError> {
        if self.wf_sets.contains_key(name) {
            Err(CheckerError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    pub 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: SetValue) -> Result<(), CheckerError> {
        match &set {
            SetValue::Concrete(set @ Set::Record(fields)) => {
                for RecordField {
                    name: rfn,
                    set: field_set,
                } in fields
                {
                    self.add_record_field(rfn, field_set, set)?;
                }
            }
            SetValue::Concrete(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, set.into());
        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.wf_elements.insert(
            name,
            Checked {
                value: element,
                container: set,
            },
        );
        Ok(())
    }

    pub fn lookup_set(&self, name: &String) -> Result<&SetValue, 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 {
    pub fn assert_unbound_signature(&self, name: &String) -> Result<(), CheckerError> {
        if self.wf_signatures.contains_key(name) {
            Err(CheckerError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    pub 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))]
    pub fn add_signature_field(
        &mut self,
        name: &String,
        field_signature: &Signature,
        owner_signature: &Signature,
        rebind: bool,
    ) -> Result<(), CheckerError> {
        if !rebind && 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,
        rebind: bool,
    ) -> Result<(), CheckerError> {
        match &signature {
            Signature::Theory(fields) => {
                for SigField {
                    name: field_name,
                    signature: field_sig,
                } in fields
                {
                    // TODO: are we supposed to recurse?
                    // let name = self.make_unique_name();
                    // self.add_signature(&name, field_sig.clone(), rebind)?;
                    self.add_signature_field(field_name, field_sig, &signature, rebind)?;
                }
            }
            // TODO: is there more?
            _ => (),
        };

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

        Ok(())
    }

    #[instrument(skip(self), level = "debug", fields(%name, %instance, %signature))]
    pub fn add_instance(
        &mut self,
        name: String,
        instance: InstanceValue,
        signature: Signature,
    ) -> Result<(), CheckerError> {
        self.wf_instances.insert(
            name,
            Checked {
                value: instance,
                container: 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)
    }

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

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

// -----------------------------------------------------------------------------
// Bindings
impl CheckerState {
    fn _make_canonical_element(&mut self, name: String, set: Set) -> Result<String, CheckerError> {
        let canonical = format!("_#{}", self.binder_element);
        self.add_element(name, Element::Var(canonical.clone()).into(), set)?;
        self.binder_element += 1;
        Ok(canonical)
    }

    #[instrument(skip(self), level = "debug", fields(%name, %set))]
    pub fn make_element_binding(&mut self, name: String, set: Set) -> Result<String, CheckerError> {
        let canonical = self._make_canonical_element(name, set.clone())?;
        self.add_element(canonical.clone(), ElementValue::Hypothetical, set)?;
        Ok(canonical)
    }

    #[instrument(skip(self), level = "debug", fields(%name, %set))]
    pub fn make_element_definition(
        &mut self,
        name: String,
        value: Element,
        set: Set,
    ) -> Result<String, CheckerError> {
        let canonical = self._make_canonical_element(name, set.clone())?;
        self.add_element(canonical.clone(), Value::Concrete(value), set)?;
        Ok(canonical)
    }

    #[instrument(skip(self))]
    pub fn make_unique_name(&mut self) -> String {
        self.unique_name += 1;
        format!("_#{}", self.unique_name)
    }
}