// Copyright (c) 2025 Redglyph (@gmail.com). All Rights Reserved.

#[cfg(any())]
use std::collections::HashMap;
use crate::{NameFixer, indent_source};
use crate::fixed_sym_table::{FixedSymTable, SymInfoTable};
use crate::{TokenId, VarId};
use crate::parser::Symbol;
#[cfg(test)]
use lexigram_core::CollectJoin;

// NOTE: nonterminal-to-ID functionality currently disabled by #[cfg(any())]

/// Stores the names of the terminal and nonterminal symbols when building a parser.
///
/// Terminals are defined in the lexicon and don't change. They have two parts to their name:
/// - the identifier in the lexicon
/// - the value, which is the source string they represent (optional)
///
/// For example:
/// ```lexicon
/// Plus : '+';
/// ...
/// ID    : [a-zA-Z][a-zA-Z_0-9]*;
/// ```
///
/// If `Arrow`'s token ID is 0 and `ID`'s is 24,
/// ```ignore
/// t[0] = ("Plus".to_string(), Some("+".to_string()));
/// t[24] = ("ID".to_string(), None);
/// ```
///
/// They're added to the symbol table with [`add_terminal()`](SymbolTable::add_terminal).
///
/// Nonterminals are defined in the grammar, and possibly completed by new ones when
/// the rules are adapted to the target parser. For example, recursive rules are
/// transformed for LL(1) parsers, which usually adds extra rules.
///
/// ```grammar
/// expr: expr Plus term | term;
/// ```
/// If `expr` is 0 and `term` is 1,
/// ```ignore
/// nt[0] = "expr".to_string();
/// nt[1] = "term".to_string();
/// ```
/// They're added with [`add_nonterminal`](SymbolTable::add_nonterminal).
///
/// The new rules are called "children" of the transformed rules, and can be added
/// with [`add_child_nonterminal()`](SymbolTable::add_child_nonterminal). The name
/// is the parent's name followed by "`_<number>`". For example, adding a child to
/// nonterminal 0 creates
///
/// ```ignore
/// nt[2] = "expr_1".to_string()
/// ```
///
#[derive(Clone, Debug)]
pub struct SymbolTable {
    t: Vec<(String, Option<String>)>,   // terminal identifiers and optional representation
    fixer_t: NameFixer,                 // keeps terminal identifiers unique
    nt: Vec<String>,                    // nt to nonterminal identifier
    #[cfg(any())]
    names: HashMap<String, VarId>,      // nonterminal identifier to nt
    fixer_nt: NameFixer,                // keeps nonterminal identifiers unique
}

impl SymbolTable {
    pub fn new() -> Self {
        SymbolTable {
            t: Vec::new(),
            fixer_t: NameFixer::new_empty(),
            nt: Vec::new(),
            #[cfg(any())]
            names: HashMap::new(),
            fixer_nt: NameFixer::new_empty(),
        }
    }

    pub fn to_fixed_sym_table(self) -> FixedSymTable {
        FixedSymTable::new(self.t, self.nt)
    }

    // -------------------------------------------------------------------------

    pub fn add_terminal<T: Into<String>>(&mut self, name: T, value_maybe: Option<T>) -> TokenId {
        let token = self.t.len();
        assert!(token < TokenId::MAX as usize);
        let unique_name = self.fixer_t.get_unique_name(name.into());
        self.t.push((unique_name, value_maybe.map(|n| n.into())));
        token as TokenId
    }

    pub fn extend_terminals<I: IntoIterator<Item=(T, Option<T>)>, T: Into<String>>(&mut self, terminals: I) {
        for (s, maybe) in terminals {
            self.add_terminal(s, maybe);
        }
    }

    pub fn get_terminals(&self) -> impl Iterator<Item = &(String, Option<String>)> {
        self.t.iter()
    }

    pub fn get_num_t(&self) -> usize {
        self.t.len()
    }

    pub fn set_t_value(&mut self, token: TokenId, name_maybe: Option<String>) {
        self.t[token as usize].1 = name_maybe
    }

    pub fn downsize_num_t(&mut self, num_t: usize) {
        if num_t < self.t.len() {
            for (name, _) in &mut self.t[num_t..] {
                self.fixer_t.remove(name);
            }
            self.t.resize(num_t, (String::new(), None));
        }
    }

    // -------------------------------------------------------------------------

    fn add_nt(&mut self, unique_name: String) -> VarId {
        let var = self.nt.len();
        assert!(var < VarId::MAX as usize);
        #[cfg(any())]
        self.names.insert(unique_name.clone(), var as VarId);
        self.nt.push(unique_name);
        var as VarId
    }

    pub fn add_nonterminal<T: Into<String>>(&mut self, name: T) -> VarId {
        let unique_name = self.fixer_nt.get_unique_name(name.into());
        self.add_nt(unique_name)
    }

