horkos 0.2.0

//! Type checking and taint analysis for Horkos.
//!
//! This module enforces:
//! - Type correctness
//! - Taint tracking (`Unverified<T>` cannot flow into secure contexts)
//! - Security invariants from smart constructors
//! - Cross-file symbol resolution

use crate::ast::{BinaryOp, UnaryOp, *};
use crate::errors::{find_similar, Diagnostic, DiagnosticBuilder};
use serde::Serialize;
use std::collections::HashMap;
use std::path::{Path, PathBuf};

// ============================================================================
// Cross-File Symbol Table
// ============================================================================

/// Symbols exported from a compiled module (file).
#[derive(Debug, Clone, Default)]
pub struct ModuleExports {
    /// The file path this module was compiled from
    pub path: PathBuf,
    /// Exported symbols: name → type
    pub symbols: HashMap<String, ResolvedType>,
}

impl ModuleExports {
    pub fn new(path: impl Into<PathBuf>) -> Self {
        Self {
            path: path.into(),
            symbols: HashMap::new(),
        }
    }

    /// Add an exported symbol
    pub fn export(&mut self, name: impl Into<String>, typ: ResolvedType) {
        self.symbols.insert(name.into(), typ);
    }

    /// Look up an exported symbol
    pub fn get(&self, name: &str) -> Option<&ResolvedType> {
        self.symbols.get(name)
    }
}

/// Global symbol table for cross-file type resolution.
///
/// Stores exports from all compiled modules, allowing one file
/// to reference types from another.
#[derive(Debug, Clone, Default)]
pub struct GlobalSymbolTable {
    /// Map from canonical file path to its exports
    modules: HashMap<PathBuf, ModuleExports>,
    /// Map from import path to canonical file path (for resolution)
    import_paths: HashMap<String, PathBuf>,
}

impl GlobalSymbolTable {
    pub fn new() -> Self {
        Self::default()
    }

    /// Register a module's exports after compilation
    pub fn register(&mut self, exports: ModuleExports) {
        self.modules.insert(exports.path.clone(), exports);
    }

    /// Register an import path mapping
    pub fn register_import_path(&mut self, import_path: &str, canonical_path: &Path) {
        self.import_paths
            .insert(import_path.to_string(), canonical_path.to_path_buf());
    }

    /// Look up exports for a file by its import path
    pub fn get_by_import(&self, import_path: &str) -> Option<&ModuleExports> {
        self.import_paths
            .get(import_path)
            .and_then(|canonical| self.modules.get(canonical))
    }

    /// Look up exports for a file by its canonical path
    pub fn get_by_path(&self, path: &Path) -> Option<&ModuleExports> {
        self.modules.get(path)
    }

    /// Look up a specific symbol from an imported module
    pub fn resolve_import(&self, import_path: &str, symbol: &str) -> Option<&ResolvedType> {
        self.get_by_import(import_path)
            .and_then(|exports| exports.get(symbol))
    }
}

// ============================================================================
// Type Checker
// ============================================================================

/// Result of type checking, including both the typed program and any warnings.
#[derive(Debug)]
pub struct TypeCheckResult {
    pub program: TypedProgram,
    pub preferred_overrides: Vec<PreferredOverride>,
}

/// Check a program for type errors and taint violations.
pub fn check(
    ast: &Program,
    source: &str,
    filename: &str,
    no_hcl: bool,
) -> Result<TypeCheckResult, Vec<Diagnostic>> {
    let mut checker = TypeChecker::new(source, filename, None);
    checker.set_strict_mode(no_hcl);
    checker.init_stdlib();
    checker.check_program(ast)
}

/// Check a program with access to global symbols (cross-file).
pub fn check_with_globals(
    ast: &Program,
    source: &str,
    filename: &str,
    globals: Option<&GlobalSymbolTable>,
    no_hcl: bool,
) -> Result<TypeCheckResult, Vec<Diagnostic>> {
    let mut checker = TypeChecker::new(source, filename, globals);
    checker.set_strict_mode(no_hcl);
    checker.init_stdlib();
    checker.check_program(ast)
}

/// Extract exports from a typed program
pub fn extract_exports(program: &TypedProgram, path: impl Into<PathBuf>) -> ModuleExports {
    let mut exports = ModuleExports::new(path);

    for stmt in &program.statements {
        match stmt {
            TypedStatement::ValDecl {
                name,
                resolved_type,
                ..
            } => {
                exports.export(name.clone(), resolved_type.clone());
            }
            TypedStatement::Module { name, body, .. } => {
                // Export module as a record of its contents
                let mut fields = HashMap::new();
                for inner in body {
                    if let TypedStatement::ValDecl {
                        name: inner_name,
                        resolved_type,
                        ..
                    } = inner
                    {
                        fields.insert(inner_name.clone(), resolved_type.clone());
                    }
                }
                exports.export(name.clone(), ResolvedType::Record(fields));
            }
            _ => {}
        }
    }

    exports
}

/// A type-checked program with resolved types.
#[derive(Debug, Serialize)]
pub struct TypedProgram {
    pub statements: Vec<TypedStatement>,
}

impl TypedProgram {
    /// Count the number of unsafe blocks in the program.
    pub fn count_unsafe_blocks(&self) -> usize {
        let mut count = 0;
        for stmt in &self.statements {
            count += count_unsafe_in_statement(stmt);
        }
        count
    }
}

fn count_unsafe_in_statement(stmt: &TypedStatement) -> usize {
    match stmt {
        TypedStatement::ValDecl { value, .. } => count_unsafe_in_expr(value),
        TypedStatement::Module { body, .. } => body.iter().map(count_unsafe_in_statement).sum(),
        TypedStatement::Assert { condition, .. } => count_unsafe_in_expr(condition),
        TypedStatement::Import { .. } => 0,
        TypedStatement::HclBlock { .. } => 1,
        TypedStatement::Unsafe { body, .. } => {
            // Unsafe block itself counts as 1, plus whatever is inside
            1 + body.iter().map(count_unsafe_in_statement).sum::<usize>()
        }
    }
}

fn count_unsafe_in_expr(expr: &TypedExpr) -> usize {
    match &expr.kind {
        TypedExprKind::Unsafe { body, .. } => 1 + count_unsafe_in_expr(body),
        TypedExprKind::FuncCall { callee, args } => {
            count_unsafe_in_expr(callee)
                + args
                    .iter()
                    .map(|a| count_unsafe_in_expr(&a.value))
                    .sum::<usize>()
        }
        TypedExprKind::MemberAccess { object, .. } => count_unsafe_in_expr(object),
        TypedExprKind::Binary { left, right, .. } => {
            count_unsafe_in_expr(left) + count_unsafe_in_expr(right)
        }
        TypedExprKind::Unary { operand, .. } => count_unsafe_in_expr(operand),
        TypedExprKind::List(items) => items.iter().map(count_unsafe_in_expr).sum(),
        TypedExprKind::Record(fields) => {
            fields.iter().map(|f| count_unsafe_in_expr(&f.value)).sum()
        }
        TypedExprKind::Lambda { body, .. } => count_unsafe_in_expr(body),
        TypedExprKind::If {
            condition,
            then_branch,
            else_branch,
        } => {
            count_unsafe_in_expr(condition)
                + count_unsafe_in_expr(then_branch)
                + count_unsafe_in_expr(else_branch)
        }
        _ => 0,
    }
}

/// A type-checked statement.
#[derive(Debug, Serialize)]
pub enum TypedStatement {
    Import {
        path: String,
        alias: String,
        span: Span,
    },
    ValDecl {
        name: String,
        resolved_type: ResolvedType,
        value: TypedExpr,
        span: Span,
    },
    Module {
        name: String,
        body: Vec<TypedStatement>,
        span: Span,
    },
    /// Compile-time assertion - validated but not emitted to Terraform
    Assert {
        condition: TypedExpr,
        message: String,
        span: Span,
    },
    /// Inline HCL block - always tainted
    HclBlock {
        reason: String,
        content: String,
        span: Span,
    },
    /// Unsafe block - explicit acknowledgement of risk
    Unsafe {
        reason: String,
        body: Vec<TypedStatement>,
        span: Span,
    },
}

/// A type-checked expression with its resolved type.
#[derive(Debug, Serialize)]
pub struct TypedExpr {
    pub kind: TypedExprKind,
    pub resolved_type: ResolvedType,
    pub span: Span,
}

/// Type-checked expression kinds.
#[derive(Debug, Serialize)]
pub enum TypedExprKind {
    Literal(Literal),
    Identifier(String),
    MemberAccess {
        object: Box<TypedExpr>,
        field: String,
    },
    FuncCall {
        callee: Box<TypedExpr>,
        args: Vec<TypedArg>,
    },
    Lambda {
        params: Vec<String>,
        body: Box<TypedExpr>,
    },
    Binary {
        left: Box<TypedExpr>,
        op: BinaryOp,
        right: Box<TypedExpr>,
    },
    Unary {
        op: UnaryOp,
        operand: Box<TypedExpr>,
    },
    List(Vec<TypedExpr>),
    Record(Vec<TypedRecordField>),
    /// Unsafe expression that unwraps taint: Unverified<T> → T
    Unsafe {
        reason: String,
        body: Box<TypedExpr>,
    },
    /// If expression: if cond then a else b
    If {
        condition: Box<TypedExpr>,
        then_branch: Box<TypedExpr>,
        else_branch: Box<TypedExpr>,
    },
}

