lux/src/symbol_table.rs

//! Symbol Table for Lux
//!
//! Provides semantic analysis infrastructure for IDE features like
//! go-to-definition, find references, and rename refactoring.

use crate::ast::*;
use std::collections::HashMap;

/// Unique identifier for a symbol
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SymbolId(pub u32);

/// Kind of symbol
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SymbolKind {
    Function,
    Variable,
    Parameter,
    Type,
    TypeParameter,
    Variant,
    Effect,
    EffectOperation,
    Field,
    Module,
}

/// A symbol definition
#[derive(Debug, Clone)]
pub struct Symbol {
    pub id: SymbolId,
    pub name: String,
    pub kind: SymbolKind,
    pub span: Span,
    /// Type signature (for display)
    pub type_signature: Option<String>,
    /// Documentation comment
    pub documentation: Option<String>,
    /// Parent symbol (e.g., type for variants, effect for operations)
    pub parent: Option<SymbolId>,
    /// Is this symbol exported (public)?
    pub is_public: bool,
}

/// A reference to a symbol
#[derive(Debug, Clone)]
pub struct Reference {
    pub symbol_id: SymbolId,
    pub span: Span,
    pub is_definition: bool,
    pub is_write: bool,
}

/// A scope in the symbol table
#[derive(Debug, Clone)]
pub struct Scope {
    /// Parent scope (None for global scope)
    pub parent: Option<usize>,
    /// Symbols defined in this scope
    pub symbols: HashMap<String, SymbolId>,
    /// Span of this scope
    pub span: Span,
}

/// The symbol table
#[derive(Debug, Clone)]
pub struct SymbolTable {
    /// All symbols
    symbols: Vec<Symbol>,
    /// All references
    references: Vec<Reference>,
    /// Scopes (index 0 is always the global scope)
    scopes: Vec<Scope>,
    /// Mapping from position to references
    position_to_reference: HashMap<(u32, u32), usize>,
    /// Next symbol ID
    next_id: u32,
}

impl SymbolTable {
    pub fn new() -> Self {
        Self {
            symbols: Vec::new(),
            references: Vec::new(),
            scopes: vec![Scope {
                parent: None,
                symbols: HashMap::new(),
                span: Span { start: 0, end: 0 },
            }],
            position_to_reference: HashMap::new(),
            next_id: 0,
        }
    }

    /// Build symbol table from a program
    pub fn build(program: &Program) -> Self {
        let mut table = Self::new();
        table.visit_program(program);
        table
    }

    /// Add a symbol to the current scope
    fn add_symbol(&mut self, scope_idx: usize, symbol: Symbol) -> SymbolId {
        let id = symbol.id;
        self.scopes[scope_idx].symbols.insert(symbol.name.clone(), id);
        self.symbols.push(symbol);
        id
    }

    /// Create a new symbol
    fn new_symbol(
        &mut self,
        name: String,
        kind: SymbolKind,
        span: Span,
        type_signature: Option<String>,
        is_public: bool,
    ) -> Symbol {
        let id = SymbolId(self.next_id);
        self.next_id += 1;
        Symbol {
            id,
            name,
            kind,
            span,
            type_signature,
            documentation: None,
            parent: None,
            is_public,
        }
    }

    /// Add a reference
    fn add_reference(&mut self, symbol_id: SymbolId, span: Span, is_definition: bool, is_write: bool) {
        let ref_idx = self.references.len();
        self.references.push(Reference {
            symbol_id,
            span,
            is_definition,
            is_write,
        });
        // Index by start position
        self.position_to_reference.insert((span.start as u32, span.end as u32), ref_idx);
    }

    /// Look up a symbol by name in the given scope and its parents
    pub fn lookup(&self, name: &str, scope_idx: usize) -> Option<SymbolId> {
        let scope = &self.scopes[scope_idx];
        if let Some(&id) = scope.symbols.get(name) {
            return Some(id);
        }
        if let Some(parent) = scope.parent {
            return self.lookup(name, parent);
        }
        None
    }