    pub fn add_child_nonterminal(&mut self, var: VarId) -> VarId {
        let unique_name = self.fixer_nt.get_unique_name_unum(self.nt[var as usize].clone());
        self.add_nt(unique_name)
    }
    
    pub fn extend_nonterminals<I: IntoIterator<Item=T>, T: Into<String>>(&mut self, nonterminals: I) {
        for s in nonterminals {
            self.add_nonterminal(s);
        }
    }

    #[cfg(any())]
    pub fn find_nonterminal(&self, name: &str) -> Option<VarId> {
        self.names.get(name).cloned()
    }

    pub fn get_nonterminals(&self) -> impl Iterator<Item = &String> {
        self.nt.iter()
    }

    pub fn get_num_nt(&self) -> usize {
        self.nt.len()
    }

    pub fn remove_nonterminal(&mut self, v: VarId) {
        let name = self.nt.remove(v as usize);
        #[cfg(any())]
        {
            for old_v in self.names.values_mut() {
                if *old_v >= v { *old_v -= 1; }
            }
            self.names.remove(&name);
        }
        self.fixer_nt.remove(&name);
    }

    pub fn set_nt_name(&mut self, var: VarId, name: String) {
        self.nt[var as usize] = name;
    }

    pub fn find_nt(&self, name: &str) -> Option<VarId> {
        self.nt.iter().position(|s| s == name).map(|n| n as VarId)
    }

    /// Removes the name assigned to NT `var` and returns it. Internally, the name of the NT is
    /// replaced by another unique string. The NT is expected to be removed later.
    pub fn remove_nt_name(&mut self, var: VarId) -> String {
        let mut removed = self.fixer_nt.get_unique_name_num(format!("{var}_removed"));
        std::mem::swap(&mut self.nt[var as usize], &mut removed);
        #[cfg(any())]
        self.names.remove(&removed);
        self.fixer_nt.remove(&removed);
        removed
    }

    pub fn gen_source_code_t(&self, indent: usize, has_enum: bool, has_array: bool) -> String {
        let mut source = Vec::<String>::new();
        let (max_n, max_option) = self.get_terminals().fold((0, 0), |(n, option), t| {
            (n.max(t.0.len()), option.max(t.1.as_ref().map_or(0, |o| o.len())))
        });
        if has_enum {
            // Arrow = 0,  // 0
            // Colon,      // 1
            source.push("#[repr(u16)]".to_string());
            source.push("enum T {".to_string());
            for (i, (n, _option)) in self.get_terminals().enumerate() {
                source.push(format!("    {n:max_n$}{} // {i}", if i == 0 { " = 0," } else { ",    " }));
            }
            source.push("}".to_string());
        }
        if has_enum && has_array {
            source.push(String::new());
        }
        if has_array {
            // ("Arrow",     Some("->")),          // 0
            // ("Colon",     Some(":")),           // 1
            source.push(format!("static TERMINALS: [(&str, Option<&str>); {}] = [", self.get_num_t()));
            for (i, (n, option)) in self.get_terminals().enumerate() {
                source.push(format!("    (\"{n}\",{:w1$}{}),{:w2$} // {i}",
                                    "",
                                    if let Some(o) = option { format!("Some(\"{o}\")") } else { "None".to_string() },
                                    "",
                                    w1 = max_n - n.len(), w2 = max_option + 4 - option.as_ref().map_or(0, |o| o.len() + 4)));
            }
            source.push("];".to_string());
        }
        indent_source(vec![source], indent)
    }

    // -------------------------------------------------------------------------

    #[cfg(test)]
    pub fn dump(&self, title: &str) {
        if !title.is_empty() {
            println!("{title}");
        }
        println!(
            "- nonterminals:\n{}",
            self.get_nonterminals().enumerate().map(|(v, s)| format!("  - NT[{v}]: {s}")).join("\n"));
        println!(
            "- terminals:\n{}",
            self.get_terminals().enumerate()
                .map(|(t, (n, v_maybe))| format!("  - T[{t}]: {n}{}", if let Some(v) = v_maybe { format!(" = {v:?}") } else { String::new() }))
                .join("\n"));

    }
}

impl SymInfoTable for SymbolTable {
    fn is_token_data(&self, token: TokenId) -> bool {
        self.t.get(token as usize).map(|t| t.1.is_none()).unwrap_or(false)
    }

    fn is_symbol_t_data(&self, symbol: &Symbol) -> bool {
        if let Symbol::T(token) = symbol {
            self.t.get(*token as usize).map(|t| t.1.is_none()).unwrap_or(false)
        } else {
            false
        }
    }

    fn is_symbol_t_fixed(&self, symbol: &Symbol) -> bool {
        if let Symbol::T(token) = symbol {
            if (*token as usize) < self.t.len() {
                self.t.get(*token as usize).map(|t| t.1.is_some()).unwrap_or(false)
            } else {
                false
            }
        } else {
            false
        }
    }

