use std::collections::{BTreeMap, BTreeSet};

use crate::sema::{err::SemaError, types::SemaType};

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SymbolId(pub usize);

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SymbolKind {
    Global,
    Local,
    Param,
}

#[derive(Clone, Debug)]
pub struct Symbol {
    pub id: SymbolId,
    pub name: String,
    pub ty: SemaType,
    pub kind: SymbolKind,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct FunctionId(pub usize);

#[derive(Clone, Debug)]
pub struct FunctionSig {
    pub id: FunctionId,
    pub name: String,
    pub return_type: SemaType,
    pub parameter_types: Vec<SemaType>,
}

pub struct SymbolTable {
    symbols: Vec<Symbol>,
    variable_map: BTreeMap<String, Vec<SymbolId>>,
    scopes: Vec<BTreeSet<String>>,
}

impl SymbolTable {
    pub fn new() -> Self {
        Self {
            symbols: vec![],
            variable_map: BTreeMap::new(),
            scopes: vec![BTreeSet::new()],
        }
    }

    pub fn enter_scope(&mut self) {
        self.scopes.push(BTreeSet::new());
    }

    pub fn exit_scope(&mut self) {
        let variables = self.scopes.pop().unwrap();
        for var in variables {
            self.variable_map.get_mut(&var).unwrap().pop();
        }
    }

    pub fn declare_variable(&mut self, name: &str, kind: SymbolKind, ty: SemaType) -> Result<SymbolId, SemaError> {
        if self.scopes.last().unwrap().contains(name) {
            return Err(SemaError::VariableHasBeenDefined(name.to_string()));
        }
        let id = SymbolId(self.symbols.len());
        self.symbols.push(Symbol {
            id,
            name: name.to_string(),
            ty,
            kind,
        });
        self.variable_map.entry(name.to_string()).or_default().push(id);
        self.scopes.last_mut().unwrap().insert(name.to_string());
        Ok(id)
    }

    pub fn get_variable(&self, name: &str) -> Option<SymbolId> {
        self.variable_map.get(name).and_then(|vars| vars.last()).cloned()
    }

    pub fn get_symbol(&self, id: SymbolId) -> &Symbol {
        &self.symbols[id.0]
    }
}