//! Static analysis: walk the function AST and compute costs.
//!
//! Counts allocations, syscalls, loop iterations, branches.
//! Computes worst-case latency estimate.

use syn::{
    Expr, ExprCall, ExprForLoop, ExprIf, ExprLoop, ExprMacro, ExprMatch, ExprMethodCall, ExprPath,
    ExprWhile, ItemFn, Lit, Pat, Stmt,
};

use crate::cost::*;
use crate::parse::{Constraint, RedlineAttr};

// Thread-local cost table set by analyze_function for the current analysis pass.
// This avoids threading the profile through every recursive call.
thread_local! {
    static ACTIVE_COSTS: std::cell::RefCell<CostTable> = const { std::cell::RefCell::new(CostTable::NS) };
    static ACTIVE_UNIT: std::cell::RefCell<String> = std::cell::RefCell::new("ns".to_string());
}

// Guard that resets thread-local cost state on drop (including panic unwind).
// Prevents cross-invocation contamination when rustc reuses threads.
struct CostGuard;
impl Drop for CostGuard {
    fn drop(&mut self) {
        ACTIVE_COSTS.with(|c| *c.borrow_mut() = CostTable::NS);
        ACTIVE_UNIT.with(|u| *u.borrow_mut() = "ns".to_string());
    }
}

fn with_costs<F, R>(f: F) -> R
where
    F: FnOnce(&CostTable) -> R,
{
    ACTIVE_COSTS.with(|c| f(&c.borrow()))
}

fn unit_name() -> String {
    ACTIVE_UNIT.with(|u| u.borrow().clone())
}

/// Analysis result for a function.
#[derive(Debug, Default)]
struct FnCost {
    /// Estimated worst-case cost in the active model's unit (ns or cycles).
    latency_cost: u64,
    /// Number of heap allocations.
    allocs: u64,
    /// Number of syscalls.
    syscalls: u64,
    /// Estimated stack usage in bytes.
    stack_bytes: u64,
}

/// Analyze a function against redline constraints.
/// Returns Ok(()) if all constraints are satisfied.
/// Returns Err(violations) with human-readable violation messages.
pub fn analyze_function(attr: &RedlineAttr, func: &ItemFn) -> Result<(), Vec<String>> {
    // Set the active cost profile for this analysis pass
    let table = attr.cost_profile.table().clone();
    let unit = attr.cost_profile.unit_name().to_string();
    ACTIVE_COSTS.with(|c| *c.borrow_mut() = table);
    ACTIVE_UNIT.with(|u| *u.borrow_mut() = unit);
    let _guard = CostGuard; // resets to NS defaults on drop (even on panic)

    let cost = analyze_block(&func.block.stmts);
    let mut violations = Vec::new();
    let unit = unit_name();

    for constraint in &attr.constraints {
        match constraint {
            Constraint::Latency(bound) => {
                if cost.latency_cost > bound.nanos {
                    violations.push(format!(
                        "estimated worst-case latency {}{} exceeds bound {}{} in `{}`",
                        cost.latency_cost, unit, bound.nanos, unit, func.sig.ident,
                    ));
                }
            }
            Constraint::Throughput(_bound) => {
                // Throughput analysis requires knowing input data size,
                // which is not available at the AST level without type info.
                // Planned for future version with LLVM integration.
            }
            Constraint::Allocs(max) => {
                if cost.allocs > *max {
                    violations.push(format!(
                        "function `{}` has {} allocation(s), limit is {}",
                        func.sig.ident, cost.allocs, max,
                    ));
                }
            }
            Constraint::Syscalls(max) => {
                if cost.syscalls > *max {
                    violations.push(format!(
                        "function `{}` has {} syscall(s), limit is {}",
                        func.sig.ident, cost.syscalls, max,
                    ));
                }
            }
            Constraint::MaxStack(max) => {
                if cost.stack_bytes > *max {
                    violations.push(format!(
                        "function `{}` uses {}B stack, limit is {}B",
                        func.sig.ident, cost.stack_bytes, max,
                    ));
                }
            }
        }
    }

    if violations.is_empty() {
        Ok(())
    } else {
        Err(violations)
    }
}

/// Analyze a block of statements.
fn analyze_block(stmts: &[Stmt]) -> FnCost {
    let mut cost = FnCost::default();

    for stmt in stmts {
        let stmt_cost = analyze_stmt(stmt);
        cost.latency_cost += stmt_cost.latency_cost;
        cost.allocs += stmt_cost.allocs;
        cost.syscalls += stmt_cost.syscalls;
        cost.stack_bytes += stmt_cost.stack_bytes;
    }

    cost
}

