#![allow(clippy::enum_variant_names)]

use std::{
    collections::{hash_map::Entry, HashMap},
    fmt,
    rc::Rc,
};

use crate::{
    ast,
    ast::ItemId,
    builtin::Builtin,
    parsing::Location,
    type_inference::{TypeContext, TypeId},
};

mod populator;
use populator::SymbolTablePopulator;
mod fqsn;
pub use fqsn::{Fqsn, ScopeSegment};
mod resolver;
mod symbol_trie;
use symbol_trie::SymbolTrie;
mod test;
use crate::identifier::{define_id_kind, Id, IdStore};

define_id_kind!(DefItem);
pub type DefId = Id<DefItem>;

#[allow(dead_code)]
#[derive(Debug, Clone)]
pub enum SymbolError {
    DuplicateName { prev_name: Fqsn, location: Location },
    DuplicateVariant { type_fqsn: Fqsn, name: String },
    DuplicateRecord { type_name: Fqsn, location: Location, member: String },
    UnknownAnnotation { name: String },
    BadAnnotation { name: String, msg: String },
}

#[allow(dead_code)]
#[derive(Debug)]
struct NameSpec<K> {
    location: Location,
    kind: K,
}

#[derive(Debug)]
enum NameKind {
    Module,
    Function,
    Binding,
}

#[derive(Debug)]
struct TypeKind;

/// Keeps track of what names were used in a given namespace.
struct NameTable<K> {
    table: HashMap<Fqsn, NameSpec<K>>,
}

impl<K> NameTable<K> {
    fn new() -> Self {
        Self { table: HashMap::new() }
    }

    fn register(&mut self, name: Fqsn, spec: NameSpec<K>) -> Result<(), SymbolError> {
        match self.table.entry(name.clone()) {
            Entry::Occupied(o) =>
                Err(SymbolError::DuplicateName { prev_name: name, location: o.get().location }),
            Entry::Vacant(v) => {
                v.insert(spec);
                Ok(())
            }
        }
    }
}

//cf. p. 150 or so of Language Implementation Patterns
pub struct SymbolTable {
    def_id_store: IdStore<DefItem>,

    /// Used for import resolution.
    symbol_trie: SymbolTrie,

    /// These tables are responsible for preventing duplicate names.
    fq_names: NameTable<NameKind>, //Note that presence of two tables implies that a type and other binding with the same name can co-exist
    types: NameTable<TypeKind>,

    id_to_def: HashMap<ItemId, DefId>,
    def_to_symbol: HashMap<DefId, Rc<Symbol>>,
}

impl SymbolTable {
    /// Create a new, empty SymbolTable
    pub fn new() -> Self {
        Self {
            def_id_store: IdStore::new(),
            symbol_trie: SymbolTrie::new(),
            fq_names: NameTable::new(),
            types: NameTable::new(),
            id_to_def: HashMap::new(),
            def_to_symbol: HashMap::new(),
        }
    }

    /// The main entry point into the symbol table. This will traverse the AST in several
    /// different ways and populate subtables with information that will be used further in the
    /// compilation process.
    pub fn process_ast(
        &mut self,
        ast: &ast::AST,
        type_context: &mut TypeContext,
    ) -> Result<(), Vec<SymbolError>> {
        let mut populator = SymbolTablePopulator { type_context, table: self };

        let errs = populator.populate_name_tables(ast);
        if !errs.is_empty() {
            return Err(errs);
        }

        // Walks the AST, matching the ID of an identifier used in some expression to
        // the corresponding Symbol.
        let mut resolver = resolver::ScopeResolver::new(self);
        resolver.resolve(ast);

        Ok(())
    }

    pub fn lookup_symbol(&self, id: &ItemId) -> Option<&Symbol> {
        let def = self.id_to_def.get(id)?;
        self.def_to_symbol.get(def).map(|s| s.as_ref())
    }

