#![doc = include_str!("../README.md")]

use axhash_map::AxHashMap;

// ── SymbolKind ────────────────────────────────────────────────────────────────

/// Coarse-grained category of a symbol, as tracked by a typical compiler
/// front-end.  Extend or replace this with your own domain types.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SymbolKind {
    /// A local variable binding.
    Variable,
    /// A function or method definition.
    Function,
    /// A type alias, struct, enum, or trait.
    Type,
    /// A compile-time constant.
    Constant,
    /// A module or namespace boundary.
    Module,
    /// Any other, domain-specific symbol kind.
    Other(String),
}

// ── SourceLocation ────────────────────────────────────────────────────────────

/// Compact source-location record attached to every symbol.
///
/// Line and column numbers are 1-based following the LSP convention.
/// `file_id` is an opaque index into a file table owned by the caller — keeping
/// it as a `u32` avoids heap allocation inside the symbol table.
#[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,
        }
    }
}

// ── SymbolData ────────────────────────────────────────────────────────────────

/// Metadata record stored for every symbol in the table.
///
/// `SymbolData` is intentionally **not** generic over the value type so that
/// the `resolve` hot path stays monomorphic and cache-friendly.  If you need
/// richer per-symbol payloads, store them in a side table indexed by `id`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SymbolData {
    /// Unique numeric id — useful as a key into type-inference arenas.
    pub id: u32,
    /// Coarse kind tag (variable, function, type, …).
    pub kind: SymbolKind,
    /// Where the symbol was first declared.
    pub location: SourceLocation,
    /// Whether the symbol has been referenced at least once.
    /// Compilers use this flag to emit "unused variable" warnings.
    pub is_used: bool,
    /// Whether the symbol is externally visible (e.g. `pub` in Rust).
    pub is_public: bool,
}

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

    /// Mark the symbol as publicly exported.
    #[inline(always)]
    pub fn public(mut self) -> Self {
        self.is_public = true;
        self
    }

    /// Mark the symbol as having been referenced at least once.
    #[inline(always)]
    pub fn mark_used(&mut self) {
        self.is_used = true;
    }
}

// ── Scope ─────────────────────────────────────────────────────────────────────

/// A single lexical scope — a flat mapping from identifier to symbol data.
///
/// Backed by [`AxHashMap`] for O(1) average-case insert and lookup with
/// minimal hashing overhead (AES-NI one-shot hash per key).
pub struct Scope {
    symbols: AxHashMap<String, SymbolData>,
}

impl Scope {
    /// Creates an empty scope with a pre-allocated bucket table.
    #[inline(always)]
    fn with_capacity(capacity: usize) -> Self {
        Self {
            symbols: AxHashMap::with_capacity(capacity),
        }
    }

    /// Inserts `data` under `name`.  Returns the previous binding if the name
    /// was already defined in **this** scope.
    #[inline(always)]
    fn define(&mut self, name: String, data: SymbolData) -> Option<SymbolData> {
        self.symbols.insert(name, data)
    }

    /// Looks up `name` in this scope only.
    #[inline(always)]
    fn get(&self, name: &str) -> Option<&SymbolData> {
        self.symbols.get(name)
    }

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

    /// Number of symbols in this scope.
    #[inline(always)]
    pub fn len(&self) -> usize {
        self.symbols.len()
    }

    /// True if the scope has no symbols.
    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.symbols.is_empty()
    }

    /// Iterate over every `(name, data)` pair in this scope.
    #[inline(always)]
    pub fn iter(&self) -> impl Iterator<Item = (&str, &SymbolData)> {
        self.symbols.iter().map(|(k, v)| (k.as_str(), v))
    }
}

// ── ResolveResult ─────────────────────────────────────────────────────────────

/// Carries a reference to the symbol and the **depth** at which it was found.
///
/// Depth 0 = innermost scope (current block), higher = further up the stack.
/// Useful for closure capture analysis and de-Bruijn index generation.
#[derive(Debug)]
pub struct ResolveResult<'sym> {
    pub data: &'sym SymbolData,
    pub depth: usize,
}

// ── ScopedSymbolTable ─────────────────────────────────────────────────────────