#[derive(Debug, Serialize)]
pub struct TypedArg {
    pub name: Option<String>,
    pub value: TypedExpr,
}

#[derive(Debug, Serialize)]
pub struct TypedRecordField {
    pub name: String,
    pub value: TypedExpr,
}

/// Resolved types in the type system.
#[derive(Debug, Clone, PartialEq, Serialize)]
pub enum ResolvedType {
    // Primitives
    String,
    Number,
    Bool,
    Void,

    // Infrastructure types
    Bucket,
    SecureBucket,
    Vpc,
    Subnet,
    SecurityGroup,
    InternetGateway,
    IamRole,
    IamPolicy,
    LogGroup, // CloudWatch Log Group
    Database, // RDS Database Instance

    // Container types
    List(Box<ResolvedType>),
    Record(HashMap<String, ResolvedType>),

    // AWS Tags: Record<String, String> - string keys and values only
    Tags,

    // Union types (for parameters that accept multiple types)
    Union(Box<ResolvedType>, Box<ResolvedType>),

    // The taint wrapper - this is the key type for security
    Unverified(Box<ResolvedType>),

    // Module/namespace reference
    Module(String),

    // Function type (for stdlib)
    Function {
        params: Vec<(String, ResolvedType)>,
        returns: Box<ResolvedType>,
    },

    // Unknown/error type (for error recovery)
    Unknown,
}

impl ResolvedType {
    /// Check if this type is tainted (unverified).
    pub fn is_tainted(&self) -> bool {
        matches!(self, ResolvedType::Unverified(_))
    }

    /// Unwrap the taint, returning the inner type.
    pub fn unwrap_taint(&self) -> &ResolvedType {
        match self {
            ResolvedType::Unverified(inner) => inner,
            other => other,
        }
    }

    /// Wrap in Unverified if not already tainted.
    pub fn taint(self) -> ResolvedType {
        if self.is_tainted() {
            self
        } else {
            ResolvedType::Unverified(Box::new(self))
        }
    }

    /// Human-readable type name for error messages.
    pub fn display_name(&self) -> String {
        match self {
            ResolvedType::String => "String".to_string(),
            ResolvedType::Number => "Number".to_string(),
            ResolvedType::Bool => "Bool".to_string(),
            ResolvedType::Void => "Void".to_string(),
            ResolvedType::Bucket => "Bucket".to_string(),
            ResolvedType::SecureBucket => "SecureBucket".to_string(),
            ResolvedType::Vpc => "Vpc".to_string(),
            ResolvedType::Subnet => "Subnet".to_string(),
            ResolvedType::SecurityGroup => "SecurityGroup".to_string(),
            ResolvedType::InternetGateway => "InternetGateway".to_string(),
            ResolvedType::IamRole => "IamRole".to_string(),
            ResolvedType::IamPolicy => "IamPolicy".to_string(),
            ResolvedType::LogGroup => "LogGroup".to_string(),
            ResolvedType::Database => "Database".to_string(),
            ResolvedType::List(inner) => format!("List<{}>", inner.display_name()),
            ResolvedType::Record(_) => "Record".to_string(),
            ResolvedType::Tags => "Tags".to_string(),
            ResolvedType::Union(a, b) => format!("{} | {}", a.display_name(), b.display_name()),
            ResolvedType::Unverified(inner) => format!("Unverified<{}>", inner.display_name()),
            ResolvedType::Module(name) => format!("module {}", name),
            ResolvedType::Function { .. } => "Function".to_string(),
            ResolvedType::Unknown => "Unknown".to_string(),
        }
    }
}

/// Info about an overridden preferred parameter (for compile-time warnings).
#[derive(Debug, Clone)]
pub struct PreferredOverride {
    pub resource_name: String,
    pub param_name: String,
    pub recommended: String,
    pub span: Span,
}

/// Type checking context.
struct TypeChecker<'a> {
    _source: &'a str,
    filename: &'a str,
    errors: Vec<Diagnostic>,

    // Symbol table: name -> type
    symbols: HashMap<String, ResolvedType>,

    // Import alias -> import path (for cross-file resolution)
    imports: HashMap<String, String>,

    // Global symbol table for cross-file imports
    globals: Option<&'a GlobalSymbolTable>,

    // Are we inside an unsafe block? If so, stores the span of the outer block.
    unsafe_span: Option<Span>,

    // Assertion failures (collected and reported at end)
    assertion_failures: Vec<(Span, String)>,

    // Preferred param overrides (collected for info messages)
    preferred_overrides: Vec<PreferredOverride>,

    /// Whether we are currently inside an unsafe block
    in_unsafe_block: bool,
    /// Whether HCL blocks are forbidden (strict mode)
    no_hcl: bool,
}

impl<'a> TypeChecker<'a> {
    fn new(source: &'a str, filename: &'a str, globals: Option<&'a GlobalSymbolTable>) -> Self {
        let mut checker = Self {
            _source: source,
            filename,
            errors: Vec::new(),
            symbols: HashMap::new(),
            imports: HashMap::new(),
            globals,
            unsafe_span: None,
            assertion_failures: Vec::new(),
            preferred_overrides: Vec::new(),
            in_unsafe_block: false,
            no_hcl: false,
        };

        // Initialize stdlib
        checker.init_stdlib();
        checker
    }

    /// Enable strict mode (forbid HCL blocks)
    pub fn set_strict_mode(&mut self, strict: bool) {
        self.no_hcl = strict;
    }

    fn check_program(&mut self, program: &Program) -> Result<TypeCheckResult, Vec<Diagnostic>> {
        let mut statements = Vec::new();

        for stmt in &program.statements {
            match self.check_statement(stmt) {
                Ok(typed) => statements.push(typed),
                Err(()) => {
                    // Error already recorded
                }
            }
        }

        // Report all assertion failures at the end
        if !self.assertion_failures.is_empty() {
            for (span, message) in std::mem::take(&mut self.assertion_failures) {
                self.errors.push(DiagnosticBuilder::assertion_failed(
                    &message,
                    span,
                    self.filename,
                ));
            }
        }

        if self.errors.is_empty() {
            Ok(TypeCheckResult {
                program: TypedProgram { statements },
                preferred_overrides: std::mem::take(&mut self.preferred_overrides),
            })
        } else {
            Err(std::mem::take(&mut self.errors))
        }
    }

    fn check_statement(&mut self, stmt: &Spanned<Statement>) -> Result<TypedStatement, ()> {
        match &stmt.node {
            Statement::Import(import) => self.check_import(import, stmt.span),
            Statement::ValDecl(decl) => self.check_val_decl(decl, stmt.span),
            Statement::Module(module) => self.check_module(module, stmt.span),
            Statement::Assert(assert_stmt) => self.check_assert(assert_stmt, stmt.span),
            Statement::HclBlock(hcl) => self.check_hcl_block(hcl, stmt.span),
            Statement::Unsafe(unsafe_stmt) => self.check_unsafe_stmt(unsafe_stmt, stmt.span),
        }
    }

    fn check_unsafe_stmt(&mut self, stmt: &UnsafeStmt, span: Span) -> Result<TypedStatement, ()> {
        let reason = stmt.reason.node.clone();

        // Enter unsafe context
        let was_unsafe = self.in_unsafe_block;
        self.in_unsafe_block = true;

        let mut body = Vec::new();
        for s in &stmt.body {
            if let Ok(typed) = self.check_statement(s) {
                body.push(typed);
            }
        }

        // Restore previous context
        self.in_unsafe_block = was_unsafe;

        Ok(TypedStatement::Unsafe { reason, body, span })
    }

    fn check_hcl_block(&mut self, hcl: &HclBlock, span: Span) -> Result<TypedStatement, ()> {
        let reason = hcl.reason.node.clone();
        let content = hcl.content.node.clone();

        // Strict Mode: Forbid all HCL blocks if enabled
        if self.no_hcl {
            self.error_diag(
                Diagnostic::error_at(
                    "HCL blocks are forbidden in strict mode",
                    span,
                    self.filename,
                )
                .with_code(crate::errors::ErrorCode::ShadowBanned)
                .with_primary_label("strict mode enabled"),
            );
            // We can return early or continue to find other errors. Returning early is safer.
            return Ok(TypedStatement::HclBlock {
                reason,
                content,
                span,
            });
        }

        // 1. Shadow Ban: Check for resources that should be native
        // Simple string parsing to find `resource "type"`
        // We look for: resource "aws_s3_bucket"
        // This is a heuristic but sufficient for the MVP "Shadow Ban"

        // List of resources that Horkos supports natively and should NOT be used via HCL
        let native_resources = [
            "aws_s3_bucket",
            "aws_vpc",
            "aws_subnet",
            "aws_security_group",
            "aws_internet_gateway",
            "aws_cloudwatch_log_group",
            // Add more as we implement them
        ];

        for resource in native_resources {
            // Check for `resource "type"` or `resource 'type'`
            // We handle basic whitespace
            let pattern1 = format!("resource \"{}\"", resource);
            let pattern2 = format!("resource '{}'", resource);

            if content.contains(&pattern1) || content.contains(&pattern2) {
                // Shadow Ban Logic:
                // 1. If in unsafe block, ALLOW (Emergency Hatch)
                // 2. Otherwise, error (Shadow Ban)

                if !self.in_unsafe_block {
                    self.error_diag(
                        Diagnostic::error_at(
                            format!("Shadow Ban: `{}` is supported natively by Horkos", resource),
                            span,
                            self.filename,
                        )
                        .with_code(crate::errors::ErrorCode::ShadowBanned)
                        .with_primary_label("use the native Horkos resource instead")
                        .with_help(format!(
                            "replace with `{}.create...` or wrap in `unsafe(\"reason\") {{ ... }}` to bypass",
                            self.get_module_for_resource(resource)
                        )),
                    );
                }
            }
        }

        // 2. Taint Tracking: HCL blocks are inherently unsafe/tainted
        // We don't need to wrap the *statement* in Unverified, but any *references* to it
        // (if we supported referencing HCL blocks directly) would be tainted.
        // For now, HCL blocks are just emitted as-is.
        // However, if they contain `data` sources that are referenced, those would be handled by `Terraform.ref` (which is separate).

        Ok(TypedStatement::HclBlock {
            reason,
            content,
            span,
        })
    }