    pub fn lookup_symbol_by_def(&self, def: &DefId) -> Option<&Symbol> {
        self.def_to_symbol.get(def).map(|s| s.as_ref())
    }

    #[allow(dead_code)]
    pub fn debug(&self) {
        println!("Symbol table:");
        println!("----------------");
        for (id, def) in self.id_to_def.iter() {
            if let Some(symbol) = self.def_to_symbol.get(def) {
                println!("{} => {}: {}", id, def, symbol);
            } else {
                println!("{} => {} <NO SYMBOL FOUND>", id, def);
            }
        }
    }

    /// Register a new mapping of a fully-qualified symbol name (e.g. `Option::Some`)
    /// to a Symbol, a descriptor of what that name refers to.
    fn add_symbol(&mut self, id: &ItemId, fqsn: Fqsn, spec: SymbolSpec) {
        let def_id = self.def_id_store.fresh();
        let symbol = Rc::new(Symbol { fully_qualified_name: fqsn.clone(), spec, def_id });
        self.symbol_trie.insert(&fqsn, def_id);
        self.id_to_def.insert(*id, def_id);
        self.def_to_symbol.insert(def_id, symbol);
    }

    fn populate_single_builtin(&mut self, fqsn: Fqsn, builtin: Builtin) {
        let def_id = self.def_id_store.fresh();
        let spec = SymbolSpec::Builtin(builtin);
        let symbol = Rc::new(Symbol { fully_qualified_name: fqsn.clone(), spec, def_id });

        self.symbol_trie.insert(&fqsn, def_id);
        self.def_to_symbol.insert(def_id, symbol);
    }
}

#[allow(dead_code)]
#[derive(Debug, Clone)]
pub struct Symbol {
    fully_qualified_name: Fqsn,
    spec: SymbolSpec,
    def_id: DefId,
}

impl Symbol {
    pub fn local_name(&self) -> Rc<String> {
        self.fully_qualified_name.last_elem()
    }

    pub fn def_id(&self) -> Option<DefId> {
        Some(self.def_id)
    }

    pub fn spec(&self) -> SymbolSpec {
        self.spec.clone()
    }
}

impl fmt::Display for Symbol {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "<Local name: {}, {}, Spec: {}>", self.local_name(), self.fully_qualified_name, self.spec)
    }
}

//TODO - I think I eventually want to draw a distinction between true global items
//i.e. global vars, and items whose definitions are scoped. Right now there's a sense
//in which Func, DataConstructor, RecordConstructor, and GlobalBinding are "globals",
//whereas LocalVarible and FunctionParam have local scope. But right now, they all
//get put into a common table, and all get DefId's from a common source.
//
//It would be good if individual functions could in parallel look up their own
//local vars without interfering with other lookups. Also some type definitions
//should be scoped in a similar way.
//
//Also it makes sense that non-globals should not use DefId's, particularly not
//function parameters (even though they are currently assigned).
#[derive(Debug, Clone)]
pub enum SymbolSpec {
    Builtin(Builtin),
    Func,
    DataConstructor { tag: u32, type_id: TypeId },
    RecordConstructor { tag: u32, type_id: TypeId },
    GlobalBinding, //Only for global variables, not for function-local ones or ones within a `let` scope context
    LocalVariable,
    FunctionParam(u8),
}

impl fmt::Display for SymbolSpec {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::SymbolSpec::*;
        match self {
            Builtin(b) => write!(f, "Builtin: {:?}", b),
            Func => write!(f, "Func"),
            DataConstructor { tag, type_id } => write!(f, "DataConstructor(tag: {}, type: {})", tag, type_id),
            RecordConstructor { type_id, tag, .. } =>
                write!(f, "RecordConstructor(tag: {})(<members> -> {})", tag, type_id),
            GlobalBinding => write!(f, "GlobalBinding"),
            LocalVariable => write!(f, "Local variable"),
            FunctionParam(n) => write!(f, "Function param: {}", n),
        }
    }
}