//! Main architectural checker — walks the AST and applies all rules.

use std::collections::HashMap;

use el_parser::{Expr, Program, Stmt, TypeExpr};

use crate::error::{ArchDiagnostic, Severity};
use crate::rule::{ArchRule, CallInfo, FnContext};
use crate::rules::{
    graph::{DuplicateActivateType, N1Detection},
    security::{AuthnWithoutAuthz, PublicFnWithActivate, SealedInLoop},
    swarm::{SwarmAgentIsolation, SwarmAgentNoSharedState, SwarmAgentNoSpawn},
    vbd::{AccessorMustNotCallManager, ExperienceMustNotCallExperience, ExperienceShouldReturnResult},
};

/// The main architectural checker. Instantiate once, call `check` per program.
pub struct ArchChecker {
    rules: Vec<Box<dyn ArchRule>>,
}

impl ArchChecker {
    /// Create an `ArchChecker` with all built-in rules registered.
    pub fn new() -> Self {
        Self {
            rules: vec![
                Box::new(AccessorMustNotCallManager),
                Box::new(ExperienceMustNotCallExperience),
                Box::new(ExperienceShouldReturnResult),
                Box::new(PublicFnWithActivate),
                Box::new(SealedInLoop),
                Box::new(AuthnWithoutAuthz),
                Box::new(N1Detection),
                Box::new(DuplicateActivateType),
                Box::new(SwarmAgentIsolation),
                Box::new(SwarmAgentNoSpawn),
                Box::new(SwarmAgentNoSharedState),
            ],
        }
    }

    /// Returns a reference to the registered rules (useful for introspection in tests).
    pub fn rules(&self) -> &[Box<dyn ArchRule>] {
        &self.rules
    }

    /// Run all rules against a parsed program and collect all diagnostics.
    pub fn check(&self, program: &Program) -> Vec<ArchDiagnostic> {
        // Step 1: build global fn_name → decorator names map (including impl methods).
        let all_fn_annotations = collect_fn_annotations(&program.stmts);

        // Step 2: gather all top-level + impl FnDef statements.
        let fn_defs = collect_fn_defs(&program.stmts);

        // Step 3: for each function, build FnContext and run every rule.
        let mut diagnostics = Vec::new();
        for (fn_name, decorators, return_type, body) in &fn_defs {
            let annotations: Vec<String> =
                decorators.iter().map(|d| d.name.clone()).collect();
            let body_calls = extract_calls(body, false);
            let activate_types = extract_activate_types(body, false);
            let has_sealed_in_loop = has_sealed_in_loop_body(body, false);
            let return_type_name = type_expr_name(return_type);

            let ctx = FnContext {
                fn_name: fn_name.as_str(),
                annotations: &annotations,
                body_calls: &body_calls,
                all_fn_annotations: &all_fn_annotations,
                activate_types: &activate_types,
                has_sealed_in_loop,
                return_type_name: &return_type_name,
            };

            for rule in &self.rules {
                diagnostics.extend(rule.check(&ctx));
            }
        }

        diagnostics
    }

    /// Returns true if any of the given diagnostics are errors.
    pub fn has_errors(diagnostics: &[ArchDiagnostic]) -> bool {
        diagnostics.iter().any(|d| d.severity == Severity::Error)
    }
}

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

// ── AST traversal helpers ─────────────────────────────────────────────────────

/// A collected function definition: (name, decorators, return_type, body).
type FnDef<'a> = (
    &'a String,
    &'a Vec<el_parser::Decorator>,
    &'a TypeExpr,
    &'a Vec<Stmt>,
);

/// Collect all FnDef statements from top-level and impl blocks.
fn collect_fn_defs<'a>(stmts: &'a [Stmt]) -> Vec<FnDef<'a>> {
    let mut out = Vec::new();
    for stmt in stmts {
        match stmt {
            Stmt::FnDef { name, decorators, return_type, body, .. } => {
                out.push((name, decorators, return_type, body));
            }
            Stmt::ImplDef { methods, .. } => {
                for m in methods {
                    if let Stmt::FnDef { name, decorators, return_type, body, .. } = m {
                        out.push((name, decorators, return_type, body));
                    }
                }
            }
            _ => {}
        }
    }
    out
}