    /// Get a symbol by ID
    pub fn get_symbol(&self, id: SymbolId) -> Option<&Symbol> {
        self.symbols.iter().find(|s| s.id == id)
    }

    /// Get the symbol at a position
    pub fn symbol_at_position(&self, offset: usize) -> Option<&Symbol> {
        // Find a reference that contains this offset
        for reference in &self.references {
            if offset >= reference.span.start && offset <= reference.span.end {
                return self.get_symbol(reference.symbol_id);
            }
        }
        None
    }

    /// Get the definition of a symbol at a position
    pub fn definition_at_position(&self, offset: usize) -> Option<&Symbol> {
        self.symbol_at_position(offset)
    }

    /// Find all references to a symbol
    pub fn find_references(&self, symbol_id: SymbolId) -> Vec<&Reference> {
        self.references
            .iter()
            .filter(|r| r.symbol_id == symbol_id)
            .collect()
    }

    /// Get all symbols of a given kind
    pub fn symbols_of_kind(&self, kind: SymbolKind) -> Vec<&Symbol> {
        self.symbols.iter().filter(|s| s.kind == kind).collect()
    }

    /// Get all symbols in the global scope
    pub fn global_symbols(&self) -> Vec<&Symbol> {
        self.scopes[0]
            .symbols
            .values()
            .filter_map(|&id| self.get_symbol(id))
            .collect()
    }

    /// Create a new scope
    fn push_scope(&mut self, parent: usize, span: Span) -> usize {
        let idx = self.scopes.len();
        self.scopes.push(Scope {
            parent: Some(parent),
            symbols: HashMap::new(),
            span,
        });
        idx
    }

    // =========================================================================
    // AST Visitors
    // =========================================================================

    fn visit_program(&mut self, program: &Program) {
        // First pass: collect all top-level declarations
        for decl in &program.declarations {
            self.visit_declaration(decl, 0);
        }
    }

    fn visit_declaration(&mut self, decl: &Declaration, scope_idx: usize) {
        match decl {
            Declaration::Function(f) => self.visit_function(f, scope_idx),
            Declaration::Type(t) => self.visit_type_decl(t, scope_idx),
            Declaration::Effect(e) => self.visit_effect(e, scope_idx),
            Declaration::Let(let_decl) => {
                let is_public = matches!(let_decl.visibility, Visibility::Public);
                let type_sig = let_decl.typ.as_ref().map(|t| self.type_expr_to_string(t));
                let mut symbol = self.new_symbol(
                    let_decl.name.name.clone(),
                    SymbolKind::Variable,
                    let_decl.span,
                    type_sig,
                    is_public,
                );
                symbol.documentation = let_decl.doc.clone();
                let id = self.add_symbol(scope_idx, symbol);
                self.add_reference(id, let_decl.name.span, true, true);

                // Visit the expression
                self.visit_expr(&let_decl.value, scope_idx);
            }
            Declaration::Handler(h) => self.visit_handler(h, scope_idx),
            Declaration::Trait(t) => self.visit_trait(t, scope_idx),
            Declaration::Impl(i) => self.visit_impl(i, scope_idx),
        }
    }