/// Analyze a single statement.
fn analyze_stmt(stmt: &Stmt) -> FnCost {
    match stmt {
        Stmt::Local(local) => {
            let mut cost = FnCost::default();
            // Local variable declaration: stack allocation
            cost.stack_bytes += estimate_pat_size(&local.pat);
            if let Some(init) = &local.init {
                let expr_cost = analyze_expr(&init.expr);
                cost.latency_cost += expr_cost.latency_cost;
                cost.allocs += expr_cost.allocs;
                cost.syscalls += expr_cost.syscalls;
            }
            cost
        }
        Stmt::Expr(expr, _) => analyze_expr(expr),
        Stmt::Item(_) => FnCost::default(), // inner items don't execute
        Stmt::Macro(mac) => analyze_macro_stmt(mac),
    }
}

/// Analyze an expression.
fn analyze_expr(expr: &Expr) -> FnCost {
    match expr {
        // Function calls
        Expr::Call(call) => analyze_call(call),

        // Method calls
        Expr::MethodCall(method) => analyze_method_call(method),

        // Macros (println!, vec!, format!, etc.)
        Expr::Macro(mac) => analyze_macro(mac),

        // For loops
        Expr::ForLoop(for_loop) => analyze_for_loop(for_loop),

        // While loops
        Expr::While(while_loop) => analyze_while_loop(while_loop),

        // Infinite loops
        Expr::Loop(loop_expr) => analyze_loop(loop_expr),

        // If/else
        Expr::If(if_expr) => analyze_if(if_expr),

        // Match
        Expr::Match(match_expr) => analyze_match(match_expr),

        // Block
        Expr::Block(block) => analyze_block(&block.block.stmts),

        // Binary operations (a + b, a * b, etc.)
        Expr::Binary(_) => with_costs(|c| FnCost {
            latency_cost: c.alu,
            ..Default::default()
        }),

        // Unary operations (!a, -a, *a)
        Expr::Unary(_) => with_costs(|c| FnCost {
            latency_cost: c.alu,
            ..Default::default()
        }),

        // Index (arr[i])
        Expr::Index(_) => with_costs(|c| FnCost {
            latency_cost: c.index,
            ..Default::default()
        }),

        // Return
        Expr::Return(ret) => {
            if let Some(expr) = &ret.expr {
                analyze_expr(expr)
            } else {
                FnCost::default()
            }
        }

        // Closures: walk the body
        Expr::Closure(closure) => analyze_expr(&closure.body),

        // Tuple: analyze each element
        Expr::Tuple(tuple) => {
            let mut cost = FnCost::default();
            for elem in &tuple.elems {
                let c = analyze_expr(elem);
                cost.latency_cost += c.latency_cost;
                cost.allocs += c.allocs;
                cost.syscalls += c.syscalls;
            }
            cost
        }

        // Everything else: assume free
        _ => FnCost::default(),
    }
}

/// Analyze a function call.
fn analyze_call(call: &ExprCall) -> FnCost {
    let mut cost = with_costs(|c| FnCost {
        latency_cost: c.call,
        ..Default::default()
    });
    let path_str = expr_to_path_string(&call.func);

    // Check if this is an allocating call
    if ALLOC_PATHS.iter().any(|p| path_str.contains(p)) {
        cost.allocs += 1;
        cost.latency_cost += with_costs(|c| c.alloc);
    }

    // Check if this is a syscall
    if SYSCALL_PATHS.iter().any(|p| path_str.contains(p)) {
        cost.syscalls += 1;
        cost.latency_cost += with_costs(|c| c.syscall);
    }

    // Check network
    if NET_PATHS.iter().any(|p| path_str.contains(p)) {
        cost.syscalls += 1;
        cost.latency_cost += with_costs(|c| c.net_syscall);
    }

    // Check sync (Mutex, RwLock)
    if path_str.contains("Mutex") || path_str.contains("lock") {
        cost.latency_cost += with_costs(|c| c.mutex_lock);
    }

    // Analyze arguments
    for arg in &call.args {
        let arg_cost = analyze_expr(arg);
        cost.latency_cost += arg_cost.latency_cost;
        cost.allocs += arg_cost.allocs;
        cost.syscalls += arg_cost.syscalls;
    }

    cost
}