    fn get_module_for_resource(&self, resource: &str) -> &str {
        match resource {
            "aws_s3_bucket" => "S3",
            "aws_vpc" | "aws_subnet" | "aws_security_group" | "aws_internet_gateway" => "Network",
            "aws_cloudwatch_log_group" => "CloudWatch",
            _ => "Unknown",
        }
    }

    fn check_import(&mut self, import: &ImportStmt, span: Span) -> Result<TypedStatement, ()> {
        let import_path = &import.path.node;
        let alias = import
            .alias
            .as_ref()
            .map(|a| a.node.clone())
            .unwrap_or_else(|| {
                // Derive alias from path
                import_path
                    .trim_end_matches(".hk")
                    .trim_end_matches(".tf")
                    .split('/')
                    .next_back()
                    .unwrap_or("import")
                    .to_string()
            });

        // Store import path for later resolution
        self.imports.insert(alias.clone(), import_path.clone());

        // Check if we have global symbols for this import (cross-file .hk imports)
        let module_type = if import_path.ends_with(".hk") {
            // Try to resolve from global symbol table
            if let Some(globals) = self.globals {
                if let Some(exports) = globals.get_by_import(import_path) {
                    // Build a record type from the exports (NOT tainted!)
                    let fields: HashMap<String, ResolvedType> = exports
                        .symbols
                        .iter()
                        .map(|(k, v)| (k.clone(), v.clone()))
                        .collect();
                    ResolvedType::Record(fields)
                } else {
                    // .hk file not in global table - treat as tainted for now
                    ResolvedType::Unverified(Box::new(ResolvedType::Record(HashMap::new())))
                }
            } else {
                // No global table - single file mode, treat as tainted
                ResolvedType::Unverified(Box::new(ResolvedType::Record(HashMap::new())))
            }
        } else {
            // External imports (.tf files) are always tainted
            ResolvedType::Unverified(Box::new(ResolvedType::Record(HashMap::new())))
        };

        self.symbols.insert(alias.clone(), module_type);

        Ok(TypedStatement::Import {
            path: import_path.clone(),
            alias,
            span,
        })
    }

    fn check_val_decl(&mut self, decl: &ValDecl, span: Span) -> Result<TypedStatement, ()> {
        let typed_value = self.check_expr(&decl.value)?;
        let resolved_type = typed_value.resolved_type.clone();

        // Check type annotation if present
        if let Some(ann) = &decl.type_ann {
            let expected = self.resolve_type_expr(ann);
            if !self.types_compatible(&expected, &resolved_type) {
                self.error_diag(
                    DiagnosticBuilder::type_mismatch(
                        &expected.display_name(),
                        &resolved_type.display_name(),
                        decl.value.span,
                        self.filename,
                    )
                    .with_label(
                        ann.span,
                        format!(
                            "expected `{}` due to type annotation",
                            expected.display_name()
                        ),
                    ),
                );
            }
        }

        // Add to symbol table
        self.symbols
            .insert(decl.name.node.clone(), resolved_type.clone());

        Ok(TypedStatement::ValDecl {
            name: decl.name.node.clone(),
            resolved_type,
            value: typed_value,
            span,
        })
    }

    fn check_module(&mut self, module: &ModuleDecl, span: Span) -> Result<TypedStatement, ()> {
        // TODO: Implement module scoping
        let mut body = Vec::new();
        for stmt in &module.body {
            if let Ok(typed) = self.check_statement(stmt) {
                body.push(typed);
            }
        }

        Ok(TypedStatement::Module {
            name: module.name.node.clone(),
            body,
            span,
        })
    }

    fn check_assert(&mut self, assert_stmt: &AssertStmt, span: Span) -> Result<TypedStatement, ()> {
        let typed_condition = self.check_expr(&assert_stmt.condition)?;

        // Condition must be a Bool
        if !self.types_compatible(&ResolvedType::Bool, &typed_condition.resolved_type) {
            self.error_diag(
                DiagnosticBuilder::type_mismatch(
                    "Bool",
                    &typed_condition.resolved_type.display_name(),
                    assert_stmt.condition.span,
                    self.filename,
                )
                .with_note("assert condition must be a boolean expression"),
            );
        }

        let message = assert_stmt.message.node.clone();

        // Try to evaluate the condition at compile time
        if let Some(false) = self.try_eval_bool(&typed_condition) {
            // Assertion failed at compile time - collect it
            self.assertion_failures.push((span, message.clone()));
        }

        Ok(TypedStatement::Assert {
            condition: typed_condition,
            message,
            span,
        })
    }

    /// Try to evaluate a boolean expression at compile time
    #[allow(clippy::only_used_in_recursion)]
    fn try_eval_bool(&self, expr: &TypedExpr) -> Option<bool> {
        match &expr.kind {
            TypedExprKind::Literal(crate::ast::Literal::Bool(b)) => Some(*b),
            TypedExprKind::Binary { left, op, right } => {
                match op {
                    crate::ast::BinaryOp::Eq => {
                        // Try to compare literals
                        if let (TypedExprKind::Literal(l), TypedExprKind::Literal(r)) =
                            (&left.kind, &right.kind)
                        {
                            Some(l == r)
                        } else {
                            None
                        }
                    }
                    crate::ast::BinaryOp::NotEq => {
                        if let (TypedExprKind::Literal(l), TypedExprKind::Literal(r)) =
                            (&left.kind, &right.kind)
                        {
                            Some(l != r)
                        } else {
                            None
                        }
                    }
                    crate::ast::BinaryOp::And => {
                        match (self.try_eval_bool(left), self.try_eval_bool(right)) {
                            (Some(l), Some(r)) => Some(l && r),
                            (Some(false), _) | (_, Some(false)) => Some(false),
                            _ => None,
                        }
                    }
                    crate::ast::BinaryOp::Or => {
                        match (self.try_eval_bool(left), self.try_eval_bool(right)) {
                            (Some(l), Some(r)) => Some(l || r),
                            (Some(true), _) | (_, Some(true)) => Some(true),
                            _ => None,
                        }
                    }
                    crate::ast::BinaryOp::Lt
                    | crate::ast::BinaryOp::LtEq
                    | crate::ast::BinaryOp::Gt
                    | crate::ast::BinaryOp::GtEq => {
                        if let (
                            TypedExprKind::Literal(crate::ast::Literal::Number(l)),
                            TypedExprKind::Literal(crate::ast::Literal::Number(r)),
                        ) = (&left.kind, &right.kind)
                        {
                            Some(match op {
                                crate::ast::BinaryOp::Lt => l < r,
                                crate::ast::BinaryOp::LtEq => l <= r,
                                crate::ast::BinaryOp::Gt => l > r,
                                crate::ast::BinaryOp::GtEq => l >= r,
                                _ => unreachable!(),
                            })
                        } else {
                            None
                        }
                    }
                    _ => None,
                }
            }
            TypedExprKind::Unary { op, operand } => {
                if matches!(op, crate::ast::UnaryOp::Not) {
                    self.try_eval_bool(operand).map(|b| !b)
                } else {
                    None
                }
            }
            _ => None, // Can't evaluate at compile time
        }
    }