    fn visit_function(&mut self, f: &FunctionDecl, scope_idx: usize) {
        let is_public = matches!(f.visibility, Visibility::Public);

        // Build type signature
        let param_types: Vec<String> = f.params.iter()
            .map(|p| format!("{}: {}", p.name.name, self.type_expr_to_string(&p.typ)))
            .collect();
        let return_type = self.type_expr_to_string(&f.return_type);
        let effects = if f.effects.is_empty() {
            String::new()
        } else {
            format!(" with {{{}}}", f.effects.iter()
                .map(|e| e.name.clone())
                .collect::<Vec<_>>()
                .join(", "))
        };
        let properties = if f.properties.is_empty() {
            String::new()
        } else {
            format!(" is {}", f.properties.iter()
                .map(|p| match p {
                    crate::ast::BehavioralProperty::Pure => "pure",
                    crate::ast::BehavioralProperty::Total => "total",
                    crate::ast::BehavioralProperty::Idempotent => "idempotent",
                    crate::ast::BehavioralProperty::Deterministic => "deterministic",
                    crate::ast::BehavioralProperty::Commutative => "commutative",
                })
                .collect::<Vec<_>>()
                .join(", "))
        };
        let type_sig = format!("fn {}({}): {}{}{}", f.name.name, param_types.join(", "), return_type, properties, effects);

        let mut symbol = self.new_symbol(
            f.name.name.clone(),
            SymbolKind::Function,
            f.name.span,
            Some(type_sig),
            is_public,
        );
        symbol.documentation = f.doc.clone();
        let fn_id = self.add_symbol(scope_idx, symbol);
        self.add_reference(fn_id, f.name.span, true, false);

        // Create scope for function body
        let body_span = f.body.span();
        let fn_scope = self.push_scope(scope_idx, body_span);

        // Add type parameters
        for tp in &f.type_params {
            let symbol = self.new_symbol(
                tp.name.clone(),
                SymbolKind::TypeParameter,
                tp.span,
                None,
                false,
            );
            self.add_symbol(fn_scope, symbol);
        }

        // Add parameters
        for param in &f.params {
            let type_sig = self.type_expr_to_string(&param.typ);
            let symbol = self.new_symbol(
                param.name.name.clone(),
                SymbolKind::Parameter,
                param.name.span,
                Some(type_sig),
                false,
            );
            self.add_symbol(fn_scope, symbol);
        }

        // Visit body
        self.visit_expr(&f.body, fn_scope);
    }

    fn visit_type_decl(&mut self, t: &TypeDecl, scope_idx: usize) {
        let is_public = matches!(t.visibility, Visibility::Public);
        let type_sig = format!("type {}", t.name.name);

        let mut symbol = self.new_symbol(
            t.name.name.clone(),
            SymbolKind::Type,
            t.name.span,
            Some(type_sig),
            is_public,
        );
        symbol.documentation = t.doc.clone();
        let type_id = self.add_symbol(scope_idx, symbol);
        self.add_reference(type_id, t.name.span, true, false);

        // Add variants
        match &t.definition {
            TypeDef::Enum(variants) => {
                for variant in variants {
                    let mut var_symbol = self.new_symbol(
                        variant.name.name.clone(),
                        SymbolKind::Variant,
                        variant.name.span,
                        None,
                        is_public,
                    );
                    var_symbol.parent = Some(type_id);
                    self.add_symbol(scope_idx, var_symbol);
                }
            }
            TypeDef::Record(fields) => {
                for field in fields {
                    let mut field_symbol = self.new_symbol(
                        field.name.name.clone(),
                        SymbolKind::Field,
                        field.name.span,
                        Some(self.type_expr_to_string(&field.typ)),
                        is_public,
                    );
                    field_symbol.parent = Some(type_id);
                    self.add_symbol(scope_idx, field_symbol);
                }
            }
            TypeDef::Alias(_) => {}
        }
    }

    fn visit_effect(&mut self, e: &EffectDecl, scope_idx: usize) {
        let is_public = true; // Effects are typically public
        let type_sig = format!("effect {}", e.name.name);

        let mut symbol = self.new_symbol(
            e.name.name.clone(),
            SymbolKind::Effect,
            e.name.span,
            Some(type_sig),
            is_public,
        );
        symbol.documentation = e.doc.clone();
        let effect_id = self.add_symbol(scope_idx, symbol);

        // Add operations
        for op in &e.operations {
            let param_types: Vec<String> = op.params.iter()
                .map(|p| format!("{}: {}", p.name.name, self.type_expr_to_string(&p.typ)))
                .collect();
            let return_type = self.type_expr_to_string(&op.return_type);
            let op_sig = format!("fn {}({}): {}", op.name.name, param_types.join(", "), return_type);

            let mut op_symbol = self.new_symbol(
                op.name.name.clone(),
                SymbolKind::EffectOperation,
                op.name.span,
                Some(op_sig),
                is_public,
            );
            op_symbol.parent = Some(effect_id);
            self.add_symbol(scope_idx, op_symbol);
        }
    }

    fn visit_handler(&mut self, _h: &HandlerDecl, _scope_idx: usize) {
        // Handlers are complex - visit their implementations
    }

