use crate::common::diagnostics::{Diagnostic, Severity};
use crate::frontends::pbs::ast::{AstArena, NodeKind};
use crate::analysis::symbols::{Symbol, SymbolArena, SymbolKind, DefIndex, DefKey, Namespace, NodeToSymbol, RefIndex};
use prometeu_analysis::{NameInterner, NodeId, ModuleId};

pub fn build_def_index(
    arena: &AstArena,
    module: ModuleId,
    _interner: &NameInterner,
    imports: Option<(&SymbolArena, &DefIndex)>,
) -> (SymbolArena, DefIndex, RefIndex, NodeToSymbol, Vec<Diagnostic>) {
    let mut symbols = SymbolArena::new();
    let mut index = DefIndex::new();
    let mut ref_index = RefIndex::new();
    let mut node_to_symbol = NodeToSymbol::new();
    let mut diagnostics = Vec::new();

    // Passo 1: Coletar definições top-level
    for &root_id in &arena.roots {
        let root_kind = arena.kind(root_id);
        if let NodeKind::File(file_node) = root_kind {
            for &decl_id in &file_node.decls {
                let decl_kind = arena.kind(decl_id);
                let span = arena.span(decl_id);

                let result = match decl_kind {
                    NodeKind::FnDecl(fn_decl) => {
                        Some((fn_decl.name, SymbolKind::Function, Namespace::Value, fn_decl.vis.as_deref() == Some("pub")))
                    }
                    NodeKind::TypeDecl(type_decl) => {
                        let kind = match type_decl.type_kind.as_str() {
                            "struct" => SymbolKind::Struct,
                            "contract" => SymbolKind::Contract,
                            "error" => SymbolKind::ErrorType,
                            _ => SymbolKind::Type,
                        };
                        Some((type_decl.name, kind, Namespace::Type, type_decl.vis.as_deref() == Some("pub")))
                    }
                    NodeKind::ServiceDecl(service_decl) => {
                        Some((service_decl.name, SymbolKind::Service, Namespace::Service, service_decl.vis.as_deref() == Some("pub")))
                    }
                    _ => None,
                };

                if let Some((name, kind, namespace, exported)) = result {
                    let symbol = Symbol {
                        name,
                        kind,
                        exported,
                        module,
                        decl_span: span.clone(),
                    };

                    let symbol_id = symbols.insert(symbol);
                    node_to_symbol.bind_node(decl_id, symbol_id);
                    let key = DefKey {
                        module,
                        name,
                        namespace,
                    };

                    if let Err(mut diag) = index.insert_symbol(key, symbol_id) {
                        diag.span = span;
                        diagnostics.push(diag);
                    }
                }
            }
        }
    }

    // Passo 2: Resolver referências (Identifiers)
    for &root_id in &arena.roots {
        walk_node(root_id, arena, module, &index, imports, &mut ref_index, &mut node_to_symbol, &mut diagnostics);
    }

    (symbols, index, ref_index, node_to_symbol, diagnostics)
}