/// Build a map from function name → list of decorator names, for all functions in the program.
fn collect_fn_annotations(stmts: &[Stmt]) -> HashMap<String, Vec<String>> {
    let mut map = HashMap::new();
    for stmt in stmts {
        match stmt {
            Stmt::FnDef { name, decorators, .. } => {
                let anns: Vec<String> = decorators.iter().map(|d| d.name.clone()).collect();
                map.insert(name.clone(), anns);
            }
            Stmt::ImplDef { methods, .. } => {
                for m in methods {
                    if let Stmt::FnDef { name, decorators, .. } = m {
                        let anns: Vec<String> = decorators.iter().map(|d| d.name.clone()).collect();
                        map.insert(name.clone(), anns);
                    }
                }
            }
            _ => {}
        }
    }
    map
}

/// Extract all function calls from a statement list.
/// `in_loop` tracks whether we are currently inside a for/while loop body.
fn extract_calls(stmts: &[Stmt], in_loop: bool) -> Vec<CallInfo> {
    let mut calls = Vec::new();
    for stmt in stmts {
        extract_calls_from_stmt(stmt, in_loop, &mut calls);
    }
    calls
}

fn extract_calls_from_stmt(stmt: &Stmt, in_loop: bool, out: &mut Vec<CallInfo>) {
    match stmt {
        Stmt::Let { value, .. } => extract_calls_from_expr(value, in_loop, out),
        Stmt::Return(expr, _) | Stmt::Expr(expr, _) | Stmt::Assert(expr, _) => {
            extract_calls_from_expr(expr, in_loop, out);
        }
        Stmt::FnDef { body, .. } => {
            // Nested function defs: walk but don't count as callee of the outer fn.
            for s in body {
                extract_calls_from_stmt(s, false, out);
            }
        }
        _ => {}
    }
}

fn extract_calls_from_expr(expr: &Expr, in_loop: bool, out: &mut Vec<CallInfo>) {
    match expr {
        Expr::Call { func, args } => {
            // Extract callee name
            let callee = expr_as_call_name(func);
            if let Some(name) = callee {
                out.push(CallInfo { callee: name, is_in_loop: in_loop });
            }
            // Recurse into func expression and args
            extract_calls_from_expr(func, in_loop, out);
            for a in args {
                extract_calls_from_expr(a, in_loop, out);
            }
        }
        Expr::Activate { type_name, .. } => {
            // Encode activate as a synthetic call so GRAPH-001 can detect in-loop activates.
            out.push(CallInfo {
                callee: format!("__activate__{type_name}"),
                is_in_loop: in_loop,
            });
        }
        Expr::BinOp { left, right, .. } => {
            extract_calls_from_expr(left, in_loop, out);
            extract_calls_from_expr(right, in_loop, out);
        }
        Expr::UnaryNot(inner) | Expr::Try(inner) => {
            extract_calls_from_expr(inner, in_loop, out);
        }
        Expr::Block(stmts) => {
            for s in stmts {
                extract_calls_from_stmt(s, in_loop, out);
            }
        }
        Expr::Sealed(stmts) => {
            // Sealed blocks are scanned but tracked separately for sealed-in-loop.
            for s in stmts {
                extract_calls_from_stmt(s, in_loop, out);
            }
        }
        Expr::If { cond, then, else_ } => {
            extract_calls_from_expr(cond, in_loop, out);
            extract_calls_from_expr(then, in_loop, out);
            if let Some(e) = else_ {
                extract_calls_from_expr(e, in_loop, out);
            }
        }
        Expr::Match { subject, arms } => {
            extract_calls_from_expr(subject, in_loop, out);
            for arm in arms {
                extract_calls_from_expr(&arm.body, in_loop, out);
            }
        }
        Expr::Field { object, .. } => extract_calls_from_expr(object, in_loop, out),
        Expr::Array(elems) => {
            for e in elems {
                extract_calls_from_expr(e, in_loop, out);
            }
        }
        Expr::Index { object, index } => {
            extract_calls_from_expr(object, in_loop, out);
            extract_calls_from_expr(index, in_loop, out);
        }
        Expr::Closure { body, .. } => {
            extract_calls_from_expr(body, in_loop, out);
        }
        Expr::MapLiteral(pairs) => {
            for (k, v) in pairs {
                extract_calls_from_expr(k, in_loop, out);
                extract_calls_from_expr(v, in_loop, out);
            }
        }
        Expr::Literal(_) | Expr::Ident(_) | Expr::Path { .. } => {}
        Expr::StructLit { fields, .. } => {
            for (_, e) in fields {
                extract_calls_from_expr(e, in_loop, out);
            }
        }
        Expr::With { base, updates } => {
            extract_calls_from_expr(base, in_loop, out);
            for (_, e) in updates {
                extract_calls_from_expr(e, in_loop, out);
            }
        }
        Expr::Reason { .. } => {}
        Expr::Parallel { entries } => {
            for (_, e) in entries {
                extract_calls_from_expr(e, in_loop, out);
            }
        }
        Expr::Trace { body, .. } => {
            for s in body {
                extract_calls_from_stmt(s, in_loop, out);
            }
        }
        Expr::UnaryBitNot(inner) => extract_calls_from_expr(inner, in_loop, out),
    }
}

