use crate::common::diagnostics::{Diagnostic, DiagnosticBundle, DiagnosticLevel};
use crate::common::spans::Span;
use crate::frontends::pbs::ast::*;
use crate::frontends::pbs::symbols::*;
use std::collections::HashMap;

pub trait ModuleProvider {
    fn get_module_symbols(&self, from_path: &str) -> Option<&ModuleSymbols>;
}

pub struct Resolver<'a> {
    module_provider: &'a dyn ModuleProvider,
    current_module: &'a ModuleSymbols,
    scopes: Vec<HashMap<String, Symbol>>,
    imported_symbols: ModuleSymbols,
    diagnostics: Vec<Diagnostic>,
}

impl<'a> Resolver<'a> {
    pub fn new(
        current_module: &'a ModuleSymbols,
        module_provider: &'a dyn ModuleProvider,
    ) -> Self {
        Self {
            module_provider,
            current_module,
            scopes: Vec::new(),
            imported_symbols: ModuleSymbols::new(),
            diagnostics: Vec::new(),
        }
    }

    pub fn resolve(&mut self, file: &FileNode) -> Result<(), DiagnosticBundle> {
        // Step 1: Process imports to populate imported_symbols
        for imp in &file.imports {
            if let Node::Import(imp_node) = imp {
                self.resolve_import(imp_node);
            }
        }

        // Step 2: Resolve all top-level declarations
        for decl in &file.decls {
            self.resolve_node(decl);
        }

        if !self.diagnostics.is_empty() {
            return Err(DiagnosticBundle {
                diagnostics: self.diagnostics.clone(),
            });
        }

        Ok(())
    }

    fn resolve_import(&mut self, imp: &ImportNode) {
        let provider = self.module_provider;
        if let Some(target_symbols) = provider.get_module_symbols(&imp.from) {
            if let Node::ImportSpec(spec) = &*imp.spec {
                for name in &spec.path {
                    // Try to find in Type namespace
                    if let Some(sym) = target_symbols.type_symbols.get(name) {
                        if sym.visibility == Visibility::Pub {
                            if let Err(_) = self.imported_symbols.type_symbols.insert(sym.clone()) {
                                self.error_duplicate_import(name, imp.span);
                            }
                        } else {
                            self.error_visibility(sym, imp.span);
                        }
                    } 
                    // Try to find in Value namespace
                    else if let Some(sym) = target_symbols.value_symbols.get(name) {
                        if sym.visibility == Visibility::Pub {
                            if let Err(_) = self.imported_symbols.value_symbols.insert(sym.clone()) {
                                self.error_duplicate_import(name, imp.span);
                            }
                        } else {
                            self.error_visibility(sym, imp.span);
                        }
                    } else {
                        self.error_undefined(name, imp.span);
                    }
                }
            }
        } else {
            self.diagnostics.push(Diagnostic {
                level: DiagnosticLevel::Error,
                code: Some("E_RESOLVE_INVALID_IMPORT".to_string()),
                message: format!("Module not found: {}", imp.from),
                span: Some(imp.span),
            });
        }
    }