fn walk_node(
    node_id: NodeId,
    arena: &AstArena,
    module: ModuleId,
    index: &DefIndex,
    imports: Option<(&SymbolArena, &DefIndex)>,
    ref_index: &mut RefIndex,
    node_to_symbol: &mut NodeToSymbol,
    diagnostics: &mut Vec<Diagnostic>,
) {
    let kind = arena.kind(node_id);
    let span = arena.span(node_id);

    match kind {
        NodeKind::File(file) => {
            for &decl in &file.decls {
                walk_node(decl, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            }
        }
        NodeKind::FnDecl(fn_decl) => {
            walk_node(fn_decl.body, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::Block(block) => {
            for &stmt in &block.stmts {
                walk_node(stmt, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            }
            if let Some(tail) = block.tail {
                walk_node(tail, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            }
        }
        NodeKind::LetStmt(let_stmt) => {
            walk_node(let_stmt.init, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::ExprStmt(expr_stmt) => {
            walk_node(expr_stmt.expr, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::ReturnStmt(ret) => {
            if let Some(expr) = ret.expr {
                walk_node(expr, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            }
        }
        NodeKind::Ident(ident) => {
            let key = DefKey {
                module,
                name: ident.name,
                namespace: Namespace::Value,
            };

            if let Some(symbol_id) = index.get(key) {
                ref_index.record_ref(symbol_id, span.clone());
                node_to_symbol.bind_node(node_id, symbol_id);
            } else if let Some((import_arena, import_index)) = imports {
                if let Some((_, symbol_id)) = import_index.get_by_name_any_module(ident.name, Namespace::Value) {
                    let symbol = import_arena.get(symbol_id);
                    if !symbol.exported && symbol.module != module {
                        diagnostics.push(Diagnostic {
                            severity: Severity::Error,
                            code: "E_RESOLVE_VISIBILITY".to_string(),
                            message: format!("Symbol is not exported from module {:?}", symbol.module),
                            span: span.clone(),
                            related: Vec::new(),
                        });
                    }
                    ref_index.record_ref(symbol_id, span.clone());
                    node_to_symbol.bind_node(node_id, symbol_id);
                } else {
                    diagnostics.push(Diagnostic {
                        severity: Severity::Error,
                        code: "E_RESOLVE_UNDEFINED_IDENTIFIER".to_string(),
                        message: "Undefined identifier".to_string(),
                        span: span.clone(),
                        related: Vec::new(),
                    });
                }
            } else {
                diagnostics.push(Diagnostic {
                    severity: Severity::Error,
                    code: "E_RESOLVE_UNDEFINED_IDENTIFIER".to_string(),
                    message: "Undefined identifier".to_string(),
                    span: span.clone(),
                    related: Vec::new(),
                });
            }
        }
        NodeKind::Call(call) => {
            walk_node(call.callee, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            for &arg in &call.args {
                walk_node(arg, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            }
        }
        NodeKind::Unary(unary) => {
            walk_node(unary.expr, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::Binary(binary) => {
            walk_node(binary.left, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            walk_node(binary.right, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::Cast(cast) => {
            walk_node(cast.expr, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::IfExpr(if_expr) => {
            walk_node(if_expr.cond, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            walk_node(if_expr.then_block, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            if let Some(else_block) = if_expr.else_block {
                walk_node(else_block, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            }
        }
        NodeKind::WhenExpr(when_expr) => {
            for &arm in &when_expr.arms {
                walk_node(arm, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            }
        }
        NodeKind::WhenArm(when_arm) => {
            walk_node(when_arm.cond, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            walk_node(when_arm.body, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::Alloc(alloc) => {
            walk_node(alloc.ty, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::Mutate(mutate) => {
            walk_node(mutate.target, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            walk_node(mutate.body, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::Borrow(borrow) => {
            walk_node(borrow.target, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            walk_node(borrow.body, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::Peek(peek) => {
            walk_node(peek.target, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            walk_node(peek.body, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::MemberAccess(member) => {
            walk_node(member.object, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
        }
        NodeKind::TypeDecl(type_decl) => {
            if let Some(body) = type_decl.body {
                walk_node(body, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            }
        }
        NodeKind::TypeBody(body) => {
            for &method in &body.methods {
                walk_node(method, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            }
        }
        NodeKind::ServiceDecl(service) => {
            for &member in &service.members {
                walk_node(member, arena, module, index, imports, ref_index, node_to_symbol, diagnostics);
            }
        }
        // Nós literais ou que não contém outros nós para percorrer (neste nível)
        NodeKind::IntLit(_) | NodeKind::FloatLit(_) | NodeKind::BoundedLit(_) | NodeKind::StringLit(_) |
        NodeKind::TypeName(_) | NodeKind::TypeApp(_) | NodeKind::Import(_) | NodeKind::ImportSpec(_) |
        NodeKind::ServiceFnSig(_) | NodeKind::ConstructorDecl(_) | NodeKind::ConstantDecl(_) => {}
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::frontends::pbs::ast::{
        AstArena, BlockNodeArena, CallNodeArena, ExprStmtNodeArena, FileNodeArena, FnDeclNodeArena,
        IdentNodeArena, TypeDeclNodeArena,
    };
    use crate::common::spans::{Span, FileId};

    #[test]
    fn test_build_def_index_success() {
        let mut arena = AstArena::default();
        let mut interner = NameInterner::new();

        // Create a dummy body for the function to avoid panic/invalid access
        let body_id = arena.push(NodeKind::Block(BlockNodeArena { stmts: vec![], tail: None }), Span::new(FileId(0), 0, 0));

        let fn_name = interner.intern("my_func");
        let fn_id = arena.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: Some("pub".to_string()),
            name: fn_name,
            params: vec![],
            ret: None,
            else_fallback: None,
            body: body_id,
        }), Span::new(FileId(0), 10, 20));

        let type_name = interner.intern("MyStruct");
        let type_id = arena.push(NodeKind::TypeDecl(TypeDeclNodeArena {
            vis: None,
            type_kind: "struct".to_string(),
            name: type_name,
            is_host: false,
            params: vec![],
            constructors: vec![],
            constants: vec![],
            body: None,
        }), Span::new(FileId(0), 30, 40));

        let file_id = arena.push(NodeKind::File(FileNodeArena {
            imports: vec![],
            decls: vec![fn_id, type_id],
        }), Span::new(FileId(0), 0, 100));

        arena.roots.push(file_id);

        let (symbols, index, _ref_index, _node_to_symbol, diagnostics) = build_def_index(&arena, ModuleId(1), &interner, None);

        assert!(diagnostics.is_empty());
        assert_eq!(symbols.symbols.len(), 2);

        let fn_sym_id = index.get(DefKey { module: ModuleId(1), name: fn_name, namespace: Namespace::Value }).unwrap();
        let fn_sym = symbols.get(fn_sym_id);
        assert_eq!(fn_sym.name, fn_name);
        assert_eq!(fn_sym.kind, SymbolKind::Function);
        assert!(fn_sym.exported);

        let type_sym_id = index.get(DefKey { module: ModuleId(1), name: type_name, namespace: Namespace::Type }).unwrap();
        let type_sym = symbols.get(type_sym_id);
        assert_eq!(type_sym.name, type_name);
        assert_eq!(type_sym.kind, SymbolKind::Struct);
        assert!(!type_sym.exported);
    }

    #[test]
    fn test_build_def_index_duplicate() {
        let mut arena = AstArena::default();
        let mut interner = NameInterner::new();

        let name = interner.intern("conflict");
        
        let body_id = arena.push(NodeKind::Block(BlockNodeArena { stmts: vec![], tail: None }), Span::new(FileId(0), 0, 0));

        let fn1_id = arena.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: None,
            name,
            params: vec![],
            ret: None,
            else_fallback: None,
            body: body_id,
        }), Span::new(FileId(0), 10, 20));

        let fn2_id = arena.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: None,
            name,
            params: vec![],
            ret: None,
            else_fallback: None,
            body: body_id,
        }), Span::new(FileId(0), 30, 40));

        let file_id = arena.push(NodeKind::File(FileNodeArena {
            imports: vec![],
            decls: vec![fn1_id, fn2_id],
        }), Span::new(FileId(0), 0, 100));

        arena.roots.push(file_id);

        let (_symbols, _index, _ref_index, _node_to_symbol, diagnostics) = build_def_index(&arena, ModuleId(1), &interner, None);

        assert_eq!(diagnostics.len(), 1);
        assert_eq!(diagnostics[0].code, "E_RESOLVE_DUPLICATE_SYMBOL");
        assert_eq!(diagnostics[0].span, Span::new(FileId(0), 30, 40));
    }

    #[test]
    fn test_node_to_symbol_binding() {
        let mut arena = AstArena::default();
        let mut interner = NameInterner::new();

        let name = interner.intern("bound_func");
        let body_id = arena.push(NodeKind::Block(BlockNodeArena { stmts: vec![], tail: None }), Span::new(FileId(0), 0, 0));
        let decl_id = arena.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: None,
            name,
            params: vec![],
            ret: None,
            else_fallback: None,
            body: body_id,
        }), Span::new(FileId(0), 10, 20));

        let file_id = arena.push(NodeKind::File(FileNodeArena {
            imports: vec![],
            decls: vec![decl_id],
        }), Span::new(FileId(0), 0, 100));

        arena.roots.push(file_id);

        let (symbols, _index, _ref_index, node_to_symbol, _diagnostics) = build_def_index(&arena, ModuleId(1), &interner, None);

        let symbol_id = node_to_symbol.get(decl_id).expect("Node should be bound to a symbol");
        let symbol = symbols.get(symbol_id);
        assert_eq!(symbol.name, name);
    }

    #[test]
    fn test_resolve_undefined_identifier() {
        let mut arena = AstArena::default();
        let mut interner = NameInterner::new();

        let ident_name = interner.intern("unknown");
        let ident_id = arena.push(NodeKind::Ident(IdentNodeArena { name: ident_name }), Span::new(FileId(0), 50, 60));
        let expr_stmt = arena.push(NodeKind::ExprStmt(ExprStmtNodeArena { expr: ident_id }), Span::new(FileId(0), 50, 60));
        let body_id = arena.push(NodeKind::Block(BlockNodeArena { stmts: vec![expr_stmt], tail: None }), Span::new(FileId(0), 40, 70));

        let fn_id = arena.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: None,
            name: interner.intern("main"),
            params: vec![],
            ret: None,
            else_fallback: None,
            body: body_id,
        }), Span::new(FileId(0), 10, 80));

        let file_id = arena.push(NodeKind::File(FileNodeArena {
            imports: vec![],
            decls: vec![fn_id],
        }), Span::new(FileId(0), 0, 100));

        arena.roots.push(file_id);

        let (_symbols, _index, _ref_index, _node_to_symbol, diagnostics) = build_def_index(&arena, ModuleId(1), &interner, None);

        assert_eq!(diagnostics.len(), 1);
        assert_eq!(diagnostics[0].code, "E_RESOLVE_UNDEFINED_IDENTIFIER");
        assert_eq!(diagnostics[0].span, Span::new(FileId(0), 50, 60));
    }

    #[test]
    fn test_resolve_reference_success() {
        let mut arena = AstArena::default();
        let mut interner = NameInterner::new();

        // fn target() {}
        let target_name = interner.intern("target");
        let target_body = arena.push(NodeKind::Block(BlockNodeArena { stmts: vec![], tail: None }), Span::new(FileId(0), 5, 5));
        let target_id = arena.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: None,
            name: target_name,
            params: vec![],
            ret: None,
            else_fallback: None,
            body: target_body,
        }), Span::new(FileId(0), 0, 10));

        // fn caller() { target(); }
        let ident_id = arena.push(NodeKind::Ident(IdentNodeArena { name: target_name }), Span::new(FileId(0), 50, 56));
        let call_id = arena.push(NodeKind::Call(CallNodeArena { callee: ident_id, args: vec![] }), Span::new(FileId(0), 50, 58));
        let expr_stmt = arena.push(NodeKind::ExprStmt(ExprStmtNodeArena { expr: call_id }), Span::new(FileId(0), 50, 58));
        let body_id = arena.push(NodeKind::Block(BlockNodeArena { stmts: vec![expr_stmt], tail: None }), Span::new(FileId(0), 40, 70));

        let caller_id = arena.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: None,
            name: interner.intern("caller"),
            params: vec![],
            ret: None,
            else_fallback: None,
            body: body_id,
        }), Span::new(FileId(0), 30, 80));

        let file_id = arena.push(NodeKind::File(FileNodeArena {
            imports: vec![],
            decls: vec![target_id, caller_id],
        }), Span::new(FileId(0), 0, 100));

        arena.roots.push(file_id);

        let (_symbols, index, ref_index, node_to_symbol, diagnostics) = build_def_index(&arena, ModuleId(1), &interner, None);

        assert!(diagnostics.is_empty(), "Diagnostics should be empty: {:?}", diagnostics);
        
        let target_sym_id = index.get(DefKey { module: ModuleId(1), name: target_name, namespace: Namespace::Value }).expect("target should be in index");
        let refs = ref_index.refs_of(target_sym_id);
        assert_eq!(refs.len(), 1);
        assert_eq!(refs[0], Span::new(FileId(0), 50, 56));

        let bound_id = node_to_symbol.get(ident_id).expect("ident should be bound to symbol");
        assert_eq!(bound_id, target_sym_id);
    }

    #[test]
    fn test_resolve_visibility_violation() {
        let mut interner = NameInterner::new();

        // Módulo 1: define função privada
        let mut arena1 = AstArena::default();
        let target_name = interner.intern("private_func");
        let body1 = arena1.push(NodeKind::Block(BlockNodeArena { stmts: vec![], tail: None }), Span::new(FileId(0), 0, 0));
        let decl1 = arena1.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: None, // Privado
            name: target_name,
            params: vec![],
            ret: None,
            else_fallback: None,
            body: body1,
        }), Span::new(FileId(0), 0, 10));
        let file1 = arena1.push(NodeKind::File(FileNodeArena {
            imports: vec![],
            decls: vec![decl1],
        }), Span::new(FileId(0), 0, 100));
        arena1.roots.push(file1);

        let (symbols1, index1, _, _, _) = build_def_index(&arena1, ModuleId(1), &interner, None);

        // Módulo 2: tenta usar função privada do Módulo 1
        let mut arena2 = AstArena::default();
        let ident_id = arena2.push(NodeKind::Ident(IdentNodeArena { name: target_name }), Span::new(FileId(0), 50, 62));
        let expr_stmt = arena2.push(NodeKind::ExprStmt(ExprStmtNodeArena { expr: ident_id }), Span::new(FileId(0), 50, 62));
        let body2 = arena2.push(NodeKind::Block(BlockNodeArena { stmts: vec![expr_stmt], tail: None }), Span::new(FileId(0), 40, 70));
        let caller = arena2.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: Some("pub".to_string()),
            name: interner.intern("caller"),
            params: vec![],
            ret: None,
            else_fallback: None,
            body: body2,
        }), Span::new(FileId(0), 30, 80));
        let file2 = arena2.push(NodeKind::File(FileNodeArena {
            imports: vec![],
            decls: vec![caller],
        }), Span::new(FileId(0), 0, 100));
        arena2.roots.push(file2);

        let (_symbols2, _index2, _ref_index2, _node_to_symbol2, diagnostics) = 
            build_def_index(&arena2, ModuleId(2), &interner, Some((&symbols1, &index1)));

        assert_eq!(diagnostics.len(), 1);
        assert_eq!(diagnostics[0].code, "E_RESOLVE_VISIBILITY");
        assert_eq!(diagnostics[0].span, Span::new(FileId(0), 50, 62));
    }

    #[test]
    fn test_determinism() {
        let mut arena = AstArena::default();
        let mut interner = NameInterner::new();

        let target_name = interner.intern("target");
        let target_body = arena.push(NodeKind::Block(BlockNodeArena { stmts: vec![], tail: None }), Span::new(FileId(0), 5, 5));
        let target_id = arena.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: Some("pub".to_string()),
            name: target_name,
            params: vec![],
            ret: None,
            else_fallback: None,
            body: target_body,
        }), Span::new(FileId(0), 0, 10));

        let caller_name = interner.intern("caller");
        let ident_id = arena.push(NodeKind::Ident(IdentNodeArena { name: target_name }), Span::new(FileId(0), 50, 56));
        let call_id = arena.push(NodeKind::Call(CallNodeArena { callee: ident_id, args: vec![] }), Span::new(FileId(0), 50, 58));
        let expr_stmt = arena.push(NodeKind::ExprStmt(ExprStmtNodeArena { expr: call_id }), Span::new(FileId(0), 50, 58));
        let body_id = arena.push(NodeKind::Block(BlockNodeArena { stmts: vec![expr_stmt], tail: None }), Span::new(FileId(0), 40, 70));

        let caller_id = arena.push(NodeKind::FnDecl(FnDeclNodeArena {
            vis: None,
            name: caller_name,
            params: vec![],
            ret: None,
            else_fallback: None,
            body: body_id,
        }), Span::new(FileId(0), 30, 80));

        let file_id = arena.push(NodeKind::File(FileNodeArena {
            imports: vec![],
            decls: vec![target_id, caller_id],
        }), Span::new(FileId(0), 0, 100));

        arena.roots.push(file_id);

        let run1 = build_def_index(&arena, ModuleId(1), &interner, None);
        let run2 = build_def_index(&arena, ModuleId(1), &interner, None);

        // runX is (SymbolArena, DefIndex, RefIndex, NodeToSymbol, Vec<Diagnostic>)
        
        let json1_symbols = serde_json::to_string(&run1.0).unwrap();
        let json2_symbols = serde_json::to_string(&run2.0).unwrap();
        assert_eq!(json1_symbols, json2_symbols, "SymbolArena should be deterministic");

        let json1_refs = serde_json::to_string(&run1.2).unwrap();
        let json2_refs = serde_json::to_string(&run2.2).unwrap();
        assert_eq!(json1_refs, json2_refs, "RefIndex should be deterministic");

        let json1_node_to_symbol = serde_json::to_string(&run1.3).unwrap();
        let json2_node_to_symbol = serde_json::to_string(&run2.3).unwrap();
        assert_eq!(json1_node_to_symbol, json2_node_to_symbol, "NodeToSymbol should be deterministic");

        let json1_diags = serde_json::to_string(&run1.4).unwrap();
        let json2_diags = serde_json::to_string(&run2.4).unwrap();
        assert_eq!(json1_diags, json2_diags, "Diagnostics should be deterministic");
    }
}