    fn check_expr(&mut self, expr: &Spanned<Expr>) -> Result<TypedExpr, ()> {
        let (kind, resolved_type) = match &expr.node {
            Expr::Literal(lit) => {
                let t = match lit {
                    Literal::String(_) => ResolvedType::String,
                    Literal::Number(_) => ResolvedType::Number,
                    Literal::Bool(_) => ResolvedType::Bool,
                };
                (TypedExprKind::Literal(lit.clone()), t)
            }

            Expr::Identifier(name) => {
                let t = self.symbols.get(name).cloned().unwrap_or_else(|| {
                    // Try to find a similar variable name for "did you mean?" suggestions
                    let symbol_names = self.get_symbol_names();
                    let similar = find_similar(name, &symbol_names, 3);
                    self.error_diag(DiagnosticBuilder::undefined_variable(
                        name,
                        expr.span,
                        self.filename,
                        similar,
                    ));
                    ResolvedType::Unknown
                });
                (TypedExprKind::Identifier(name.clone()), t)
            }

            Expr::MemberAccess { object, field } => {
                let typed_obj = self.check_expr(object)?;
                let (field_type, allow) =
                    self.resolve_member(&typed_obj.resolved_type, &field.node, field.span);

                if !allow && self.unsafe_span.is_none() {
                    self.error_diag(DiagnosticBuilder::tainted_access(
                        &field.node,
                        &typed_obj.resolved_type.display_name(),
                        expr.span,
                        self.filename,
                    ));
                }

                (
                    TypedExprKind::MemberAccess {
                        object: Box::new(typed_obj),
                        field: field.node.clone(),
                    },
                    field_type,
                )
            }

            Expr::FuncCall { callee, args } => {
                let typed_callee = self.check_expr(callee)?;
                let mut typed_args = Vec::new();

                for arg in args {
                    let typed_value = self.check_expr(&arg.value)?;
                    typed_args.push(TypedArg {
                        name: arg.name.as_ref().map(|n| n.node.clone()),
                        value: typed_value,
                    });
                }

                // Check for security-weakening parameters (requires unsafe)
                if self.unsafe_span.is_none() {
                    if let Some((module, function)) = self.extract_module_function(callee) {
                        if let Some(info) =
                            crate::stdlib::check_security_params(&module, &function, &typed_args)
                        {
                            self.error_diag(
                                Diagnostic::error_at(
                                    format!(
                                        "`{}` weakens security and requires `unsafe`",
                                        info.param_name
                                    ),
                                    info.span, // Point to the specific argument
                                    self.filename,
                                )
                                .with_code(crate::errors::ErrorCode::SecurityViolation)
                                .with_primary_label(format!(
                                    "`{}` is a security-sensitive parameter",
                                    info.param_name
                                ))
                                .with_help("wrap in unsafe with justification"),
                            );
                        }
                    }
                }

                // Check for preferred param overrides (collect for info messages)
                if let Some((module, function)) = self.extract_module_function(callee) {
                    let overrides =
                        crate::stdlib::check_preferred_params(&module, &function, &typed_args);
                    for info in overrides {
                        // Get the resource name from the first positional arg or use function name
                        let resource_name = typed_args
                            .first()
                            .and_then(|a| {
                                if let TypedExprKind::Literal(crate::ast::Literal::String(s)) =
                                    &a.value.kind
                                {
                                    Some(s.clone())
                                } else {
                                    None
                                }
                            })
                            .unwrap_or_else(|| function.clone());

                        self.preferred_overrides.push(PreferredOverride {
                            resource_name,
                            param_name: info.param_name.to_string(),
                            recommended: info.recommended,
                            span: info.span, // Use the argument's span, not the whole call
                        });
                    }
                }

                // Check for unknown parameters (typos, etc.)
                if let Some((module, function)) = self.extract_module_function(callee) {
                    let unknown =
                        crate::stdlib::check_unknown_params(&module, &function, &typed_args);
                    for info in unknown {
                        let mut diag = Diagnostic::error_at(
                            format!("unknown parameter `{}`", info.param_name),
                            info.span,
                            self.filename,
                        )
                        .with_code(crate::errors::ErrorCode::UnknownParameter)
                        .with_primary_label("not a valid parameter");

                        // Add "did you mean" suggestion if available
                        if let Some(suggestion) = info.suggestion {
                            diag = diag.with_help(format!("did you mean `{}`?", suggestion));
                        }

                        // Add note with available parameters
                        let available = info.known_params.join(", ");
                        diag = diag.with_note(format!("available parameters: {}", available));

                        self.error_diag(diag);
                    }
                }

                // Resolve function return type
                let return_type = self.resolve_call_type(&typed_callee, &typed_args, expr.span);

                (
                    TypedExprKind::FuncCall {
                        callee: Box::new(typed_callee),
                        args: typed_args,
                    },
                    return_type,
                )
            }

            Expr::Lambda { params, body } => {
                // Save current symbols to restore after
                let saved_symbols: Vec<_> = params
                    .iter()
                    .filter_map(|p| {
                        self.symbols
                            .get(&p.node)
                            .cloned()
                            .map(|t| (p.node.clone(), t))
                    })
                    .collect();

                // Bind parameters with Unknown type for now
                // (In a full implementation, we'd infer from usage or require annotations)
                for param in params {
                    self.symbols
                        .insert(param.node.clone(), ResolvedType::Unknown);
                }

                let typed_body = self.check_expr(body)?;

                // Remove parameter bindings
                for param in params {
                    self.symbols.remove(&param.node);
                }

                // Restore any shadowed symbols
                for (name, typ) in saved_symbols {
                    self.symbols.insert(name, typ);
                }

                let param_types: Vec<_> = params
                    .iter()
                    .map(|p| (p.node.clone(), ResolvedType::Unknown))
                    .collect();

                let func_type = ResolvedType::Function {
                    params: param_types,
                    returns: Box::new(typed_body.resolved_type.clone()),
                };

                (
                    TypedExprKind::Lambda {
                        params: params.iter().map(|p| p.node.clone()).collect(),
                        body: Box::new(typed_body),
                    },
                    func_type,
                )
            }

            Expr::Binary { left, op, right } => {
                let typed_left = self.check_expr(left)?;
                let typed_right = self.check_expr(right)?;

                let result_type = match op {
                    // Arithmetic operators return Number
                    BinaryOp::Add | BinaryOp::Sub | BinaryOp::Mul | BinaryOp::Div => {
                        // Type check: both operands should be numbers
                        if typed_left.resolved_type != ResolvedType::Number {
                            self.error_diag(DiagnosticBuilder::invalid_operator(
                                op.name(),
                                &typed_left.resolved_type.display_name(),
                                &typed_right.resolved_type.display_name(),
                                left.span,
                                self.filename,
                            ));
                        } else if typed_right.resolved_type != ResolvedType::Number {
                            self.error_diag(DiagnosticBuilder::invalid_operator(
                                op.name(),
                                &typed_left.resolved_type.display_name(),
                                &typed_right.resolved_type.display_name(),
                                right.span,
                                self.filename,
                            ));
                        }
                        ResolvedType::Number
                    }

                    // Comparison operators return Bool
                    BinaryOp::Eq | BinaryOp::NotEq => ResolvedType::Bool,

                    BinaryOp::Lt | BinaryOp::LtEq | BinaryOp::Gt | BinaryOp::GtEq => {
                        // Operands should be comparable (numbers)
                        if typed_left.resolved_type != ResolvedType::Number {
                            self.error_diag(DiagnosticBuilder::invalid_operator(
                                op.name(),
                                &typed_left.resolved_type.display_name(),
                                &typed_right.resolved_type.display_name(),
                                left.span,
                                self.filename,
                            ));
                        } else if typed_right.resolved_type != ResolvedType::Number {
                            self.error_diag(DiagnosticBuilder::invalid_operator(
                                op.name(),
                                &typed_left.resolved_type.display_name(),
                                &typed_right.resolved_type.display_name(),
                                right.span,
                                self.filename,
                            ));
                        }
                        ResolvedType::Bool
                    }

                    // Logical operators return Bool
                    BinaryOp::And | BinaryOp::Or => {
                        if typed_left.resolved_type != ResolvedType::Bool {
                            self.error_diag(DiagnosticBuilder::invalid_operator(
                                op.name(),
                                &typed_left.resolved_type.display_name(),
                                &typed_right.resolved_type.display_name(),
                                left.span,
                                self.filename,
                            ));
                        } else if typed_right.resolved_type != ResolvedType::Bool {
                            self.error_diag(DiagnosticBuilder::invalid_operator(
                                op.name(),
                                &typed_left.resolved_type.display_name(),
                                &typed_right.resolved_type.display_name(),
                                right.span,
                                self.filename,
                            ));
                        }
                        ResolvedType::Bool
                    }
                };

                (
                    TypedExprKind::Binary {
                        left: Box::new(typed_left),
                        op: *op,
                        right: Box::new(typed_right),
                    },
                    result_type,
                )
            }

            Expr::Unary { op, operand } => {
                let typed_operand = self.check_expr(operand)?;

                let result_type = match op {
                    UnaryOp::Not => {
                        if typed_operand.resolved_type != ResolvedType::Bool {
                            self.error_diag(
                                DiagnosticBuilder::type_mismatch(
                                    "Bool",
                                    &typed_operand.resolved_type.display_name(),
                                    operand.span,
                                    self.filename,
                                )
                                .with_note("the `!` operator requires a Bool operand"),
                            );
                        }
                        ResolvedType::Bool
                    }
                    UnaryOp::Neg => {
                        if typed_operand.resolved_type != ResolvedType::Number {
                            self.error_diag(
                                DiagnosticBuilder::type_mismatch(
                                    "Number",
                                    &typed_operand.resolved_type.display_name(),
                                    operand.span,
                                    self.filename,
                                )
                                .with_note("the `-` operator requires a Number operand"),
                            );
                        }
                        ResolvedType::Number
                    }
                };

                (
                    TypedExprKind::Unary {
                        op: *op,
                        operand: Box::new(typed_operand),
                    },
                    result_type,
                )
            }

            Expr::List(elements) => {
                let mut typed_elements = Vec::new();
                let mut element_type = ResolvedType::Unknown;

                for elem in elements {
                    let typed = self.check_expr(elem)?;
                    if element_type == ResolvedType::Unknown {
                        element_type = typed.resolved_type.clone();
                    }
                    typed_elements.push(typed);
                }

                (
                    TypedExprKind::List(typed_elements),
                    ResolvedType::List(Box::new(element_type)),
                )
            }

            Expr::Record(fields) => {
                let mut typed_fields = Vec::new();
                let mut field_types = HashMap::new();

                for field in fields {
                    let typed_value = self.check_expr(&field.value)?;
                    field_types.insert(field.name.node.clone(), typed_value.resolved_type.clone());
                    typed_fields.push(TypedRecordField {
                        name: field.name.node.clone(),
                        value: typed_value,
                    });
                }

                (
                    TypedExprKind::Record(typed_fields),
                    ResolvedType::Record(field_types),
                )
            }

            Expr::Unsafe { reason, body } => {
                // Reject nested unsafe blocks for clear accountability
                if let Some(outer_span) = self.unsafe_span {
                    let mut diag = Diagnostic::error_at(
                        "nested unsafe blocks are not allowed",
                        expr.span,
                        self.filename,
                    )
                    .with_code(crate::errors::ErrorCode::NestedUnsafe)
                    .with_primary_label("this unsafe block is nested");

                    // Show where the outer block is
                    diag = diag.with_label(outer_span, "outer unsafe block starts here");

                    diag = diag.with_help(
                        "split into separate statements, each with its own unsafe block and justification",
                    );

                    self.error_diag(diag);
                }

                // Enter unsafe context (save previous state for error recovery)
                let was_unsafe = self.unsafe_span;
                self.unsafe_span = Some(expr.span);

                // Type check the body
                let typed_body = self.check_expr(body)?;

                // Restore context
                self.unsafe_span = was_unsafe;

                // Unwrap taint if present - this is the key operation!
                // unsafe("reason") { x } where x: Unverified<T> produces T
                let result_type = match &typed_body.resolved_type {
                    ResolvedType::Unverified(inner) => (**inner).clone(),
                    other => other.clone(),
                };

                (
                    TypedExprKind::Unsafe {
                        reason: reason.node.clone(),
                        body: Box::new(typed_body),
                    },
                    result_type,
                )
            }

            Expr::If {
                condition,
                then_branch,
                else_branch,
            } => {
                // Type check condition - must be Bool
                let typed_condition = self.check_expr(condition)?;
                if typed_condition.resolved_type != ResolvedType::Bool {
                    self.error_diag(
                        DiagnosticBuilder::type_mismatch(
                            "Bool",
                            &typed_condition.resolved_type.display_name(),
                            condition.span,
                            self.filename,
                        )
                        .with_note("if condition must be a boolean expression"),
                    );
                }

                // Type check both branches
                let typed_then = self.check_expr(then_branch)?;
                let typed_else = self.check_expr(else_branch)?;

                // Both branches must have compatible types
                let result_type = if self
                    .types_compatible(&typed_then.resolved_type, &typed_else.resolved_type)
                {
                    typed_then.resolved_type.clone()
                } else if self
                    .types_compatible(&typed_else.resolved_type, &typed_then.resolved_type)
                {
                    typed_else.resolved_type.clone()
                } else {
                    self.error_diag(
                        DiagnosticBuilder::type_mismatch(
                            &typed_then.resolved_type.display_name(),
                            &typed_else.resolved_type.display_name(),
                            else_branch.span,
                            self.filename,
                        )
                        .with_label(
                            then_branch.span,
                            format!(
                                "then branch has type `{}`",
                                typed_then.resolved_type.display_name()
                            ),
                        )
                        .with_note("if branches must have compatible types"),
                    );
                    typed_then.resolved_type.clone()
                };

                (
                    TypedExprKind::If {
                        condition: Box::new(typed_condition),
                        then_branch: Box::new(typed_then),
                        else_branch: Box::new(typed_else),
                    },
                    result_type,
                )
            }
        };

        Ok(TypedExpr {
            kind,
            resolved_type,
            span: expr.span,
        })
    }