/// Try to extract a simple callee name from a Call's `func` expression.
fn expr_as_call_name(expr: &Expr) -> Option<String> {
    match expr {
        Expr::Ident(name) => Some(name.clone()),
        Expr::Field { field, .. } => Some(field.clone()),
        Expr::Path { segments } => segments.last().cloned(),
        _ => None,
    }
}

/// Collect all `activate TypeName` type names from a statement list.
/// `in_loop` indicates whether we're currently inside a loop.
fn extract_activate_types(stmts: &[Stmt], in_loop: bool) -> Vec<String> {
    let mut types = Vec::new();
    for stmt in stmts {
        extract_activate_types_stmt(stmt, in_loop, &mut types);
    }
    types
}

fn extract_activate_types_stmt(stmt: &Stmt, in_loop: bool, out: &mut Vec<String>) {
    match stmt {
        Stmt::Let { value, .. } => extract_activate_types_expr(value, in_loop, out),
        Stmt::Return(expr, _) | Stmt::Expr(expr, _) | Stmt::Assert(expr, _) => {
            extract_activate_types_expr(expr, in_loop, out);
        }
        Stmt::FnDef { body, .. } => {
            for s in body {
                extract_activate_types_stmt(s, false, out);
            }
        }
        _ => {}
    }
}

fn extract_activate_types_expr(expr: &Expr, in_loop: bool, out: &mut Vec<String>) {
    match expr {
        Expr::Activate { type_name, .. } => {
            out.push(type_name.clone());
        }
        Expr::Call { func, args } => {
            extract_activate_types_expr(func, in_loop, out);
            for a in args {
                extract_activate_types_expr(a, in_loop, out);
            }
        }
        Expr::BinOp { left, right, .. } => {
            extract_activate_types_expr(left, in_loop, out);
            extract_activate_types_expr(right, in_loop, out);
        }
        Expr::UnaryNot(inner) | Expr::Try(inner) => {
            extract_activate_types_expr(inner, in_loop, out);
        }
        Expr::Block(stmts) => {
            for s in stmts {
                extract_activate_types_stmt(s, in_loop, out);
            }
        }
        Expr::Sealed(stmts) => {
            for s in stmts {
                extract_activate_types_stmt(s, in_loop, out);
            }
        }
        Expr::If { cond, then, else_ } => {
            extract_activate_types_expr(cond, in_loop, out);
            extract_activate_types_expr(then, in_loop, out);
            if let Some(e) = else_ {
                extract_activate_types_expr(e, in_loop, out);
            }
        }
        Expr::Match { subject, arms } => {
            extract_activate_types_expr(subject, in_loop, out);
            for arm in arms {
                extract_activate_types_expr(&arm.body, in_loop, out);
            }
        }
        Expr::Field { object, .. } => extract_activate_types_expr(object, in_loop, out),
        Expr::Array(elems) => {
            for e in elems {
                extract_activate_types_expr(e, in_loop, out);
            }
        }
        Expr::Index { object, index } => {
            extract_activate_types_expr(object, in_loop, out);
            extract_activate_types_expr(index, in_loop, out);
        }
        Expr::Closure { body, .. } => {
            extract_activate_types_expr(body, in_loop, out);
        }
        Expr::MapLiteral(pairs) => {
            for (k, v) in pairs {
                extract_activate_types_expr(k, in_loop, out);
                extract_activate_types_expr(v, in_loop, out);
            }
        }
        Expr::Literal(_) | Expr::Ident(_) | Expr::Path { .. } => {}
        Expr::StructLit { fields, .. } => {
            for (_, e) in fields {
                extract_activate_types_expr(e, in_loop, out);
            }
        }
        Expr::With { base, updates } => {
            extract_activate_types_expr(base, in_loop, out);
            for (_, e) in updates {
                extract_activate_types_expr(e, in_loop, out);
            }
        }
        Expr::Reason { .. } => {}
        Expr::Parallel { entries } => {
            for (_, e) in entries {
                extract_activate_types_expr(e, in_loop, out);
            }
        }
        Expr::Trace { body, .. } => {
            for s in body {
                extract_activate_types_stmt(s, in_loop, out);
            }
        }
        Expr::UnaryBitNot(inner) => extract_activate_types_expr(inner, in_loop, out),
    }
}