/// Analyze a method call.
fn analyze_method_call(method: &ExprMethodCall) -> FnCost {
    let mut cost = with_costs(|c| FnCost {
        latency_cost: c.call,
        ..Default::default()
    });
    let method_name = method.method.to_string();
    let receiver_str = expr_to_path_string(&method.receiver);

    // Allocating methods (only on heap types, not Copy types)
    let heap_alloc_methods = ["to_string", "to_owned", "collect", "extend", "append"];
    if heap_alloc_methods.iter().any(|m| method_name == *m) {
        cost.allocs += 1;
        cost.latency_cost += with_costs(|c| c.alloc);
    }

    // clone: only counts as alloc if receiver looks like a heap type
    if method_name == "clone" && !receiver_str.is_empty() {
        let heap_types = ["String", "Vec", "Box", "HashMap", "BTreeMap", "Rc", "Arc"];
        if heap_types.iter().any(|t| receiver_str.contains(t)) {
            cost.allocs += 1;
            cost.latency_cost += with_costs(|c| c.alloc);
        }
    }

    // push/insert: may or may not allocate, count conservatively
    if method_name == "push" || method_name == "insert" {
        cost.latency_cost += with_costs(|c| c.alloc);
    }

    // Sync methods (only for known sync types)
    if (method_name == "lock" || method_name == "try_lock")
        && (receiver_str.contains("Mutex") || receiver_str.contains("mutex"))
    {
        cost.latency_cost += with_costs(|c| c.mutex_lock);
    }

    // HashMap/BTreeMap specific methods
    if method_name == "contains_key" || method_name == "entry" {
        cost.latency_cost += with_costs(|c| c.hashmap_lookup);
    }

    // Analyze receiver
    let recv_cost = analyze_expr(&method.receiver);
    cost.latency_cost += recv_cost.latency_cost;
    cost.allocs += recv_cost.allocs;
    cost.syscalls += recv_cost.syscalls;

    // Analyze arguments
    for arg in &method.args {
        let arg_cost = analyze_expr(arg);
        cost.latency_cost += arg_cost.latency_cost;
        cost.allocs += arg_cost.allocs;
        cost.syscalls += arg_cost.syscalls;
    }

    cost
}

/// Analyze a for loop. Tries to extract the iteration count from range literals.
fn analyze_for_loop(for_loop: &ExprForLoop) -> FnCost {
    let iterations = estimate_range_size(&for_loop.expr);
    let body_cost = analyze_block(&for_loop.body.stmts);
    let branch = with_costs(|c| c.branch);

    FnCost {
        latency_cost: body_cost.latency_cost.saturating_mul(iterations) + branch,
        allocs: body_cost.allocs.saturating_mul(iterations),
        syscalls: body_cost.syscalls.saturating_mul(iterations),
        stack_bytes: body_cost.stack_bytes,
    }
}

/// Analyze a while loop. Assumes worst case of 1000 iterations (configurable).
fn analyze_while_loop(while_loop: &ExprWhile) -> FnCost {
    let iterations = 1000; // pessimistic default
    let body_cost = analyze_block(&while_loop.body.stmts);
    let branch = with_costs(|c| c.branch);

    FnCost {
        latency_cost: body_cost.latency_cost.saturating_mul(iterations) + branch,
        allocs: body_cost.allocs.saturating_mul(iterations),
        syscalls: body_cost.syscalls.saturating_mul(iterations),
        stack_bytes: body_cost.stack_bytes,
    }
}

/// Analyze an infinite loop. Assumes worst case of 1000 iterations.
fn analyze_loop(loop_expr: &ExprLoop) -> FnCost {
    let iterations = 1000;
    let body_cost = analyze_block(&loop_expr.body.stmts);

    FnCost {
        latency_cost: body_cost.latency_cost.saturating_mul(iterations),
        allocs: body_cost.allocs.saturating_mul(iterations),
        syscalls: body_cost.syscalls.saturating_mul(iterations),
        stack_bytes: body_cost.stack_bytes,
    }
}

/// Analyze if/else. Takes the worst-case branch.
fn analyze_if(if_expr: &ExprIf) -> FnCost {
    let cond_cost = analyze_expr(&if_expr.cond);
    let then_cost = analyze_block(&if_expr.then_branch.stmts);
    let branch = with_costs(|c| c.branch);

    let else_cost = if let Some((_, else_expr)) = &if_expr.else_branch {
        analyze_expr(else_expr)
    } else {
        FnCost::default()
    };

    // Worst case: whichever branch is more expensive
    FnCost {
        latency_cost: cond_cost.latency_cost
            + branch
            + then_cost.latency_cost.max(else_cost.latency_cost),
        allocs: cond_cost.allocs + then_cost.allocs.max(else_cost.allocs),
        syscalls: cond_cost.syscalls + then_cost.syscalls.max(else_cost.syscalls),
        stack_bytes: cond_cost.stack_bytes + then_cost.stack_bytes.max(else_cost.stack_bytes),
    }
}