    fn visit_trait(&mut self, t: &TraitDecl, scope_idx: usize) {
        let is_public = matches!(t.visibility, Visibility::Public);
        let type_sig = format!("trait {}", t.name.name);

        let mut symbol = self.new_symbol(
            t.name.name.clone(),
            SymbolKind::Type, // Traits are like types
            t.name.span,
            Some(type_sig),
            is_public,
        );
        symbol.documentation = t.doc.clone();
        self.add_symbol(scope_idx, symbol);
    }

    fn visit_impl(&mut self, _i: &ImplDecl, _scope_idx: usize) {
        // Impl blocks add methods to types
    }

    fn visit_expr(&mut self, expr: &Expr, scope_idx: usize) {
        match expr {
            Expr::Var(ident) => {
                // Look up the identifier and add a reference
                if let Some(id) = self.lookup(&ident.name, scope_idx) {
                    self.add_reference(id, ident.span, false, false);
                }
            }
            Expr::Let { name, value, body, span, .. } => {
                // Visit the value first
                self.visit_expr(value, scope_idx);

                // Create a new scope for the let binding
                let let_scope = self.push_scope(scope_idx, *span);

                // Add the variable
                let symbol = self.new_symbol(
                    name.name.clone(),
                    SymbolKind::Variable,
                    name.span,
                    None,
                    false,
                );
                let var_id = self.add_symbol(let_scope, symbol);
                self.add_reference(var_id, name.span, true, true);

                // Visit the body
                self.visit_expr(body, let_scope);
            }
            Expr::Lambda { params, body, span, .. } => {
                let lambda_scope = self.push_scope(scope_idx, *span);

                for param in params {
                    let symbol = self.new_symbol(
                        param.name.name.clone(),
                        SymbolKind::Parameter,
                        param.name.span,
                        None,
                        false,
                    );
                    self.add_symbol(lambda_scope, symbol);
                }

                self.visit_expr(body, lambda_scope);
            }
            Expr::Call { func, args, .. } => {
                self.visit_expr(func, scope_idx);
                for arg in args {
                    self.visit_expr(arg, scope_idx);
                }
            }
            Expr::EffectOp { args, .. } => {
                for arg in args {
                    self.visit_expr(arg, scope_idx);
                }
            }
            Expr::Field { object, .. } | Expr::TupleIndex { object, .. } => {
                self.visit_expr(object, scope_idx);
            }
            Expr::If { condition, then_branch, else_branch, .. } => {
                self.visit_expr(condition, scope_idx);
                self.visit_expr(then_branch, scope_idx);
                self.visit_expr(else_branch, scope_idx);
            }
            Expr::Match { scrutinee, arms, .. } => {
                self.visit_expr(scrutinee, scope_idx);
                for arm in arms {
                    // Each arm may bind variables
                    let arm_scope = self.push_scope(scope_idx, arm.body.span());
                    self.visit_pattern(&arm.pattern, arm_scope);
                    if let Some(ref guard) = arm.guard {
                        self.visit_expr(guard, arm_scope);
                    }
                    self.visit_expr(&arm.body, arm_scope);
                }
            }
            Expr::Block { statements, result, .. } => {
                for stmt in statements {
                    self.visit_statement(stmt, scope_idx);
                }
                self.visit_expr(result, scope_idx);
            }
            Expr::BinaryOp { left, right, .. } => {
                self.visit_expr(left, scope_idx);
                self.visit_expr(right, scope_idx);
            }
            Expr::UnaryOp { operand, .. } => {
                self.visit_expr(operand, scope_idx);
            }
            Expr::List { elements, .. } => {
                for e in elements {
                    self.visit_expr(e, scope_idx);
                }
            }
            Expr::Tuple { elements, .. } => {
                for e in elements {
                    self.visit_expr(e, scope_idx);
                }
            }
            Expr::Record { fields, .. } => {
                for (_, e) in fields {
                    self.visit_expr(e, scope_idx);
                }
            }
            Expr::Run { expr, handlers, .. } => {
                self.visit_expr(expr, scope_idx);
                for (_effect, handler_expr) in handlers {
                    self.visit_expr(handler_expr, scope_idx);
                }
            }
            Expr::Resume { value, .. } => {
                self.visit_expr(value, scope_idx);
            }
            // Literals don't need symbol resolution
            Expr::Literal(_) => {}
        }
    }

