use axhash_map::AxHashMap;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SymbolKind {
    Variable,
    Function,
    Type,
    Constant,
    Module,
    Other(String),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SourceLocation {
    pub file_id: u32,
    pub line: u32,
    pub column: u32,
}

impl SourceLocation {
    #[inline(always)]
    pub fn new(file_id: u32, line: u32, column: u32) -> Self {
        Self {
            file_id,
            line,
            column,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SymbolData {
    pub id: u32,
    pub kind: SymbolKind,
    pub location: SourceLocation,
    pub is_used: bool,
    pub is_public: bool,
}

impl SymbolData {
    #[inline(always)]
    pub fn new(id: u32, kind: SymbolKind, location: SourceLocation) -> Self {
        Self {
            id,
            kind,
            location,
            is_used: false,
            is_public: false,
        }
    }

    #[inline(always)]
    pub fn public(mut self) -> Self {
        self.is_public = true;
        self
    }

    #[inline(always)]
    pub fn mark_used(&mut self) {
        self.is_used = true;
    }
}

pub struct Scope {
    symbols: AxHashMap<String, SymbolData>,
}

impl Scope {
    #[inline(always)]
    fn with_capacity(capacity: usize) -> Self {
        Self {
            symbols: AxHashMap::with_capacity(capacity),
        }
    }

    #[inline(always)]
    fn define(&mut self, name: String, data: SymbolData) -> Option<SymbolData> {
        self.symbols.insert(name, data)
    }

    #[inline(always)]
    fn get(&self, name: &str) -> Option<&SymbolData> {
        self.symbols.get(name)
    }

    #[inline(always)]
    fn get_mut(&mut self, name: &str) -> Option<&mut SymbolData> {
        self.symbols.get_mut(name)
    }

    #[inline(always)]
    pub fn len(&self) -> usize {
        self.symbols.len()
    }

    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.symbols.is_empty()
    }

    #[inline(always)]
    pub fn iter(&self) -> impl Iterator<Item = (&str, &SymbolData)> {
        self.symbols.iter().map(|(k, v)| (k.as_str(), v))
    }
}

#[derive(Debug)]
pub struct ResolveResult<'sym> {
    pub data: &'sym SymbolData,
    pub depth: usize,
}

pub struct ScopedSymbolTable {
    scopes: Vec<Scope>,

    next_id: u32,

    scope_capacity: usize,
}

impl ScopedSymbolTable {
    #[inline(always)]
    pub fn new() -> Self {
        Self::with_scope_capacity(64)
    }

    #[inline(always)]
    pub fn with_scope_capacity(scope_capacity: usize) -> Self {
        let mut scopes = Vec::with_capacity(16);
        scopes.push(Scope::with_capacity(scope_capacity));
        Self {
            scopes,
            next_id: 0,
            scope_capacity,
        }
    }

    #[inline(always)]
    pub fn with_capacity(max_depth: usize, scope_capacity: usize) -> Self {
        let mut scopes = Vec::with_capacity(max_depth.max(1));
        scopes.push(Scope::with_capacity(scope_capacity));
        Self {
            scopes,
            next_id: 0,
            scope_capacity,
        }
    }

    #[inline(always)]
    pub fn push_scope(&mut self) {
        self.scopes.push(Scope::with_capacity(self.scope_capacity));
    }

    #[inline(always)]
    pub fn pop_scope(&mut self) -> bool {
        if self.scopes.len() > 1 {
            self.scopes.pop();
            true
        } else {
            false
        }
    }

    #[inline(always)]
    pub fn depth(&self) -> usize {
        self.scopes.len() - 1
    }

    #[inline(always)]
    pub fn define(&mut self, name: String, mut data: SymbolData) -> Option<SymbolData> {
        data.id = self.next_id;
        self.next_id += 1;
        self.scopes.last_mut().unwrap().define(name, data)
    }

    #[inline(always)]
    pub fn define_symbol(
        &mut self,
        name: String,
        kind: SymbolKind,
        location: SourceLocation,
    ) -> u32 {
        let id = self.next_id;
        let data = SymbolData::new(id, kind, location);
        self.define(name, data);
        id
    }

    #[inline(always)]
    pub fn resolve(&self, name: &str) -> Option<&SymbolData> {
        for scope in self.scopes.iter().rev() {
            if let Some(sym) = scope.get(name) {
                return Some(sym);
            }
        }
        None
    }

    #[inline(always)]
    pub fn resolve_with_depth(&self, name: &str) -> Option<ResolveResult<'_>> {
        for (depth, scope) in self.scopes.iter().rev().enumerate() {
            if let Some(data) = scope.get(name) {
                return Some(ResolveResult { data, depth });
            }
        }
        None
    }

    #[inline(always)]
    pub fn resolve_mut(&mut self, name: &str) -> Option<&mut SymbolData> {
        for scope in self.scopes.iter_mut().rev() {
            if let Some(sym) = scope.get_mut(name) {
                return Some(sym);
            }
        }
        None
    }

    #[inline(always)]
    pub fn is_defined_locally(&self, name: &str) -> bool {
        self.scopes.last().map_or(false, |s| s.get(name).is_some())
    }

    #[inline(always)]
    pub fn current_scope(&self) -> &Scope {
        self.scopes.last().unwrap()
    }

    #[inline(always)]
    pub fn current_scope_mut(&mut self) -> &mut Scope {
        self.scopes.last_mut().unwrap()
    }

    #[inline(always)]
    pub fn total_symbols(&self) -> usize {
        self.scopes.iter().map(|s| s.len()).sum()
    }

    #[inline(always)]
    pub fn scope_count(&self) -> usize {
        self.scopes.len()
    }

    pub fn unused_symbols(&self) -> Vec<(usize, &str, &SymbolData)> {
        let max_depth = self.scopes.len() - 1;
        let mut result = Vec::new();
        for (i, scope) in self.scopes.iter().enumerate() {
            let depth = max_depth - i;
            for (name, data) in scope.iter() {
                if !data.is_used {
                    result.push((depth, name, data));
                }
            }
        }
        result
    }
}

impl Default for ScopedSymbolTable {
    #[inline(always)]
    fn default() -> Self {
        Self::new()
    }
}

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

    fn loc(line: u32) -> SourceLocation {
        SourceLocation::new(0, line, 1)
    }
    fn var(line: u32) -> SymbolData {
        SymbolData::new(0, SymbolKind::Variable, loc(line))
    }
    fn func(line: u32) -> SymbolData {
        SymbolData::new(0, SymbolKind::Function, loc(line))
    }

    #[test]
    fn global_define_and_resolve() {
        let mut t = ScopedSymbolTable::new();
        t.define("println".to_string(), func(1).public());
        let r = t.resolve("println").unwrap();
        assert_eq!(r.kind, SymbolKind::Function);
        assert!(r.is_public);
    }

    #[test]
    fn resolve_none_for_undefined() {
        let t = ScopedSymbolTable::new();
        assert!(t.resolve("undefined").is_none());
    }

    #[test]
    fn push_pop_scope() {
        let mut t = ScopedSymbolTable::new();
        assert_eq!(t.depth(), 0);
        t.push_scope();
        assert_eq!(t.depth(), 1);
        t.define("x".to_string(), var(5));
        assert!(t.resolve("x").is_some());
        t.pop_scope();
        assert_eq!(t.depth(), 0);
        assert!(t.resolve("x").is_none());
    }

    #[test]
    fn pop_global_scope_is_noop() {
        let mut t = ScopedSymbolTable::new();
        assert!(!t.pop_scope());
        assert_eq!(t.depth(), 0);
    }

    #[test]
    fn inner_scope_sees_outer_symbols() {
        let mut t = ScopedSymbolTable::new();
        t.define("global_fn".to_string(), func(1).public());
        t.push_scope();
        t.define("arg".to_string(), var(3));
        t.push_scope();
        assert!(t.resolve("global_fn").is_some());
        assert!(t.resolve("arg").is_some());
    }

    #[test]
    fn outer_scope_does_not_see_inner_symbols() {
        let mut t = ScopedSymbolTable::new();
        t.push_scope();
        t.define("inner".to_string(), var(2));
        t.pop_scope();
        assert!(t.resolve("inner").is_none());
    }

    #[test]
    fn shadowing_inner_scope() {
        let mut t = ScopedSymbolTable::new();
        t.define("x".to_string(), func(1));
        t.push_scope();
        t.define("x".to_string(), var(3));
        assert_eq!(t.resolve("x").unwrap().kind, SymbolKind::Variable);
        t.pop_scope();
        assert_eq!(t.resolve("x").unwrap().kind, SymbolKind::Function);
    }

    #[test]
    fn resolve_with_depth_current() {
        let mut t = ScopedSymbolTable::new();
        t.push_scope();
        t.define("local".to_string(), var(1));
        assert_eq!(t.resolve_with_depth("local").unwrap().depth, 0);
    }

    #[test]
    fn resolve_with_depth_outer() {
        let mut t = ScopedSymbolTable::new();
        t.define("outer".to_string(), func(1));
        t.push_scope();
        t.push_scope();
        assert_eq!(t.resolve_with_depth("outer").unwrap().depth, 2);
    }

    #[test]
    fn is_defined_locally() {
        let mut t = ScopedSymbolTable::new();
        t.define("outer".to_string(), func(1));
        t.push_scope();
        assert!(!t.is_defined_locally("outer"));
        t.define("local".to_string(), var(2));
        assert!(t.is_defined_locally("local"));
    }

    #[test]
    fn resolve_mut_marks_used() {
        let mut t = ScopedSymbolTable::new();
        t.define("x".to_string(), var(1));
        assert!(!t.resolve("x").unwrap().is_used);
        t.resolve_mut("x").unwrap().mark_used();
        assert!(t.resolve("x").unwrap().is_used);
    }

    #[test]
    fn unused_symbols_report() {
        let mut t = ScopedSymbolTable::new();
        t.define("used_fn".to_string(), func(1));
        t.define("unused_fn".to_string(), func(2));
        t.resolve_mut("used_fn").unwrap().mark_used();
        let unused = t.unused_symbols();
        assert_eq!(unused.len(), 1);
        assert_eq!(unused[0].1, "unused_fn");
    }

    #[test]
    fn ids_are_unique_and_monotonic() {
        let mut t = ScopedSymbolTable::new();
        let id0 = t.define_symbol("a".to_string(), SymbolKind::Variable, loc(1));
        let id1 = t.define_symbol("b".to_string(), SymbolKind::Variable, loc(2));
        let id2 = t.define_symbol("c".to_string(), SymbolKind::Function, loc(3));
        assert!(id0 < id1 && id1 < id2);
        assert_eq!(t.total_symbols(), 3);
    }

    #[test]
    fn with_capacity_constructor() {
        let mut t = ScopedSymbolTable::with_capacity(32, 128);
        for i in 0..128u32 {
            t.define(format!("sym_{i}"), var(i));
        }
        assert_eq!(t.total_symbols(), 128);
    }

    #[test]
    fn deep_nesting_resolve() {
        let mut t = ScopedSymbolTable::new();
        t.define("root".to_string(), func(1));
        for _ in 0..50 {
            t.push_scope();
        }
        let r = t.resolve_with_depth("root").unwrap();
        assert_eq!(r.depth, 50);
        assert_eq!(r.data.kind, SymbolKind::Function);
    }

    #[test]
    fn redeclaration_returns_previous() {
        let mut t = ScopedSymbolTable::new();
        t.define("x".to_string(), func(1));
        let prev = t.define("x".to_string(), var(2));
        assert!(prev.is_some());
        assert_eq!(prev.unwrap().kind, SymbolKind::Function);
        assert_eq!(t.resolve("x").unwrap().kind, SymbolKind::Variable);
    }
}