    #[allow(clippy::only_used_in_recursion)]
    fn resolve_type_expr(&self, expr: &Spanned<TypeExpr>) -> ResolvedType {
        match &expr.node {
            TypeExpr::Named(name) => match name.as_str() {
                "String" => ResolvedType::String,
                "Number" => ResolvedType::Number,
                "Bool" => ResolvedType::Bool,
                "Bucket" => ResolvedType::Bucket,
                "SecureBucket" => ResolvedType::SecureBucket,
                "Vpc" => ResolvedType::Vpc,
                "Subnet" => ResolvedType::Subnet,
                "SecurityGroup" => ResolvedType::SecurityGroup,
                "InternetGateway" => ResolvedType::InternetGateway,
                "LogGroup" => ResolvedType::LogGroup,
                "Tags" => ResolvedType::Tags,
                _ => ResolvedType::Unknown,
            },
            TypeExpr::Generic { name, args } => match name.as_str() {
                "Unverified" if args.len() == 1 => {
                    let inner = self.resolve_type_expr(&args[0]);
                    ResolvedType::Unverified(Box::new(inner))
                }
                "List" if args.len() == 1 => {
                    let inner = self.resolve_type_expr(&args[0]);
                    ResolvedType::List(Box::new(inner))
                }
                _ => ResolvedType::Unknown,
            },
        }
    }

    /// Extract module and function names from a callee expression like `S3.createBucket`.
    fn extract_module_function(&self, callee: &Spanned<Expr>) -> Option<(String, String)> {
        if let Expr::MemberAccess { object, field } = &callee.node {
            if let Expr::Identifier(module) = &object.node {
                return Some((module.clone(), field.node.clone()));
            }
        }
        None
    }