/// Returns true if any `sealed { }` block appears inside a loop in the given body.
fn has_sealed_in_loop_body(stmts: &[Stmt], in_loop: bool) -> bool {
    stmts.iter().any(|s| has_sealed_in_loop_stmt(s, in_loop))
}

fn has_sealed_in_loop_stmt(stmt: &Stmt, in_loop: bool) -> bool {
    match stmt {
        Stmt::Let { value, .. } => has_sealed_in_loop_expr(value, in_loop),
        Stmt::Return(expr, _) | Stmt::Expr(expr, _) | Stmt::Assert(expr, _) => {
            has_sealed_in_loop_expr(expr, in_loop)
        }
        Stmt::FnDef { body, .. } => {
            // Inner function defs reset loop context.
            body.iter().any(|s| has_sealed_in_loop_stmt(s, false))
        }
        _ => false,
    }
}

fn has_sealed_in_loop_expr(expr: &Expr, in_loop: bool) -> bool {
    match expr {
        Expr::Sealed(_) => in_loop,
        Expr::Call { func, args } => {
            has_sealed_in_loop_expr(func, in_loop)
                || args.iter().any(|a| has_sealed_in_loop_expr(a, in_loop))
        }
        Expr::BinOp { left, right, .. } => {
            has_sealed_in_loop_expr(left, in_loop) || has_sealed_in_loop_expr(right, in_loop)
        }
        Expr::UnaryNot(inner) | Expr::Try(inner) => has_sealed_in_loop_expr(inner, in_loop),
        Expr::Block(stmts) => stmts.iter().any(|s| has_sealed_in_loop_stmt(s, in_loop)),
        Expr::If { cond, then, else_ } => {
            has_sealed_in_loop_expr(cond, in_loop)
                || has_sealed_in_loop_expr(then, in_loop)
                || else_.as_deref().is_some_and(|e| has_sealed_in_loop_expr(e, in_loop))
        }
        Expr::Match { subject, arms } => {
            has_sealed_in_loop_expr(subject, in_loop)
                || arms.iter().any(|a| has_sealed_in_loop_expr(&a.body, in_loop))
        }
        Expr::Field { object, .. } => has_sealed_in_loop_expr(object, in_loop),
        Expr::Array(elems) => elems.iter().any(|e| has_sealed_in_loop_expr(e, in_loop)),
        Expr::Index { object, index } => {
            has_sealed_in_loop_expr(object, in_loop) || has_sealed_in_loop_expr(index, in_loop)
        }
        Expr::Closure { body, .. } => has_sealed_in_loop_expr(body, in_loop),
        Expr::MapLiteral(pairs) => pairs
            .iter()
            .any(|(k, v)| has_sealed_in_loop_expr(k, in_loop) || has_sealed_in_loop_expr(v, in_loop)),
        Expr::Activate { .. } | Expr::Literal(_) | Expr::Ident(_) | Expr::Path { .. } => false,
        Expr::StructLit { fields, .. } => fields
            .iter()
            .any(|(_, e)| has_sealed_in_loop_expr(e, in_loop)),
        Expr::With { base, updates } => {
            has_sealed_in_loop_expr(base, in_loop)
                || updates.iter().any(|(_, e)| has_sealed_in_loop_expr(e, in_loop))
        }
        Expr::Reason { .. } => false,
        Expr::Parallel { entries } => entries.iter().any(|(_, e)| has_sealed_in_loop_expr(e, in_loop)),
        Expr::Trace { body, .. } => body.iter().any(|s| has_sealed_in_loop_stmt(s, in_loop)),
        Expr::UnaryBitNot(inner) => has_sealed_in_loop_expr(inner, in_loop),
    }
}

/// Convert a `TypeExpr` to a display string for the return-type name check.
fn type_expr_name(te: &TypeExpr) -> String {
    match te {
        TypeExpr::Named(n) => n.clone(),
        TypeExpr::Result { .. } => "Result".to_string(),
        TypeExpr::Array(inner) => format!("[{}]", type_expr_name(inner)),
        TypeExpr::Optional(inner) => format!("{}?", type_expr_name(inner)),
        TypeExpr::Map { key, value } => {
            format!("Map<{}, {}>", type_expr_name(key), type_expr_name(value))
        }
        TypeExpr::Fn { .. } => "fn".to_string(),
        TypeExpr::TypeParam(n) => n.clone(),
    }
}