    fn visit_statement(&mut self, stmt: &Statement, scope_idx: usize) {
        match stmt {
            Statement::Expr(e) => self.visit_expr(e, scope_idx),
            Statement::Let { name, value, .. } => {
                self.visit_expr(value, scope_idx);
                let symbol = self.new_symbol(
                    name.name.clone(),
                    SymbolKind::Variable,
                    name.span,
                    None,
                    false,
                );
                let id = self.add_symbol(scope_idx, symbol);
                self.add_reference(id, name.span, true, true);
            }
        }
    }

    fn visit_pattern(&mut self, pattern: &Pattern, scope_idx: usize) {
        match pattern {
            Pattern::Var(ident) => {
                let symbol = self.new_symbol(
                    ident.name.clone(),
                    SymbolKind::Variable,
                    ident.span,
                    None,
                    false,
                );
                let id = self.add_symbol(scope_idx, symbol);
                self.add_reference(id, ident.span, true, true);
            }
            Pattern::Constructor { fields, .. } => {
                for p in fields {
                    self.visit_pattern(p, scope_idx);
                }
            }
            Pattern::Tuple { elements, .. } => {
                for p in elements {
                    self.visit_pattern(p, scope_idx);
                }
            }
            Pattern::Record { fields, .. } => {
                for (_, p) in fields {
                    self.visit_pattern(p, scope_idx);
                }
            }
            Pattern::Wildcard(_) => {}
            Pattern::Literal(_) => {}
        }
    }

    fn type_expr_to_string(&self, typ: &TypeExpr) -> String {
        match typ {
            TypeExpr::Named(ident) => ident.name.clone(),
            TypeExpr::App(base, args) => {
                let base_str = self.type_expr_to_string(base);
                if args.is_empty() {
                    base_str
                } else {
                    let args_str: Vec<String> = args.iter()
                        .map(|a| self.type_expr_to_string(a))
                        .collect();
                    format!("{}<{}>", base_str, args_str.join(", "))
                }
            }
            TypeExpr::Function { params, return_type, .. } => {
                let params_str: Vec<String> = params.iter()
                    .map(|p| self.type_expr_to_string(p))
                    .collect();
                format!("fn({}): {}", params_str.join(", "), self.type_expr_to_string(return_type))
            }
            TypeExpr::Tuple(types) => {
                let types_str: Vec<String> = types.iter()
                    .map(|t| self.type_expr_to_string(t))
                    .collect();
                format!("({})", types_str.join(", "))
            }
            TypeExpr::Record(fields) => {
                let fields_str: Vec<String> = fields.iter()
                    .map(|f| format!("{}: {}", f.name, self.type_expr_to_string(&f.typ)))
                    .collect();
                format!("{{ {} }}", fields_str.join(", "))
            }
            TypeExpr::Unit => "Unit".to_string(),
            TypeExpr::Versioned { base, .. } => {
                format!("{}@versioned", self.type_expr_to_string(base))
            }
        }
    }
}

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

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

    #[test]
    fn test_symbol_table_basic() {
        let source = r#"
            fn add(a: Int, b: Int): Int = a + b
            let x = 42
        "#;

        let program = Parser::parse_source(source).unwrap();
        let table = SymbolTable::build(&program);

        // Should have add function and x variable
        let globals = table.global_symbols();
        assert!(globals.iter().any(|s| s.name == "add"));
        assert!(globals.iter().any(|s| s.name == "x"));
    }

    #[test]
    fn test_symbol_lookup() {
        let source = r#"
            fn foo(x: Int): Int = x + 1
        "#;

        let program = Parser::parse_source(source).unwrap();
        let table = SymbolTable::build(&program);

        // Should be able to find foo
        assert!(table.lookup("foo", 0).is_some());
    }
}