    fn resolve_member(
        &mut self,
        base: &ResolvedType,
        field: &str,
        span: Span,
    ) -> (ResolvedType, bool) {
        match base {
            // Stdlib modules
            ResolvedType::Module(name) => {
                let result = crate::stdlib::resolve_member(name, field);
                (result, true)
            }

            // Tainted values propagate taint
            ResolvedType::Unverified(inner) => {
                match inner.as_ref() {
                    // For tainted records, allow any field access and return Unverified<Unknown>
                    // This is for imported legacy infrastructure where we don't know the schema
                    ResolvedType::Record(_) => (
                        ResolvedType::Unverified(Box::new(ResolvedType::Unknown)),
                        false,
                    ),
                    // For other tainted types, recurse
                    _ => {
                        let (inner_type, _) = self.resolve_member(inner, field, span);
                        (inner_type.taint(), false) // false = needs unsafe
                    }
                }
            }

            // Record field access
            ResolvedType::Record(fields) => {
                if let Some(t) = fields.get(field) {
                    (t.clone(), true)
                } else {
                    self.error_diag(DiagnosticBuilder::unknown_field(
                        field,
                        "Record",
                        span,
                        self.filename,
                    ));
                    (ResolvedType::Unknown, true)
                }
            }

            // List methods
            ResolvedType::List(element_type) => {
                match field {
                    "map" => {
                        // map takes a function (T) -> U and returns List<U>
                        // For now, we accept any function and return List<Unknown>
                        (
                            ResolvedType::Function {
                                params: vec![(
                                    "fn".to_string(),
                                    ResolvedType::Function {
                                        params: vec![(
                                            "item".to_string(),
                                            (**element_type).clone(),
                                        )],
                                        returns: Box::new(ResolvedType::Unknown),
                                    },
                                )],
                                returns: Box::new(ResolvedType::List(Box::new(
                                    ResolvedType::Unknown,
                                ))),
                            },
                            true,
                        )
                    }
                    "filter" => {
                        // filter takes a predicate (T) -> Bool and returns List<T>
                        (
                            ResolvedType::Function {
                                params: vec![(
                                    "predicate".to_string(),
                                    ResolvedType::Function {
                                        params: vec![(
                                            "item".to_string(),
                                            (**element_type).clone(),
                                        )],
                                        returns: Box::new(ResolvedType::Bool),
                                    },
                                )],
                                returns: Box::new(ResolvedType::List(element_type.clone())),
                            },
                            true,
                        )
                    }
                    "find" => {
                        // find takes a predicate (T) -> Bool and returns T (or Unknown if not found)
                        (
                            ResolvedType::Function {
                                params: vec![(
                                    "predicate".to_string(),
                                    ResolvedType::Function {
                                        params: vec![(
                                            "item".to_string(),
                                            (**element_type).clone(),
                                        )],
                                        returns: Box::new(ResolvedType::Bool),
                                    },
                                )],
                                returns: element_type.clone(),
                            },
                            true,
                        )
                    }
                    "forEach" => {
                        // forEach takes a function (T) -> Void and returns Void
                        (
                            ResolvedType::Function {
                                params: vec![(
                                    "fn".to_string(),
                                    ResolvedType::Function {
                                        params: vec![(
                                            "item".to_string(),
                                            (**element_type).clone(),
                                        )],
                                        returns: Box::new(ResolvedType::Void),
                                    },
                                )],
                                returns: Box::new(ResolvedType::Void),
                            },
                            true,
                        )
                    }
                    "any" | "all" => {
                        // any/all take a predicate (T) -> Bool and return Bool
                        (
                            ResolvedType::Function {
                                params: vec![(
                                    "predicate".to_string(),
                                    ResolvedType::Function {
                                        params: vec![(
                                            "item".to_string(),
                                            (**element_type).clone(),
                                        )],
                                        returns: Box::new(ResolvedType::Bool),
                                    },
                                )],
                                returns: Box::new(ResolvedType::Bool),
                            },
                            true,
                        )
                    }
                    "reduce" => {
                        // reduce takes initial value and reducer function
                        (
                            ResolvedType::Function {
                                params: vec![
                                    ("initial".to_string(), ResolvedType::Unknown),
                                    (
                                        "reducer".to_string(),
                                        ResolvedType::Function {
                                            params: vec![
                                                ("acc".to_string(), ResolvedType::Unknown),
                                                ("item".to_string(), (**element_type).clone()),
                                            ],
                                            returns: Box::new(ResolvedType::Unknown),
                                        },
                                    ),
                                ],
                                returns: Box::new(ResolvedType::Unknown),
                            },
                            true,
                        )
                    }
                    "length" | "size" | "count" => (ResolvedType::Number, true),
                    "isEmpty" => (ResolvedType::Bool, true),
                    "first" | "last" => ((**element_type).clone(), true),
                    "concat" => {
                        // concat takes another list and returns List<T>
                        (
                            ResolvedType::Function {
                                params: vec![(
                                    "other".to_string(),
                                    ResolvedType::List(element_type.clone()),
                                )],
                                returns: Box::new(ResolvedType::List(element_type.clone())),
                            },
                            true,
                        )
                    }
                    _ => {
                        self.error_diag(DiagnosticBuilder::unknown_field(
                            field,
                            &format!("List<{}>", element_type.display_name()),
                            span,
                            self.filename,
                        ).with_help("available methods: map, filter, find, forEach, any, all, reduce, length, isEmpty, first, last, concat"));
                        (ResolvedType::Unknown, true)
                    }
                }
            }

            // String methods
            ResolvedType::String => {
                match field {
                    "length" | "size" => (ResolvedType::Number, true),
                    "concat" => {
                        // concat takes another string and returns String
                        (
                            ResolvedType::Function {
                                params: vec![("other".to_string(), ResolvedType::String)],
                                returns: Box::new(ResolvedType::String),
                            },
                            true,
                        )
                    }
                    "toUpper" | "toLower" | "trim" => (ResolvedType::String, true),
                    "contains" | "startsWith" | "endsWith" => (
                        ResolvedType::Function {
                            params: vec![("substr".to_string(), ResolvedType::String)],
                            returns: Box::new(ResolvedType::Bool),
                        },
                        true,
                    ),
                    _ => {
                        self.error_diag(DiagnosticBuilder::unknown_field(
                            field,
                            "String",
                            span,
                            self.filename,
                        ).with_help("available methods: length, concat, toUpper, toLower, trim, contains, startsWith, endsWith"));
                        (ResolvedType::Unknown, true)
                    }
                }
            }

            _ => {
                self.error_diag(DiagnosticBuilder::unknown_field(
                    field,
                    &base.display_name(),
                    span,
                    self.filename,
                ));
                (ResolvedType::Unknown, true)
            }
        }
    }

    fn resolve_call_type(
        &mut self,
        callee: &TypedExpr,
        args: &[TypedArg],
        span: Span,
    ) -> ResolvedType {
        let callee_type = &callee.resolved_type;

        match callee_type {
            ResolvedType::Function { params, returns } => {
                // Try to get resource definition for better arity checking
                let (min_args, max_args) = if let Some((module, function)) =
                    self.extract_module_function_from_typed(callee)
                {
                    if let Some(def) = crate::stdlib::get_resource_definition(&module, &function) {
                        (def.min_args(), def.max_args())
                    } else {
                        // Legacy: all params required
                        (params.len(), params.len())
                    }
                } else {
                    // Non-module function: all params required
                    (params.len(), params.len())
                };

                // Check arity using resource definition
                if args.len() < min_args {
                    self.error_diag(DiagnosticBuilder::arity_mismatch(
                        min_args,
                        args.len(),
                        span,
                        self.filename,
                    ));
                    return ResolvedType::Unknown;
                }

                if args.len() > max_args {
                    self.error_diag(DiagnosticBuilder::arity_mismatch(
                        max_args,
                        args.len(),
                        span,
                        self.filename,
                    ));
                    return ResolvedType::Unknown;
                }

                // Check positional argument types (only check up to number of args provided)
                for (i, arg) in args.iter().enumerate() {
                    // Get expected type - first try by name, then by position
                    let expected_ty = if let Some(ref name) = arg.name {
                        params.iter().find(|(n, _)| n == name).map(|(_, t)| t)
                    } else if i < params.len() {
                        Some(&params[i].1)
                    } else {
                        None
                    };

                    if let Some(expected_ty) = expected_ty {
                        let arg_ty = &arg.value.resolved_type;
                        let param_name = arg
                            .name
                            .as_deref()
                            .unwrap_or(params.get(i).map(|(n, _)| n.as_str()).unwrap_or("arg"));

                        // Check if argument is tainted
                        if arg_ty.is_tainted() && !matches!(arg_ty, ResolvedType::Function { .. }) {
                            self.error_diag(DiagnosticBuilder::tainted_argument(
                                param_name,
                                &arg_ty.display_name(),
                                &expected_ty.display_name(),
                                arg.value.span,
                                self.filename,
                            ));
                            continue;
                        }

                        // Check type compatibility
                        if !self.types_compatible_for_call(arg_ty, expected_ty) {
                            self.error_diag(DiagnosticBuilder::type_mismatch(
                                &self.type_to_string(expected_ty),
                                &self.type_to_string(arg_ty),
                                arg.value.span,
                                self.filename,
                            ));
                        }
                    }
                }

                // Check for required security parameters (e.g., flowLogs for VPC)
                if let Some((module, function)) = self.extract_module_function_from_typed(callee) {
                    self.check_required_security_params(&module, &function, args, span);
                }

                (**returns).clone()
            }
            ResolvedType::Unknown => ResolvedType::Unknown,
            _ => {
                self.error_diag(DiagnosticBuilder::not_callable(
                    &self.type_to_string(callee_type),
                    span,
                    self.filename,
                ));
                ResolvedType::Unknown
            }
        }
    }

    /// Extract module and function name from a typed callee expression.
    fn extract_module_function_from_typed(&self, callee: &TypedExpr) -> Option<(String, String)> {
        if let TypedExprKind::MemberAccess { object, field } = &callee.kind {
            if let TypedExprKind::Identifier(module) = &object.kind {
                return Some((module.clone(), field.clone()));
            }
        }
        None
    }

    /// Check for required security parameters that must be provided (unless in unsafe).
    /// For example, Network.createVpc requires flowLogs to be provided.
    fn check_required_security_params(
        &mut self,
        module: &str,
        function: &str,
        args: &[TypedArg],
        span: Span,
    ) {
        // VPC requires flowLogs unless in unsafe
        if module == "Network" && function == "createVpc" && self.unsafe_span.is_none() {
            let has_flow_logs = args
                .iter()
                .any(|arg| arg.name.as_deref() == Some("flowLogs"));

            if !has_flow_logs {
                self.error_diag(
                    Diagnostic::error_at(
                        "VPC requires `flowLogs` parameter for network auditing",
                        span,
                        self.filename,
                    )
                    .with_code(crate::errors::ErrorCode::SecurityViolation)
                    .with_primary_label("missing required `flowLogs` parameter")
                    .with_help("Provide flowLogs with a Bucket or LogGroup:\n\n    Network.createVpc(\"name\", cidr: \"...\", flowLogs: logBucket)")
                );
            }
        }

        // Subnet: public: true requires gateway parameter
        if module == "Network" && function == "createSubnet" {
            let public_arg = args.iter().find(|arg| {
                arg.name.as_deref() == Some("public")
                    && matches!(
                        &arg.value.kind,
                        TypedExprKind::Literal(crate::ast::Literal::Bool(true))
                    )
            });

            if let Some(public_arg) = public_arg {
                let has_gateway = args
                    .iter()
                    .any(|arg| arg.name.as_deref() == Some("gateway"));

                if !has_gateway {
                    self.error_diag(
                        Diagnostic::error_at(
                            "`public: true` requires `gateway` parameter",
                            public_arg.value.span,  // Point to the `true` value
                            self.filename,
                        )
                        .with_code(crate::errors::ErrorCode::MissingRequiredParam)
                        .with_primary_label("public subnet requires an internet gateway")
                        .with_help("Create an internet gateway:\n\n    val igw = Network.createInternetGateway(vpc: vpc)\n    val subnet = Network.createSubnet(..., public: true, gateway: igw)")
                    );
                }
            }
        }

        // SecurityGroup: check ingress rules for dangerous patterns
        if module == "Network" && function == "createSecurityGroup" && self.unsafe_span.is_none() {
            if let Some(ingress_arg) = args
                .iter()
                .find(|a| a.name.as_deref() == Some("ingressRules"))
            {
                self.check_security_group_ingress_rules(&ingress_arg.value, span);
            }
        }
    }

