use std::collections::HashMap;

use crate::ast::*;
use crate::diagnostic::Diagnostic;
use crate::identity::{BindingId, BindingKind, ExprId, ModuleId};
use crate::package_signature::{
    PackageModuleSignatureEnvironment, PackageSignatureEnvironment, PackageSignatureSource,
};
use crate::prelude::{self, BuiltinKind};
use crate::span::Span;
use crate::symbol::{Symbol, SymbolTable};

#[derive(Clone, Debug)]
pub struct ResolveOutput {
    pub diagnostics: Vec<Diagnostic>,
    pub bindings: Vec<BindingInfo>,
    pub identifier_refs: Vec<ResolvedIdentifier>,
    pub symbols: SymbolTable,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BindingInfo {
    pub id: BindingId,
    pub symbol: Symbol,
    pub kind: BindingKind,
    pub span: Span,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ResolvedIdentifier {
    pub expr_id: ExprId,
    pub name: Symbol,
    pub span: Span,
    pub binding: BindingId,
}

pub fn resolve(program: &Program) -> Vec<Diagnostic> {
    resolve_program(program).diagnostics
}

pub fn resolve_program(program: &Program) -> ResolveOutput {
    let mut resolver = Resolver::new();
    resolver.install_prelude();
    resolver.predeclare_records(&program.statements);
    resolver.predeclare_enums(&program.statements);
    resolver.resolve_scope_statements(&program.statements);
    resolver.into_output()
}

pub fn resolve_package_module(
    program: &Program,
    signatures: &PackageSignatureEnvironment,
    module: ModuleId,
) -> ResolveOutput {
    let mut resolver = Resolver::new();
    resolver.install_prelude();
    if program
        .package
        .as_ref()
        .is_some_and(|package| crate::std_package::allows_internal_builtins(&package.path))
    {
        resolver.install_internal_builtins();
    }
    if let Some(environment) = signatures.module(module) {
        resolver.install_package_module_signatures(signatures, environment);
    }
    resolver.predeclare_records(&program.statements);
    resolver.predeclare_enums(&program.statements);
    resolver.resolve_scope_statements(&program.statements);
    resolver.into_output()
}

#[derive(Clone, Copy, Debug)]
struct Binding {
    id: BindingId,
    symbol: Symbol,
    kind: BindingKind,
    span: Span,
}

#[derive(Clone, Debug)]
struct EnumResolveInfo {
    span: Span,
    variants: HashMap<Symbol, Span>,
}

struct ScopeFrame {
    bindings: HashMap<Symbol, BindingId>,
    function_boundary: bool,
}

impl ScopeFrame {
    fn new(function_boundary: bool) -> Self {
        Self {
            bindings: HashMap::new(),
            function_boundary,
        }
    }
}

struct Resolver {
    scopes: Vec<ScopeFrame>,
    records: HashMap<Symbol, Span>,
    enums: HashMap<Symbol, EnumResolveInfo>,
    bindings: Vec<Binding>,
    identifier_refs: Vec<ResolvedIdentifier>,
    symbols: SymbolTable,
    diagnostics: Vec<Diagnostic>,
}

impl Resolver {
    fn new() -> Self {
        Self {
            scopes: vec![ScopeFrame::new(true)],
            records: HashMap::new(),
            enums: HashMap::new(),
            bindings: Vec::new(),
            identifier_refs: Vec::new(),
            symbols: SymbolTable::default(),
            diagnostics: Vec::new(),
        }
    }

    fn into_output(self) -> ResolveOutput {
        let bindings = self
            .bindings
            .iter()
            .map(|binding| BindingInfo {
                id: binding.id,
                symbol: binding.symbol,
                kind: binding.kind,
                span: binding.span,
            })
            .collect();
        ResolveOutput {
            diagnostics: self.diagnostics,
            bindings,
            identifier_refs: self.identifier_refs,
            symbols: self.symbols,
        }
    }

    fn install_prelude(&mut self) {
        for builtin in prelude::builtins() {
            self.install_builtin(*builtin);
        }
    }

    fn install_internal_builtins(&mut self) {
        for builtin in prelude::internal_builtins() {
            self.install_builtin(*builtin);
        }
    }

    fn install_builtin(&mut self, builtin: prelude::Builtin) {
        let kind = match builtin.kind {
            BuiltinKind::Function => BindingKind::Function,
            BuiltinKind::Value => BindingKind::Immutable,
        };
        let symbol = self.symbol(builtin.name);
        self.insert_current(symbol, kind, Span::default());
    }

    fn install_package_module_signatures(
        &mut self,
        signatures: &PackageSignatureEnvironment,
        environment: &PackageModuleSignatureEnvironment,
    ) {
        for visible in &environment.records {
            if visible.source == PackageSignatureSource::ModuleLocal {
                continue;
            }
            let Some(record) = signatures.record(visible.item) else {
                continue;
            };
            let symbol = self.symbol(&visible.name);
            self.records.insert(symbol, record.span);
        }

        for visible in &environment.enums {
            if visible.source == PackageSignatureSource::ModuleLocal {
                continue;
            }
            let Some(enumeration) = signatures.enumeration(visible.item) else {
                continue;
            };
            let enum_symbol = self.symbol(&visible.name);
            let mut variants = HashMap::new();
            for variant in &enumeration.variants {
                let variant_symbol = self.symbol(&variant.name);
                variants.insert(variant_symbol, variant.span);
                let qualified = self.symbol(&format!("{}::{}", visible.name, variant.name));
                let kind = if variant.payload.is_some() {
                    BindingKind::Function
                } else {
                    BindingKind::Immutable
                };
                self.insert_current(qualified, kind, variant.span);
            }
            self.enums.insert(
                enum_symbol,
                EnumResolveInfo {
                    span: enumeration.span,
                    variants,
                },
            );
        }

        for visible in &environment.functions {
            if visible.source == PackageSignatureSource::ModuleLocal {
                continue;
            }
            let Some(function) = signatures.function(visible.item) else {
                continue;
            };
            let symbol = self.symbol(&visible.name);
            self.insert_current(symbol, BindingKind::Function, function.span);
        }
    }

    fn resolve_scope_statements(&mut self, statements: &[Stmt]) {
        self.predeclare_functions(statements);
        for statement in statements {
            self.resolve_stmt(statement);
        }
    }

    fn resolve_block(&mut self, block: &Block) {
        self.push_scope(false);
        self.resolve_scope_statements(&block.statements);
        self.pop_scope();
    }

    fn resolve_value_block(&mut self, block: &ValueBlock) {
        self.push_scope(false);
        self.predeclare_functions(&block.statements);
        for statement in &block.statements {
            self.resolve_stmt(statement);
        }
        self.resolve_expr(&block.expr);
        self.pop_scope();
    }

    fn resolve_stmt(&mut self, statement: &Stmt) {
        match statement {
            Stmt::Assign(stmt) => self.resolve_assign(stmt),
            Stmt::RecordDecl(_) => {}
            Stmt::EnumDecl(_) => {}
            Stmt::FuncDecl(stmt) => self.resolve_func_decl(stmt),
            Stmt::If(stmt) => {
                self.resolve_expr(&stmt.condition);
                self.resolve_block(&stmt.then_branch);
                if let Some(else_branch) = &stmt.else_branch {
                    self.resolve_block(else_branch);
                }
            }
            Stmt::While(stmt) => {
                self.resolve_expr(&stmt.condition);
                self.resolve_block(&stmt.body);
            }
            Stmt::Expr(stmt) => self.resolve_expr(&stmt.expr),
        }
    }

    fn resolve_assign(&mut self, stmt: &AssignStmt) {
        self.resolve_expr(&stmt.value);
        let name = self.symbol(&stmt.name);
        if stmt.mutable {
            if let Some(binding) = self.current_scope_binding(name) {
                self.diagnostics.push(
                    Diagnostic::new(
                        "E002",
                        format!("duplicate binding `{}` in the current scope", stmt.name),
                        stmt.span,
                    )
                    .with_related("previous binding is here", binding.span),
                );
            } else if let Some(binding) = self.any_enclosing_scope_lookup(name).copied() {
                self.diagnostics.push(
                    Diagnostic::new(
                        "E003",
                        format!("shadowing is prohibited for `{}`", stmt.name),
                        stmt.span,
                    )
                    .with_related("shadowed binding is here", binding.span),
                );
            } else {
                self.insert_current(name, BindingKind::Mutable, stmt.span);
            }
            return;
        }

        if let Some(binding) = self.lookup_in_current_function(name) {
            match binding.kind {
                BindingKind::Mutable => {}
                BindingKind::Immutable | BindingKind::Function | BindingKind::Parameter => {
                    self.diagnostics.push(
                        Diagnostic::new(
                            "E001",
                            format!("cannot update immutable binding `{}`", stmt.name),
                            stmt.span,
                        )
                        .with_related("binding is declared here", binding.span),
                    );
                }
            }
            return;
        }

        if let Some(binding) = self.lookup_beyond_current_function(name) {
            match binding.kind {
                BindingKind::Mutable => {
                    self.diagnostics.push(
                        Diagnostic::new(
                            "E004",
                            format!(
                                "cannot update outer-scope mutable binding `{}` in v1",
                                stmt.name
                            ),
                            stmt.span,
                        )
                        .with_related("outer mutable binding is declared here", binding.span)
                        .with_suggestion(
                            "pass the value into this function and return the updated value",
                        ),
                    );
                }
                BindingKind::Immutable | BindingKind::Function | BindingKind::Parameter => {
                    self.diagnostics.push(
                        Diagnostic::new(
                            "E003",
                            format!("shadowing is prohibited for `{}`", stmt.name),
                            stmt.span,
                        )
                        .with_related("shadowed binding is here", binding.span),
                    );
                }
            }
            return;
        }

        self.insert_current(name, BindingKind::Immutable, stmt.span);
    }

    fn resolve_func_decl(&mut self, stmt: &FuncDecl) {
        self.push_scope(true);
        for param in &stmt.params {
            let name = self.symbol(&param.name);
            if let Some(binding) = self.current_scope_binding(name) {
                self.diagnostics.push(
                    Diagnostic::new(
                        "E002",
                        format!("duplicate binding `{}` in the current scope", param.name),
                        param.span,
                    )
                    .with_related("previous binding is here", binding.span),
                );
                continue;
            }
            if let Some(binding) = self.any_enclosing_scope_lookup(name).copied() {
                self.diagnostics.push(
                    Diagnostic::new(
                        "E003",
                        format!("shadowing is prohibited for `{}`", param.name),
                        param.span,
                    )
                    .with_related("shadowed binding is here", binding.span),
                );
                continue;
            }
            self.insert_current(name, BindingKind::Parameter, param.span);
        }
        self.resolve_scope_statements(&stmt.body.statements);
        self.resolve_expr(&stmt.body.expr);
        self.pop_scope();
    }

    fn resolve_expr(&mut self, expr: &Expr) {
        match expr {
            Expr::Int(_) | Expr::Bool(_) | Expr::String(_) | Expr::Unit(_) => {}
            Expr::ListLit(expr) => {
                for item in &expr.items {
                    self.resolve_expr(item);
                }
            }
            Expr::Index(expr) => {
                self.resolve_expr(&expr.base);
                self.resolve_expr(&expr.index);
            }
            Expr::Ident(expr) => {
                let name = self.symbol(&expr.name);
                if let Some(binding) = self.any_scope_lookup(name).copied() {
                    self.identifier_refs.push(ResolvedIdentifier {
                        expr_id: expr.id,
                        name,
                        span: expr.span,
                        binding: binding.id,
                    });
                } else if let Some((enum_name, variant_name)) = split_variant_name(&expr.name) {
                    self.push_unresolved_variant_diagnostic(enum_name, variant_name, expr.span);
                } else {
                    self.diagnostics.push(Diagnostic::new(
                        "N001",
                        format!("unresolved name `{}`", expr.name),
                        expr.span,
                    ));
                }
            }
            Expr::RecordLit(expr) => {
                for field in &expr.fields {
                    self.resolve_expr(&field.value);
                }
            }
            Expr::Field(expr) => self.resolve_expr(&expr.base),
            Expr::RecordUpdate(expr) => {
                self.resolve_expr(&expr.base);
                for field in &expr.fields {
                    self.resolve_expr(&field.value);
                }
            }
            Expr::Unary(expr) => self.resolve_expr(&expr.expr),
            Expr::Binary(expr) => {
                self.resolve_expr(&expr.left);
                self.resolve_expr(&expr.right);
            }
            Expr::Call(expr) => {
                self.resolve_expr(&expr.callee);
                for arg in &expr.args {
                    self.resolve_expr(arg);
                }
            }
            Expr::Try(expr) => self.resolve_expr(&expr.expr),
            Expr::If(expr) => {
                self.resolve_expr(&expr.condition);
                self.resolve_value_block(&expr.then_branch);
                self.resolve_value_block(&expr.else_branch);
            }
            Expr::Match(expr) => {
                self.resolve_expr(&expr.value);
                for arm in &expr.arms {
                    self.push_scope(false);
                    let MatchPattern::Variant(pattern) = &arm.pattern;
                    if let Some(binding) = &pattern.binding {
                        self.resolve_pattern_binding(binding, pattern.span);
                    }
                    self.resolve_expr(&arm.value);
                    self.pop_scope();
                }
            }
            Expr::Fn(expr) => {
                self.push_scope(true);
                for param in &expr.params {
                    let name = self.symbol(&param.name);
                    if let Some(binding) = self.current_scope_binding(name) {
                        self.diagnostics.push(
                            Diagnostic::new(
                                "E002",
                                format!("duplicate binding `{}` in the current scope", param.name),
                                param.span,
                            )
                            .with_related("previous binding is here", binding.span),
                        );
                        continue;
                    }
                    if let Some(binding) = self.any_enclosing_scope_lookup(name).copied() {
                        self.diagnostics.push(
                            Diagnostic::new(
                                "E003",
                                format!("shadowing is prohibited for `{}`", param.name),
                                param.span,
                            )
                            .with_related("shadowed binding is here", binding.span),
                        );
                        continue;
                    }
                    self.insert_current(name, BindingKind::Parameter, param.span);
                }
                self.resolve_scope_statements(&expr.body.statements);
                self.resolve_expr(&expr.body.expr);
                self.pop_scope();
            }
        }
    }

    fn resolve_pattern_binding(&mut self, name: &str, span: Span) {
        let symbol = self.symbol(name);
        if let Some(binding) = self.current_scope_binding(symbol) {
            self.diagnostics.push(
                Diagnostic::new(
                    "E002",
                    format!("duplicate binding `{name}` in the current scope"),
                    span,
                )
                .with_related("previous binding is here", binding.span),
            );
            return;
        }
        if let Some(binding) = self.any_enclosing_scope_lookup(symbol).copied() {
            self.diagnostics.push(
                Diagnostic::new(
                    "E003",
                    format!("shadowing is prohibited for `{name}`"),
                    span,
                )
                .with_related("shadowed binding is here", binding.span),
            );
            return;
        }
        self.insert_current(symbol, BindingKind::Immutable, span);
    }

    fn push_unresolved_variant_diagnostic(
        &mut self,
        enum_name: &str,
        variant_name: &str,
        span: Span,
    ) {
        let enum_symbol = self.symbol(enum_name);
        let variant_symbol = self.symbol(variant_name);
        if let Some(enumeration) = self.enums.get(&enum_symbol) {
            if enumeration.variants.contains_key(&variant_symbol) {
                self.diagnostics.push(Diagnostic::new(
                    "N001",
                    format!("unresolved name `{enum_name}::{variant_name}`"),
                    span,
                ));
                return;
            }
            self.diagnostics.push(
                Diagnostic::new(
                    "N001",
                    format!("unknown variant `{variant_name}` for enum `{enum_name}`"),
                    span,
                )
                .with_related("enum is declared here", enumeration.span),
            );
            return;
        }

        self.diagnostics.push(Diagnostic::new(
            "N001",
            format!("unknown enum `{enum_name}` in variant reference"),
            span,
        ));
    }

    fn predeclare_functions(&mut self, statements: &[Stmt]) {
        for statement in statements {
            if let Stmt::FuncDecl(func) = statement {
                let name = self.symbol(&func.name);
                if let Some(binding) = self.current_scope_binding(name) {
                    self.diagnostics.push(
                        Diagnostic::new(
                            "E002",
                            format!("duplicate binding `{}` in the current scope", func.name),
                            func.span,
                        )
                        .with_related("previous binding is here", binding.span),
                    );
                } else if let Some(binding) = self.any_enclosing_scope_lookup(name).copied() {
                    self.diagnostics.push(
                        Diagnostic::new(
                            "E003",
                            format!("shadowing is prohibited for `{}`", func.name),
                            func.span,
                        )
                        .with_related("shadowed binding is here", binding.span),
                    );
                } else {
                    self.insert_current(name, BindingKind::Function, func.span);
                }
            }
        }
    }

    fn predeclare_records(&mut self, statements: &[Stmt]) {
        for statement in statements {
            if let Stmt::RecordDecl(record) = statement {
                let name = self.symbol(&record.name);
                if let Some(span) = self.records.get(&name).copied() {
                    self.diagnostics.push(
                        Diagnostic::new(
                            "E002",
                            format!("duplicate record `{}` in the current scope", record.name),
                            record.span,
                        )
                        .with_related("previous record declaration is here", span),
                    );
                } else {
                    self.records.insert(name, record.span);
                }
            }
        }
    }

    fn predeclare_enums(&mut self, statements: &[Stmt]) {
        for statement in statements {
            if let Stmt::EnumDecl(enumeration) = statement {
                let name = self.symbol(&enumeration.name);
                if let Some(existing) = self.enums.get(&name) {
                    self.diagnostics.push(
                        Diagnostic::new(
                            "E002",
                            format!("duplicate enum `{}` in the current scope", enumeration.name),
                            enumeration.span,
                        )
                        .with_related("previous enum declaration is here", existing.span),
                    );
                } else if let Some(span) = self.records.get(&name).copied() {
                    self.diagnostics.push(
                        Diagnostic::new(
                            "E002",
                            format!("duplicate type `{}` in the current scope", enumeration.name),
                            enumeration.span,
                        )
                        .with_related("previous type declaration is here", span),
                    );
                } else {
                    let mut variants = HashMap::new();
                    for variant in &enumeration.variants {
                        let variant_name = self.symbol(&variant.name);
                        variants.insert(variant_name, variant.span);
                        let qualified =
                            self.symbol(&format!("{}::{}", enumeration.name, variant.name));
                        let kind = if variant.payload.is_some() {
                            BindingKind::Function
                        } else {
                            BindingKind::Immutable
                        };
                        self.insert_current(qualified, kind, variant.span);
                    }
                    self.enums.insert(
                        name,
                        EnumResolveInfo {
                            span: enumeration.span,
                            variants,
                        },
                    );
                }
            }
        }
    }

    fn push_scope(&mut self, function_boundary: bool) {
        self.scopes.push(ScopeFrame::new(function_boundary));
    }

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

    fn current_scope_binding(&self, name: Symbol) -> Option<Binding> {
        self.scopes
            .last()
            .and_then(|scope| scope.bindings.get(&name).map(|id| *self.binding(*id)))
    }

    fn lookup_in_current_function(&self, name: Symbol) -> Option<&Binding> {
        for scope in self.scopes.iter().rev() {
            if let Some(id) = scope.bindings.get(&name) {
                return Some(self.binding(*id));
            }
            if scope.function_boundary {
                break;
            }
        }
        None
    }

    fn lookup_beyond_current_function(&self, name: Symbol) -> Option<&Binding> {
        let boundary_index = self
            .scopes
            .iter()
            .rposition(|scope| scope.function_boundary)
            .unwrap_or(0);
        self.scopes[..boundary_index]
            .iter()
            .rev()
            .find_map(|scope| scope.bindings.get(&name).map(|id| self.binding(*id)))
    }

    fn any_enclosing_scope_lookup(&self, name: Symbol) -> Option<&Binding> {
        self.scopes
            .iter()
            .rev()
            .skip(1)
            .find_map(|scope| scope.bindings.get(&name).map(|id| self.binding(*id)))
    }

    fn any_scope_lookup(&self, name: Symbol) -> Option<&Binding> {
        self.scopes
            .iter()
            .rev()
            .find_map(|scope| scope.bindings.get(&name).map(|id| self.binding(*id)))
    }

    fn insert_current(&mut self, name: Symbol, kind: BindingKind, span: Span) -> BindingId {
        let id = BindingId::new(self.bindings.len() as u32);
        self.bindings.push(Binding {
            id,
            symbol: name,
            kind,
            span,
        });
        if let Some(scope) = self.scopes.last_mut() {
            scope.bindings.insert(name, id);
        }
        id
    }

    fn binding(&self, id: BindingId) -> &Binding {
        &self.bindings[id.as_u32() as usize]
    }

    fn symbol(&mut self, name: &str) -> Symbol {
        self.symbols.intern(name)
    }
}

fn split_variant_name(name: &str) -> Option<(&str, &str)> {
    let (enum_name, variant_name) = name.rsplit_once("::")?;
    if enum_name.is_empty() || variant_name.is_empty() {
        None
    } else {
        Some((enum_name, variant_name))
    }
}