    fn get_t_str(&self, token: TokenId) -> String {
        match token {
            _ if (token as usize) < self.t.len() => {
                let (name, literal) = &self.t[token as usize];
                literal.as_ref().unwrap_or(name).clone()
            }
            TokenId::MAX => "<bad character>".to_string(),
            _ => format!("T({token}?)")
        }
    }

    fn get_t_name(&self, token: TokenId) -> String {
        if token as usize >= self.t.len() {
            format!("T({token}?)")
        } else {
            self.t[token as usize].0.clone()
        }
    }

    fn get_nt_name(&self, var: VarId) -> String {
        if var as usize >= self.nt.len() { return format!("NT({var}?)") }
        self.nt[var as usize].clone()
    }

    fn get_name(&self, symbol: &Symbol) -> String {
        match symbol {
            Symbol::Empty | Symbol::End => symbol.to_string(),
            Symbol::T(token) => self.get_t_name(*token),
            Symbol::NT(var) => self.get_nt_name(*var),
        }
    }

    fn get_str(&self, symbol: &Symbol) -> String {
        match symbol {
            Symbol::Empty | Symbol::End => symbol.to_string(),
            Symbol::T(token) => self.get_t_str(*token),
            Symbol::NT(var) => self.get_nt_name(*var),
        }
    }

    fn get_name_quote(&self, symbol: &Symbol) -> String {
        match symbol {
            Symbol::Empty | Symbol::End => symbol.to_string(),
            Symbol::T(token) => if self.is_symbol_t_fixed(symbol) { format!("{:?}", self.get_t_str(*token)) } else { self.get_t_str(*token) },
            Symbol::NT(var) => self.get_nt_name(*var),
        }
    }
}

impl Default for SymbolTable {
    fn default() -> Self {
        SymbolTable::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn general() {
        let mut st = SymbolTable::new();
        st.extend_nonterminals(["A".to_string(), "NOT_USED".to_string(), "E".to_string()]);
        assert_eq!(st.add_child_nonterminal(0), 3);
        assert_eq!(st.get_str(&Symbol::NT(3)), "A_1");
        assert_eq!(st.add_child_nonterminal(0), 4);
        assert_eq!(st.get_str(&Symbol::NT(4)), "A_2");
        assert_eq!(st.add_child_nonterminal(2), 5);
        assert_eq!(st.get_str(&Symbol::NT(5)), "E_1");
        assert_eq!(st.add_child_nonterminal(1), 6);
        st.remove_nonterminal(1);
        assert_eq!(st.get_str(&Symbol::NT(2)), "A_1");
        assert_eq!(st.get_str(&Symbol::NT(3)), "A_2");
        assert_eq!(st.get_str(&Symbol::NT(4)), "E_1");
        #[cfg(any())]
        assert_eq!(st.find_nonterminal("A_1"), Some(2));
        assert!(st.nt.contains(&"A_2".to_string()));
        #[cfg(any())]
        assert!(st.names.contains_key("A_2"));
        assert!(st.fixer_nt.contains("A_2"));
        st.remove_nt_name(3);
        assert!(!st.nt.contains(&"A_2".to_string()));
        #[cfg(any())]
        assert!(!st.names.contains_key("A_2"));
        assert!(!st.fixer_nt.contains("A_2"));
    }

    #[test]
    fn terminals() {
        let mut st = SymbolTable::new();
        let tnames = vec![("a", Some("aa")), ("b", Some("bb")), ("a", Some("A")), ("c", None), ("d", None)];
        st.extend_terminals(tnames);
        assert_eq!(st.get_num_t(), 5);
        let result = st.get_terminals().map(|(s, v)| (s.as_str(), v.as_ref().map(|s| s.as_str()))).to_vec();
        assert_eq!(result, vec![("a", Some("aa")), ("b", Some("bb")), ("a1", Some("A")), ("c", None), ("d", None)]);
        for name in vec!["a1", "c", "d"] {
            assert_eq!(st.fixer_t.contains(name), true);
        }
        st.downsize_num_t(2);
        let result = st.get_terminals().map(|(s, v)| (s.as_str(), v.as_ref().map(|s| s.as_str()))).to_vec();
        assert_eq!(result, vec![("a", Some("aa")), ("b", Some("bb"))]);
        assert_eq!(st.get_num_t(), 2);
        for name in vec!["a1", "c", "d"] {
            assert_eq!(st.fixer_t.contains(name), false);
        }
        let extra_tnames = vec![("a", Some("A")), ("c", None), ("d", None)];
        st.extend_terminals(extra_tnames);
        let result = st.get_terminals().map(|(s, v)| (s.as_str(), v.as_ref().map(|s| s.as_str()))).to_vec();
        assert_eq!(result, vec![("a", Some("aa")), ("b", Some("bb")), ("a1", Some("A")), ("c", None), ("d", None)]);
    }
}