    /// Check security group ingress rules for dangerous patterns
    ///
    /// Safe patterns (no unsafe needed):
    /// - Private CIDRs: 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8
    /// - Security group references: sourceSecurityGroup
    /// - Specific public IPs or ranges (e.g., 203.0.113.10/32, 198.51.100.0/24)
    ///
    /// Unsafe required for:
    /// - 0.0.0.0/0 (entire internet)
    /// - Critical ports (SSH, RDP, DB) get stronger warnings
    fn check_security_group_ingress_rules(&mut self, rules_expr: &TypedExpr, _span: Span) {
        if let TypedExprKind::List(rules) = &rules_expr.kind {
            for rule in rules {
                if let TypedExprKind::Record(fields) = &rule.kind {
                    let port = self.get_record_number_field(fields, "port");
                    let to_port = self.get_record_number_field(fields, "toPort").or(port);
                    let cidr = self.get_record_string_field(fields, "cidr");
                    let has_source_sg = fields.iter().any(|f| f.name == "sourceSecurityGroup");

                    // Security group reference is always safe
                    if has_source_sg {
                        continue;
                    }

                    // No CIDR means must have sourceSecurityGroup (handled above)
                    let cidr = match cidr {
                        Some(c) => c,
                        None => continue,
                    };

                    // Check for dangerous patterns:
                    // 1. 0.0.0.0/0 (entire internet) - always requires unsafe
                    // 2. IP range (not /32) with all ports - requires unsafe

                    let protocol = self
                        .get_record_string_field(fields, "protocol")
                        .unwrap_or_else(|| "tcp".to_string());
                    let is_all_ports = port == Some(0) && (protocol == "-1" || protocol == "all");
                    let is_range = !cidr.ends_with("/32"); // Not a single IP

                    // Specific single IPs (/32) are always safe - explicit whitelist
                    // Ranges are safe UNLESS they open all ports
                    if cidr != "0.0.0.0/0" && !(is_range && is_all_ports) {
                        continue;
                    }

                    // Special case: IP range with all ports
                    if is_range && is_all_ports && cidr != "0.0.0.0/0" {
                        self.error_diag(
                            Diagnostic::error_at(
                                format!("all ports open to IP range `{}`", cidr),
                                rule.span,  // Point to the specific rule
                                self.filename,
                            )
                            .with_code(crate::errors::ErrorCode::PortExposed)
                            .with_primary_label("opens all ports to a range of IPs")
                            .with_help("Restrict to specific ports, or use unsafe:\n\n    { port: 443, cidr: \"...\", description: \"HTTPS only\" }")
                        );
                        continue;
                    }

                    // 0.0.0.0/0 - check for critical ports vs regular ports
                    let critical_ports: &[(i64, &str)] = &[
                        (22, "SSH"),
                        (3389, "RDP"),
                        (3306, "MySQL"),
                        (5432, "PostgreSQL"),
                        (1433, "MSSQL"),
                        (27017, "MongoDB"),
                        (6379, "Redis"),
                        (11211, "Memcached"),
                        (9200, "Elasticsearch"),
                    ];

                    let mut is_critical = false;
                    let mut critical_name = "";

                    for (crit_port, name) in critical_ports {
                        if self.port_in_range(port, to_port, *crit_port) {
                            is_critical = true;
                            critical_name = name;
                            break;
                        }
                    }

                    if is_critical {
                        // Critical port open to entire internet
                        self.error_diag(
                            Diagnostic::error_at(
                                format!(
                                    "{} (port {}) open to entire internet",
                                    critical_name,
                                    port.unwrap_or(0)
                                ),
                                rule.span, // Point to the specific rule
                                self.filename,
                            )
                            .with_code(crate::errors::ErrorCode::CriticalPortExposed)
                            .with_primary_label(format!("{} exposed to 0.0.0.0/0", critical_name))
                            .with_help("Restrict to specific IPs, or wrap in unsafe"),
                        );
                    } else {
                        // Non-critical port open to entire internet
                        self.error_diag(
                            Diagnostic::error_at(
                                format!(
                                    "port {} open to entire internet (0.0.0.0/0)",
                                    port.unwrap_or(0)
                                ),
                                rule.span, // Point to the specific rule
                                self.filename,
                            )
                            .with_code(crate::errors::ErrorCode::PortExposed)
                            .with_primary_label("exposes port to 0.0.0.0/0")
                            .with_help("Wrap in unsafe with justification"),
                        );
                    }
                }
            }
        }
    }

    /// Check if a CIDR is entirely within a private network (RFC 1918)
    ///
    /// Uses the `ipnetwork` crate for robust CIDR parsing and validation.
    /// Returns true only if the ENTIRE CIDR range falls within:
    /// - 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16 (RFC 1918)
    /// - 127.0.0.0/8 (localhost)
    ///
    /// Note: Currently not used since we only flag 0.0.0.0/0 as unsafe,
    /// but keeping for potential future stricter modes.
    #[allow(dead_code)]
    fn is_private_cidr(&self, cidr: &str) -> bool {
        use ipnetwork::Ipv4Network;
        use std::net::Ipv4Addr;

        // Parse the input CIDR
        let network: Ipv4Network = match cidr.parse() {
            Ok(n) => n,
            Err(_) => {
                // Try parsing as just an IP (assume /32)
                match cidr.parse::<Ipv4Addr>() {
                    Ok(ip) => Ipv4Network::new(ip, 32).unwrap(),
                    Err(_) => return false,
                }
            }
        };

        // Define private network ranges
        let private_ranges: [Ipv4Network; 4] = [
            "10.0.0.0/8".parse().unwrap(),
            "172.16.0.0/12".parse().unwrap(),
            "192.168.0.0/16".parse().unwrap(),
            "127.0.0.0/8".parse().unwrap(),
        ];

        // Check if the input CIDR is entirely contained within any private range
        for private in &private_ranges {
            if private.contains(network.network()) && private.contains(network.broadcast()) {
                return true;
            }
        }

        false
    }

    fn get_record_number_field(&self, fields: &[TypedRecordField], name: &str) -> Option<i64> {
        fields.iter().find(|f| f.name == name).and_then(|f| {
            if let TypedExprKind::Literal(crate::ast::Literal::Number(n)) = &f.value.kind {
                Some(*n as i64)
            } else {
                None
            }
        })
    }

    fn get_record_string_field(&self, fields: &[TypedRecordField], name: &str) -> Option<String> {
        fields.iter().find(|f| f.name == name).and_then(|f| {
            if let TypedExprKind::Literal(crate::ast::Literal::String(s)) = &f.value.kind {
                Some(s.clone())
            } else {
                None
            }
        })
    }

    fn port_in_range(&self, port: Option<i64>, to_port: Option<i64>, target: i64) -> bool {
        match (port, to_port) {
            (Some(p), Some(t)) => p <= target && target <= t,
            (Some(p), None) => p == target,
            _ => false,
        }
    }

    /// Type compatibility check for function call arguments
    /// More relaxed than general type compatibility - allows lambdas to match function parameters
    fn types_compatible_for_call(&self, actual: &ResolvedType, expected: &ResolvedType) -> bool {
        match (actual, expected) {
            // Exact match
            (a, b) if a == b => true,

            // Unknown is compatible with anything (for error recovery and inference)
            (ResolvedType::Unknown, _) | (_, ResolvedType::Unknown) => true,

            // Union type: actual is compatible if it matches either side
            (actual, ResolvedType::Union(left, right)) => {
                self.types_compatible_for_call(actual, left)
                    || self.types_compatible_for_call(actual, right)
            }

            // Record type: any record is compatible with empty record (generic record param)
            (ResolvedType::Record(_), ResolvedType::Record(expected_fields))
                if expected_fields.is_empty() =>
            {
                true
            }

            // Tags type: record with all string values (AWS compatible)
            (ResolvedType::Record(fields), ResolvedType::Tags) => {
                fields.values().all(|v| matches!(v, ResolvedType::String))
            }

            // Function types: if expected is a function and actual is a function,
            // be lenient (allow lambda inference)
            (ResolvedType::Function { .. }, ResolvedType::Function { .. }) => true,

            // List type compatibility
            (ResolvedType::List(a), ResolvedType::List(b)) => self.types_compatible_for_call(a, b),

            // Fall back to regular compatibility
            _ => self.types_compatible(actual, expected),
        }
    }

    fn init_stdlib(&mut self) {
        self.symbols
            .insert("S3".to_string(), ResolvedType::Module("S3".to_string()));
        self.symbols.insert(
            "Network".to_string(),
            ResolvedType::Module("Network".to_string()),
        );
        self.symbols.insert(
            "CloudWatch".to_string(),
            ResolvedType::Module("CloudWatch".to_string()),
        );
        self.symbols
            .insert("IAM".to_string(), ResolvedType::Module("IAM".to_string()));
        self.symbols
            .insert("EC2".to_string(), ResolvedType::Module("EC2".to_string()));
        self.symbols
            .insert("RDS".to_string(), ResolvedType::Module("RDS".to_string()));
    }