/// Analyze match. Takes the worst-case arm.
fn analyze_match(match_expr: &ExprMatch) -> FnCost {
    let scrutinee_cost = analyze_expr(&match_expr.expr);
    let branch = with_costs(|c| c.branch);
    let mut worst_arm = FnCost::default();

    for arm in &match_expr.arms {
        let arm_cost = analyze_expr(&arm.body);
        worst_arm.latency_cost = worst_arm.latency_cost.max(arm_cost.latency_cost);
        worst_arm.allocs = worst_arm.allocs.max(arm_cost.allocs);
        worst_arm.syscalls = worst_arm.syscalls.max(arm_cost.syscalls);
        worst_arm.stack_bytes = worst_arm.stack_bytes.max(arm_cost.stack_bytes);
    }

    FnCost {
        latency_cost: scrutinee_cost.latency_cost + branch + worst_arm.latency_cost,
        allocs: scrutinee_cost.allocs + worst_arm.allocs,
        syscalls: scrutinee_cost.syscalls + worst_arm.syscalls,
        stack_bytes: scrutinee_cost.stack_bytes + worst_arm.stack_bytes,
    }
}

/// Analyze a macro by name.
fn analyze_macro_by_name(name: &str) -> FnCost {
    with_costs(|c| match name {
        "println" | "print" | "eprintln" | "eprint" => FnCost {
            latency_cost: c.println,
            syscalls: 1,
            ..Default::default()
        },
        "format" => FnCost {
            latency_cost: c.format,
            allocs: 1,
            ..Default::default()
        },
        "vec" => FnCost {
            latency_cost: c.alloc,
            allocs: 1,
            ..Default::default()
        },
        "write" | "writeln" => FnCost {
            latency_cost: c.format + c.syscall,
            syscalls: 1,
            ..Default::default()
        },
        "panic" | "unreachable" | "todo" | "unimplemented" => FnCost {
            latency_cost: c.format,
            allocs: 1,
            ..Default::default()
        },
        _ => FnCost {
            latency_cost: c.call,
            ..Default::default()
        },
    })
}

/// Analyze a macro expression.
fn analyze_macro(mac: &ExprMacro) -> FnCost {
    let name = mac
        .mac
        .path
        .segments
        .last()
        .map(|s| s.ident.to_string())
        .unwrap_or_default();
    analyze_macro_by_name(&name)
}

/// Analyze a macro statement.
fn analyze_macro_stmt(mac: &syn::StmtMacro) -> FnCost {
    let name = mac
        .mac
        .path
        .segments
        .last()
        .map(|s| s.ident.to_string())
        .unwrap_or_default();
    analyze_macro_by_name(&name)
}

/// Try to extract the size of a range expression (0..100 -> 100).
fn estimate_range_size(expr: &Expr) -> u64 {
    match expr {
        Expr::Range(range) => {
            let start = range
                .start
                .as_ref()
                .and_then(|e| extract_lit_int(e))
                .unwrap_or(0u64);
            let end = range
                .end
                .as_ref()
                .and_then(|e| extract_lit_int(e))
                .unwrap_or(1000u64);
            end.saturating_sub(start)
        }
        // Method call like (0..n).into_iter() or iter()
        Expr::MethodCall(mc) => estimate_range_size(&mc.receiver),
        _ => 1000, // unknown, assume pessimistic
    }
}

/// Try to extract an integer literal from an expression.
fn extract_lit_int(expr: &Expr) -> Option<u64> {
    if let Expr::Lit(lit) = expr {
        if let Lit::Int(int) = &lit.lit {
            return int.base10_parse().ok();
        }
    }
    None
}

/// Estimate stack size of a pattern binding.
fn estimate_pat_size(pat: &Pat) -> u64 {
    match pat {
        Pat::Ident(_) => USIZE_SIZE, // default: one word
        Pat::Tuple(tuple) => tuple.elems.len() as u64 * USIZE_SIZE,
        Pat::Type(pt) => estimate_pat_size(&pt.pat),
        _ => USIZE_SIZE,
    }
}

/// Convert an expression to a dotted path string for pattern matching.
fn expr_to_path_string(expr: &Expr) -> String {
    match expr {
        Expr::Path(ExprPath { path, .. }) => path
            .segments
            .iter()
            .map(|s| s.ident.to_string())
            .collect::<Vec<_>>()
            .join("::"),
        _ => String::new(),
    }
}