/// High-performance scoped symbol table for compiler front-ends.
///
/// # Architecture
///
/// Internally a `Vec<Scope>` is used as a scope stack.  Each `Scope` is a
/// flat [`AxHashMap`] — one AES-NI hash per lookup, no indirection.
/// `resolve` walks the stack from top (innermost) to bottom (global) and
/// returns on the first hit, matching standard lexical scoping.
///
/// ```text
/// ┌──────────────────────────────────┐
/// │  ScopedSymbolTable               │
/// │  ┌────────────────────────────┐  │
/// │  │ scope[2]  (innermost)      │  │  ← push_scope / pop_scope
/// │  │  "x" → SymbolData { … }   │  │
/// │  ├────────────────────────────┤  │
/// │  │ scope[1]  (function body)  │  │
/// │  │  "n" → SymbolData { … }   │  │
/// │  ├────────────────────────────┤  │
/// │  │ scope[0]  (global / module)│  │
/// │  │  "println" → …            │  │
/// │  └────────────────────────────┘  │
/// └──────────────────────────────────┘
/// ```
///
/// # Performance characteristics
///
/// | Operation | Complexity |
/// |---|---|
/// | `define` | O(1) amortised — single AxHashMap insert |
/// | `resolve` at depth 0 | O(1) — single hash probe |
/// | `resolve` at depth d | O(d) — d hash probes |
/// | `push_scope` | O(1) amortised |
/// | `pop_scope` | O(n) where n = symbols in top scope |
///
/// # Example
///
/// ```rust
/// use ax_sym::{ScopedSymbolTable, SymbolData, SymbolKind, SourceLocation};
///
/// let mut table = ScopedSymbolTable::new();
///
/// let loc = SourceLocation::new(0, 1, 1);
/// table.define("println".to_string(),
///     SymbolData::new(0, SymbolKind::Function, loc).public());
///
/// table.push_scope();
/// table.define("x".to_string(), SymbolData::new(1, SymbolKind::Variable, loc));
///
/// assert!(table.resolve("x").is_some());
/// assert!(table.resolve("println").is_some());
/// assert!(table.resolve("y").is_none());
///
/// table.pop_scope();
/// assert!(table.resolve("x").is_none());
/// ```
pub struct ScopedSymbolTable {
    /// The scope stack.  `scopes[0]` is always the global scope.
    scopes: Vec<Scope>,
    /// Monotonically increasing counter for auto-assigned symbol IDs.
    next_id: u32,
    /// Default per-scope bucket count to pre-allocate on `push_scope`.
    scope_capacity: usize,
}

impl ScopedSymbolTable {
    // ── Constructors ───────────────────────────────────────────────────────

    /// Creates a new symbol table with a single global scope (capacity = 64).
    #[inline(always)]
    pub fn new() -> Self {
        Self::with_scope_capacity(64)
    }

    /// Creates a new symbol table where each scope is pre-allocated for
    /// `scope_capacity` symbols.  Tune to the average symbols-per-block in
    /// your language to eliminate per-scope rehashing.
    #[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,
        }
    }

    /// Creates a new symbol table pre-allocated to hold at least `max_depth`
    /// scopes without reallocating the outer `Vec`.
    #[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,
        }
    }

    // ── Scope management ───────────────────────────────────────────────────

    /// Pushes a new empty scope, entering a deeper nesting level.
    #[inline(always)]
    pub fn push_scope(&mut self) {
        self.scopes.push(Scope::with_capacity(self.scope_capacity));
    }

    /// Pops the innermost scope, discarding all symbols in it.
    /// Returns `false` (no-op) when the global scope is reached.
    #[inline(always)]
    pub fn pop_scope(&mut self) -> bool {
        if self.scopes.len() > 1 {
            self.scopes.pop();
            true
        } else {
            false
        }
    }

    /// Current nesting depth.  0 = global scope only.
    #[inline(always)]
    pub fn depth(&self) -> usize {
        self.scopes.len() - 1
    }

    // ── Symbol operations ──────────────────────────────────────────────────

    /// Defines a symbol in the **innermost** scope, auto-assigning a unique id.
    ///
    /// Returns the previous binding if the name was already defined in the
    /// **same** scope (re-declaration error in most languages).
    #[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)
    }

    /// Shorthand for `define` — constructs `SymbolData` inline and returns the
    /// assigned `id`.
    #[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
    }

    /// Resolves `name` by searching innermost → outermost scope.
    ///
    /// This is the **hot path** in type-checking / name-resolution passes.
    #[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
    }

    /// Like `resolve`, but also returns the depth at which the symbol was found.
    ///
    /// Depth 0 = innermost scope, depth 1 = enclosing, etc.
    #[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
    }

    /// Mutable resolve — used to update symbol metadata (e.g. mark as used,
    /// attach inferred type information).
    #[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
    }

    /// Returns `true` if `name` is defined in the innermost scope only.
    /// Use this to detect same-scope re-declarations.
    #[inline(always)]
    pub fn is_defined_locally(&self, name: &str) -> bool {
        self.scopes.last().map_or(false, |s| s.get(name).is_some())
    }

    /// Returns a reference to the current (innermost) scope.
    #[inline(always)]
    pub fn current_scope(&self) -> &Scope {
        self.scopes.last().unwrap()
    }

    /// Returns a mutable reference to the current (innermost) scope.
    #[inline(always)]
    pub fn current_scope_mut(&mut self) -> &mut Scope {
        self.scopes.last_mut().unwrap()
    }

    /// Total number of symbols across all active scopes.
    #[inline(always)]
    pub fn total_symbols(&self) -> usize {
        self.scopes.iter().map(|s| s.len()).sum()
    }

    /// Number of active scopes, including the global scope.
    #[inline(always)]
    pub fn scope_count(&self) -> usize {
        self.scopes.len()
    }

    // ── Diagnostics ────────────────────────────────────────────────────────

    /// Returns all symbols that were defined but never marked as `is_used`.
    ///
    /// Yields `(depth, name, &SymbolData)` triples from innermost to outermost.
    /// Filter by `depth > 0` to skip the global scope.
    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()
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[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);
    }
}