    #[allow(clippy::only_used_in_recursion)]
    fn types_compatible(&self, actual: &ResolvedType, expected: &ResolvedType) -> bool {
        match (actual, expected) {
            // Exact match
            (a, b) if a == b => true,

            // Unknown is compatible with anything (for error recovery)
            (ResolvedType::Unknown, _) | (_, ResolvedType::Unknown) => true,

            // List type compatibility
            (ResolvedType::List(a), ResolvedType::List(b)) => self.types_compatible(a, b),

            // Record type compatibility
            (ResolvedType::Record(_), ResolvedType::Record(b_fields)) if b_fields.is_empty() => {
                // Empty record means "any record" - used for generic params
                true
            }
            (ResolvedType::Record(a_fields), ResolvedType::Record(b_fields)) => {
                // Structural compatibility: all fields in b must exist in a with compatible types
                b_fields.iter().all(|(k, v)| {
                    a_fields
                        .get(k)
                        .is_some_and(|av| self.types_compatible(av, v))
                })
            }

            // Tags type: record with all string values (AWS compatible)
            (ResolvedType::Record(fields), ResolvedType::Tags) => {
                fields.values().all(|v| matches!(v, ResolvedType::String))
            }

            // Everything else is incompatible
            _ => false,
        }
    }

    #[allow(clippy::only_used_in_recursion)]
    fn type_to_string(&self, ty: &ResolvedType) -> String {
        match ty {
            ResolvedType::String => "String".to_string(),
            ResolvedType::Number => "Number".to_string(),
            ResolvedType::Bool => "Bool".to_string(),
            ResolvedType::SecureBucket => "SecureBucket".to_string(),
            ResolvedType::Bucket => "Bucket".to_string(),
            ResolvedType::Vpc => "Vpc".to_string(),
            ResolvedType::Subnet => "Subnet".to_string(),
            ResolvedType::SecurityGroup => "SecurityGroup".to_string(),
            ResolvedType::InternetGateway => "InternetGateway".to_string(),
            ResolvedType::IamRole => "IamRole".to_string(),
            ResolvedType::IamPolicy => "IamPolicy".to_string(),
            ResolvedType::LogGroup => "LogGroup".to_string(),
            ResolvedType::Database => "Database".to_string(),
            ResolvedType::List(inner) => format!("List<{}>", self.type_to_string(inner)),
            ResolvedType::Record(_) => "Record".to_string(),
            ResolvedType::Tags => "Tags".to_string(),
            ResolvedType::Union(a, b) => {
                format!("{} | {}", self.type_to_string(a), self.type_to_string(b))
            }
            ResolvedType::Function { params, returns } => {
                let param_strs: Vec<_> = params
                    .iter()
                    .map(|(name, ty)| format!("{}: {}", name, self.type_to_string(ty)))
                    .collect();
                format!(
                    "({}) -> {}",
                    param_strs.join(", "),
                    self.type_to_string(returns)
                )
            }
            ResolvedType::Unverified(inner) => {
                format!("Unverified<{}>", self.type_to_string(inner))
            }
            ResolvedType::Module(name) => format!("Module({})", name),
            ResolvedType::Void => "Void".to_string(),
            ResolvedType::Unknown => "Unknown".to_string(),
        }
    }

    fn error_diag(&mut self, diag: Diagnostic) {
        self.errors.push(diag);
    }

    /// Get all symbol names for "did you mean" suggestions
    fn get_symbol_names(&self) -> Vec<&str> {
        self.symbols.keys().map(|s| s.as_str()).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::lexer::tokenize;
    use crate::parser::parse;

    fn check_str(source: &str) -> Result<TypedProgram, Vec<Diagnostic>> {
        let tokens = tokenize(source, "test.hk").unwrap();
        let ast = parse(tokens, source, "test.hk").unwrap();
        check(&ast, source, "test.hk", false).map(|r| r.program)
    }

    #[test]
    fn test_literal_types() {
        let program = check_str(r#"val x = "hello""#).unwrap();
        match &program.statements[0] {
            TypedStatement::ValDecl { resolved_type, .. } => {
                assert_eq!(*resolved_type, ResolvedType::String);
            }
            _ => panic!("expected ValDecl"),
        }
    }

    #[test]
    fn test_import_is_tainted() {
        let program = check_str(r#"import "legacy.tf" as legacy"#).unwrap();
        assert_eq!(program.statements.len(), 1);
    }

    #[test]
    fn test_tainted_value_blocked() {
        let source = r#"
            import "legacy.tf" as legacy
            val x = legacy.bucket
        "#;
        let result = check_str(source);
        assert!(
            result.is_err(),
            "should fail: accessing tainted value outside unsafe"
        );
    }

    #[test]
    fn test_unsafe_allows_tainted() {
        let source = r#"
            import "legacy.tf" as legacy
            val x = unsafe("migration") { legacy.bucket }
        "#;
        let result = check_str(source);
        assert!(result.is_ok(), "should succeed: unsafe unwraps taint");
    }

    #[test]
    fn test_unsafe_unwraps_taint() {
        let source = r#"
            import "legacy.tf" as legacy
            val bucket = unsafe("verified") { legacy.bucket }
        "#;
        let program = check_str(source).unwrap();
        match &program.statements[1] {
            TypedStatement::ValDecl { resolved_type, .. } => {
                assert!(!resolved_type.is_tainted(), "unsafe should unwrap taint");
            }
            _ => panic!("expected ValDecl"),
        }
    }

    #[test]
    fn test_binary_arithmetic() {
        let program = check_str("val x = 1 + 2").unwrap();
        match &program.statements[0] {
            TypedStatement::ValDecl { resolved_type, .. } => {
                assert_eq!(*resolved_type, ResolvedType::Number);
            }
            _ => panic!("expected ValDecl"),
        }
    }

    #[test]
    fn test_binary_comparison() {
        let program = check_str("val x = 1 < 2").unwrap();
        match &program.statements[0] {
            TypedStatement::ValDecl { resolved_type, .. } => {
                assert_eq!(*resolved_type, ResolvedType::Bool);
            }
            _ => panic!("expected ValDecl"),
        }
    }

    #[test]
    fn test_binary_logical() {
        let program = check_str("val x = true && false").unwrap();
        match &program.statements[0] {
            TypedStatement::ValDecl { resolved_type, .. } => {
                assert_eq!(*resolved_type, ResolvedType::Bool);
            }
            _ => panic!("expected ValDecl"),
        }
    }

    #[test]
    fn test_unary_not() {
        let program = check_str("val x = !true").unwrap();
        match &program.statements[0] {
            TypedStatement::ValDecl { resolved_type, .. } => {
                assert_eq!(*resolved_type, ResolvedType::Bool);
            }
            _ => panic!("expected ValDecl"),
        }
    }

    #[test]
    fn test_unary_neg() {
        let program = check_str("val x = -5").unwrap();
        match &program.statements[0] {
            TypedStatement::ValDecl { resolved_type, .. } => {
                assert_eq!(*resolved_type, ResolvedType::Number);
            }
            _ => panic!("expected ValDecl"),
        }
    }

    #[test]
    fn test_lambda_type() {
        let program = check_str("val f = x => 42").unwrap();
        match &program.statements[0] {
            TypedStatement::ValDecl { resolved_type, .. } => match resolved_type {
                ResolvedType::Function { returns, .. } => {
                    assert_eq!(**returns, ResolvedType::Number);
                }
                _ => panic!("expected Function type"),
            },
            _ => panic!("expected ValDecl"),
        }
    }
    #[test]
    fn test_shadow_ban() {
        let source = r#"hcl("bad") { resource "aws_s3_bucket" "b" {} }"#;
        let result = check_str(source);
        assert!(result.is_err(), "should fail: shadow banned resource");
        let errs = result.err().unwrap();
        assert!(errs[0].message.contains("Shadow Ban"));
    }

    #[test]
    fn test_unsafe_hatch_for_shadow_ban() {
        let source = r#"unsafe("bypass") { hcl("ok") { resource "aws_s3_bucket" "b" {} } }"#;
        let result = check_str(source);
        assert!(result.is_ok(), "should succeed: unsafe bypasses shadow ban");
    }

    #[test]
    fn test_strict_mode() {
        let source = r#"hcl("any") { resource "foo" "bar" {} }"#;
        let tokens = tokenize(source, "test.hk").unwrap();
        let ast = parse(tokens, source, "test.hk").unwrap();
        // Manually call check with strict=true
        let result = check(&ast, source, "test.hk", true);

        assert!(result.is_err(), "should fail: strict mode forbids HCL");
        let errs = result.err().unwrap();
        assert!(errs[0].message.contains("forbidden in strict mode"));
    }

    #[test]
    fn test_strict_mode_overrides_unsafe() {
        let source = r#"unsafe("bypass") { hcl("ok") { resource "foo" "bar" {} } }"#;
        let tokens = tokenize(source, "test.hk").unwrap();
        let ast = parse(tokens, source, "test.hk").unwrap();
        // Manually call check with strict=true
        let result = check(&ast, source, "test.hk", true);

        assert!(result.is_err(), "should fail: strict mode overrides unsafe");
        let errs = result.err().unwrap();
        assert!(errs[0].message.contains("forbidden in strict mode"));
    }
}