    fn resolve_node(&mut self, node: &Node) {
        match node {
            Node::FnDecl(n) => self.resolve_fn_decl(n),
            Node::ServiceDecl(n) => self.resolve_service_decl(n),
            Node::TypeDecl(n) => self.resolve_type_decl(n),
            Node::Block(n) => self.resolve_block(n),
            Node::LetStmt(n) => self.resolve_let_stmt(n),
            Node::ExprStmt(n) => self.resolve_node(&n.expr),
            Node::ReturnStmt(n) => {
                if let Some(expr) = &n.expr {
                    self.resolve_node(expr);
                }
            }
            Node::Call(n) => {
                self.resolve_node(&n.callee);
                for arg in &n.args {
                    self.resolve_node(arg);
                }
            }
            Node::Unary(n) => self.resolve_node(&n.expr),
            Node::Binary(n) => {
                self.resolve_node(&n.left);
                self.resolve_node(&n.right);
            }
            Node::Cast(n) => {
                self.resolve_node(&n.expr);
                self.resolve_type_ref(&n.ty);
            }
            Node::IfExpr(n) => {
                self.resolve_node(&n.cond);
                self.resolve_node(&n.then_block);
                if let Some(else_block) = &n.else_block {
                    self.resolve_node(else_block);
                }
            }
            Node::WhenExpr(n) => {
                for arm in &n.arms {
                    if let Node::WhenArm(arm_node) = arm {
                        self.resolve_node(&arm_node.cond);
                        self.resolve_node(&arm_node.body);
                    }
                }
            }
            Node::Ident(n) => {
                self.resolve_identifier(&n.name, n.span, Namespace::Value);
            }
            Node::TypeName(n) => {
                self.resolve_identifier(&n.name, n.span, Namespace::Type);
            }
            Node::TypeApp(n) => {
                self.resolve_identifier(&n.base, n.span, Namespace::Type);
                for arg in &n.args {
                    self.resolve_type_ref(arg);
                }
            }
            Node::Alloc(n) => {
                self.resolve_type_ref(&n.ty);
            }
            Node::Mutate(n) => {
                self.resolve_node(&n.target);
                self.enter_scope();
                self.define_local(&n.binding, n.span, SymbolKind::Local);
                self.resolve_node(&n.body);
                self.exit_scope();
            }
            Node::Borrow(n) => {
                self.resolve_node(&n.target);
                self.enter_scope();
                self.define_local(&n.binding, n.span, SymbolKind::Local);
                self.resolve_node(&n.body);
                self.exit_scope();
            }
            Node::Peek(n) => {
                self.resolve_node(&n.target);
                self.enter_scope();
                self.define_local(&n.binding, n.span, SymbolKind::Local);
                self.resolve_node(&n.body);
                self.exit_scope();
            }
            Node::MemberAccess(n) => {
                self.resolve_node(&n.object);
                // For member access, the member name itself isn't resolved in the value namespace
                // unless it's a property. In v0, we mostly care about host calls.
            }
            _ => {}
        }
    }

    fn resolve_fn_decl(&mut self, n: &FnDeclNode) {
        self.enter_scope();
        for param in &n.params {
            self.resolve_type_ref(&param.ty);
            self.define_local(&param.name, param.span, SymbolKind::Local);
        }
        if let Some(ret) = &n.ret {
            self.resolve_type_ref(ret);
        }
        self.resolve_node(&n.body);
        self.exit_scope();
    }

    fn resolve_service_decl(&mut self, n: &ServiceDeclNode) {
        if let Some(ext) = &n.extends {
            self.resolve_identifier(ext, n.span, Namespace::Type);
        }
        for member in &n.members {
            if let Node::ServiceFnSig(sig) = member {
                for param in &sig.params {
                    self.resolve_type_ref(&param.ty);
                }
                self.resolve_type_ref(&sig.ret);
            }
        }
    }

    fn resolve_type_decl(&mut self, n: &TypeDeclNode) {
        if let Node::TypeBody(body) = &*n.body {
            for member in &body.members {
                self.resolve_type_ref(&member.ty);
            }
        }
    }

    fn resolve_block(&mut self, n: &BlockNode) {
        self.enter_scope();
        for stmt in &n.stmts {
            self.resolve_node(stmt);
        }
        self.exit_scope();
    }

    fn resolve_let_stmt(&mut self, n: &LetStmtNode) {
        if let Some(ty) = &n.ty {
            self.resolve_type_ref(ty);
        }
        self.resolve_node(&n.init);
        self.define_local(&n.name, n.span, SymbolKind::Local);
    }

    fn resolve_type_ref(&mut self, node: &Node) {
        self.resolve_node(node);
    }

