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

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

#[derive(Display)]
pub enum CheckError {
    #[display("Unbound: {_0}")]
    Unbound(String),
    #[display("Duplicate field: {_0}")]
    DuplicateField(String),
    #[display("Rebinding: {_0}")]
    Rebinding(String),
    #[display("The following functionality is unimplemented: {_0}")]
    Unimplemented(String),
}

#[derive(Debug)]
struct SetRef {
    set: Set,
    belongs_to: String,
}

#[derive(Debug)]
struct ElRef {
    set: Set,
    belongs_to: String,
}

#[derive(Debug, Default)]
struct CheckState {
    sets: HashMap<String, Set>,
    elements: HashMap<String, Element>,
    record_fields: HashMap<String, SetRef>,
    variant_fields: HashMap<String, SetRef>,
    signatures: HashMap<String, Signature>,
    instances: HashMap<String, Instance>,
}

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

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

    #[instrument(skip(self), level = "debug")]
    fn assert_unbound_variant_field(&self, name: &String) -> Result<(), CheckError> {
        if self.variant_fields.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.elements.contains_key(name) {
            Err(CheckError::Rebinding(name.clone()))
        } else {
            Ok(())
        }
    }

    #[instrument(skip(self), level = "debug")]
    fn add_record_field(
        &mut self,
        name: &String,
        set: &Set,
        belongs_to: &String,
    ) -> Result<(), CheckError> {
        self.assert_unbound_record_field(name)?;
        self.record_fields.insert(
            name.clone(),
            SetRef {
                set: set.clone(),
                belongs_to: belongs_to.clone(),
            },
        );
        Ok(())
    }

    #[instrument(skip(self), level = "debug")]
    fn add_variant_field(
        &mut self,
        name: &String,
        set: &Set,
        belongs_to: &String,
    ) -> Result<(), CheckError> {
        self.assert_unbound_variant_field(name)?;
        self.variant_fields.insert(
            name.clone(),
            SetRef {
                set: set.clone(),
                belongs_to: belongs_to.clone(),
            },
        );
        Ok(())
    }

    #[instrument(skip(self), level = "debug")]
    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 } in fields {
                    self.add_record_field(rfn, set, name)?;
                }
            }
            Set::Variant(fields) => {
                for VariantField { name: vfn, set } in fields {
                    self.add_variant_field(vfn, set, name)?;
                }
            }
            _ => (),
        };
        self.sets.insert(name.clone(), set.clone());
        Ok(())
    }

    #[instrument(skip(self), level = "debug")]
    fn add_element(&mut self, name: &String, element: &Element) -> Result<(), CheckError> {
        self.assert_unbound_element(name)?;
        self.elements.insert(name.clone(), element.clone());
        Ok(())
    }
}

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

        for decl in decls {
            debug!(%decl, "checking declaration");
            match decl {
                Decl::Set { name, set } => {
                    self.assert_unbound_set(name)?;
                    self.check_set(set)?;
                    // One catch, prohibit "let .. X = X"
                    if let Set::Var(v) = set
                        && v == name
                    {
                        return Err(CheckError::Rebinding(v.clone()));
                    };
                    self.add_set(name, set)
                }

                Decl::Element { name, set, element } => {
                    self.check_element(name, set, element)?;
                    // TODO: don't drop the set?
                    self.add_element(name, element)
                }
                Decl::Signature { name, signature } => Ok(()),
                Decl::Instance {
                    name,
                    signature,
                    instance,
                } => Ok(()),
            }?;
        }
        Ok(())
    }

    #[instrument(skip(self), level = "debug")]
    fn check_set(&self, set: &Set) -> Result<(), CheckError> {
        match set {
            Set::BuiltIn(_) => Ok(()),
            Set::Record(fields) => self.check_record(fields),
            Set::Variant(fields) => Err(CheckError::Unimplemented("variants".to_string())),
            Set::ClaimedSet(instance) => {
                Err(CheckError::Unimplemented("instances as sets".to_string()))
            }
            Set::Var(v) => {
                if self.sets.contains_key(v) {
                    Ok(())
                } else {
                    Err(CheckError::Unbound(v.clone()))
                }
            }
        }
    }

    #[instrument(skip(self), level = "debug")]
    fn check_record(&self, fields: &Vec<RecordField>) -> Result<(), CheckError> {
        for RecordField { name, set } in fields {
            self.assert_unbound_record_field(name)?;
            self.check_set(set)?;
        }
        Ok(())
    }

    #[instrument(skip(self), level = "debug")]
    fn check_element(&self, name: &String, set: &Set, element: &Element) -> Result<(), CheckError> {
        self.assert_unbound_element(name)?;
        self.check_set(set)?;
        Ok(())
    }
}

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