    fn resolve_identifier(&mut self, name: &str, span: Span, namespace: Namespace) -> Option<Symbol> {
        // Built-ins (minimal for v0)
        if namespace == Namespace::Type {
            match name {
                "int" | "float" | "string" | "bool" | "void" | "optional" | "result" => return None,
                _ => {}
            }
        }

        if namespace == Namespace::Value {
            match name {
                "none" | "some" | "ok" | "err" | "true" | "false" => return None,
                _ => {}
            }
        }

        // 1. local bindings
        if namespace == Namespace::Value {
            for scope in self.scopes.iter().rev() {
                if let Some(sym) = scope.get(name) {
                    return Some(sym.clone());
                }
            }
        }

        let table = if namespace == Namespace::Type {
            &self.current_module.type_symbols
        } else {
            &self.current_module.value_symbols
        };

        // 2 & 3. file-private and module symbols
        if let Some(sym) = table.get(name) {
            return Some(sym.clone());
        }

        // 4. imported symbols
        let imp_table = if namespace == Namespace::Type {
            &self.imported_symbols.type_symbols
        } else {
            &self.imported_symbols.value_symbols
        };
        if let Some(sym) = imp_table.get(name) {
            return Some(sym.clone());
        }

        if namespace == Namespace::Type {
            self.diagnostics.push(Diagnostic {
                level: DiagnosticLevel::Error,
                code: Some("E_TYPE_UNKNOWN_TYPE".to_string()),
                message: format!("Unknown type: {}", name),
                span: Some(span),
            });
        } else {
            self.error_undefined(name, span);
        }
        None
    }

    fn define_local(&mut self, name: &str, span: Span, kind: SymbolKind) {
        let scope = self.scopes.last_mut().expect("No scope to define local");
        
        // Check for collision in Type namespace at top-level?
        // Actually, the spec says "A name may not exist in both namespaces".
        // If we want to be strict, we check current module's type symbols too.
        if self.current_module.type_symbols.get(name).is_some() {
             self.diagnostics.push(Diagnostic {
                level: DiagnosticLevel::Error,
                code: Some("E_RESOLVE_NAMESPACE_COLLISION".to_string()),
                message: format!("Local variable '{}' collides with a type name", name),
                span: Some(span),
            });
            return;
        }

        if scope.contains_key(name) {
            self.diagnostics.push(Diagnostic {
                level: DiagnosticLevel::Error,
                code: Some("E_RESOLVE_DUPLICATE_SYMBOL".to_string()),
                message: format!("Duplicate local variable '{}'", name),
                span: Some(span),
            });
        } else {
            scope.insert(name.to_string(), Symbol {
                name: name.to_string(),
                kind,
                namespace: Namespace::Value,
                visibility: Visibility::FilePrivate,
                ty: None, // Will be set by TypeChecker
                span,
            });
        }
    }

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

    fn exit_scope(&mut self) {
        self.scopes.pop();
    }

    fn error_undefined(&mut self, name: &str, span: Span) {
        self.diagnostics.push(Diagnostic {
            level: DiagnosticLevel::Error,
            code: Some("E_RESOLVE_UNDEFINED".to_string()),
            message: format!("Undefined identifier: {}", name),
            span: Some(span),
        });
    }

    fn error_duplicate_import(&mut self, name: &str, span: Span) {
        self.diagnostics.push(Diagnostic {
            level: DiagnosticLevel::Error,
            code: Some("E_RESOLVE_DUPLICATE_SYMBOL".to_string()),
            message: format!("Duplicate import: {}", name),
            span: Some(span),
        });
    }

    fn error_visibility(&mut self, sym: &Symbol, span: Span) {
        self.diagnostics.push(Diagnostic {
            level: DiagnosticLevel::Error,
            code: Some("E_RESOLVE_VISIBILITY".to_string()),
            message: format!("Symbol '{}' is not visible here", sym.name),
            span: Some(span),
        });
    }
}