pmcp-code-mode 0.5.0

//! JavaScript-specific validation for Code Mode (OpenAPI servers).
//!
//! This module validates JavaScript code generated by LLMs for REST API interactions.
//! It uses SWC (Speedy Web Compiler) for production-grade parsing and enforces a safe
//! subset of JavaScript that prevents malicious operations while enabling powerful
//! API orchestration.
//!
//! ## Safe Subset
//!
//! Allowed:
//! - async/await for API calls
//! - api.get(), api.post(), api.put(), api.delete(), api.patch() calls
//! - const/let variable declarations
//! - Arrow functions for callbacks (including block bodies with nested callbacks)
//! - Array methods: map, filter, reduce, find, some, every, slice
//! - Object destructuring and spread
//! - Template literals (string interpolation)
//! - Bounded for...of loops (with .slice() limits)
//! - if/else conditionals
//! - try/catch for error handling
//! - Logical operators (&&, ||)
//!
//! Blocked:
//! - import/export statements
//! - eval(), Function(), new Function()
//! - while/do-while loops (unbounded)
//! - Regular function declarations (only arrow functions)
//! - new keyword (except specific built-ins)
//! - this keyword
//! - class declarations
//! - Generators/iterators
//! - with statement
//! - delete operator
//! - Prototype manipulation

use crate::types::{
    CodeLocation, CodeType, Complexity, SecurityAnalysis, SecurityIssue, SecurityIssueType,
    ValidationError,
};
use std::collections::HashSet;
use swc_common::{sync::Lrc, SourceMap, Span};
use swc_ecma_ast::*;
use swc_ecma_parser::{lexer::Lexer, Parser, StringInput, Syntax};
use swc_ecma_visit::{Visit, VisitWith};

/// HTTP methods that can be called via the api object.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HttpMethod {
    Get,
    Post,
    Put,
    Delete,
    Patch,
    Head,
    Options,
}

impl HttpMethod {
    /// Whether this method is read-only (safe).
    pub fn is_read_only(&self) -> bool {
        matches!(
            self,
            HttpMethod::Get | HttpMethod::Head | HttpMethod::Options
        )
    }

    /// Parse from string.
    pub fn from_str(s: &str) -> Option<Self> {
        match s.to_lowercase().as_str() {
            "get" => Some(HttpMethod::Get),
            "post" => Some(HttpMethod::Post),
            "put" => Some(HttpMethod::Put),
            "delete" => Some(HttpMethod::Delete),
            "patch" => Some(HttpMethod::Patch),
            "head" => Some(HttpMethod::Head),
            "options" => Some(HttpMethod::Options),
            _ => None,
        }
    }
}

/// An API call extracted from the JavaScript code.
#[derive(Debug, Clone)]
pub struct ApiCall {
    /// The HTTP method
    pub method: HttpMethod,
    /// The path template (may contain interpolations)
    pub path: String,
    /// Whether the path is dynamic (contains template expressions)
    pub is_dynamic_path: bool,
    /// Line number in the source
    pub line: u32,
    /// Column number in the source
    pub column: u32,
}

/// Declared output type from @returns annotation.
#[derive(Debug, Clone, Default)]
pub struct OutputDeclaration {
    /// Whether a @returns annotation was found
    pub has_declaration: bool,

    /// The raw type string from the annotation (e.g., "{ users: Array<{ id: string, name: string }> }")
    pub type_string: Option<String>,

    /// Fields mentioned in the output type (extracted for field blocklist checking)
    pub declared_fields: HashSet<String>,

    /// Whether the declaration uses spread operators (potential field leakage)
    pub has_spread_risk: bool,
}

/// Information extracted from parsed JavaScript code.
#[derive(Debug, Clone, Default)]
pub struct JavaScriptCodeInfo {
    /// All API calls in the code
    pub api_calls: Vec<ApiCall>,

    /// Whether the code is read-only (only GET/HEAD/OPTIONS calls)
    pub is_read_only: bool,

    /// All endpoints accessed
    pub endpoints_accessed: HashSet<String>,

    /// All HTTP methods used
    pub methods_used: HashSet<String>,

    /// Whether the code uses async/await
    pub uses_async: bool,

    /// Variable names declared
    pub variable_names: Vec<String>,

    /// Maximum nesting depth
    pub max_depth: usize,

    /// Number of for...of loops
    pub loop_count: usize,

    /// Whether all loops are bounded (use .slice())
    pub all_loops_bounded: bool,

    /// Policy violations found during parsing
    pub violations: Vec<SafetyViolation>,

    /// Total number of statements
    pub statement_count: usize,

    /// Output declaration from @returns annotation
    pub output_declaration: OutputDeclaration,

    /// Whether the script contains spread operators that could leak fields
    pub has_output_spread_risk: bool,
}

/// A safety violation found during JavaScript validation.
#[derive(Debug, Clone)]
pub struct SafetyViolation {
    /// Type of violation
    pub violation_type: SafetyViolationType,
    /// Human-readable message
    pub message: String,
    /// Location in source
    pub location: Option<CodeLocation>,
}

/// Types of safety violations in JavaScript code.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SafetyViolationType {
    /// import/export statement
    ImportExport,
    /// eval() or Function() call
    DynamicCodeExecution,
    /// while/do-while loop (unbounded)
    UnboundedLoop,
    /// Regular function declaration
    FunctionDeclaration,
    /// try/catch statement
    TryCatch,
    /// new keyword (except allowed)
    NewKeyword,
    /// this keyword
    ThisKeyword,
    /// class declaration
    ClassDeclaration,
    /// Generator function
    Generator,
    /// with statement
    WithStatement,
    /// delete operator
    DeleteOperator,
    /// Prototype manipulation
    PrototypeManipulation,
    /// Unbounded for loop (no .slice())
    UnboundedForLoop,
    /// Unknown API call (not api.method())
    UnknownApiCall,
}

/// JavaScript code validator for OpenAPI Code Mode.
pub struct JavaScriptValidator {
    /// Sensitive path patterns (e.g., "/admin", "/internal")
    sensitive_paths: Vec<String>,

    /// Maximum allowed nesting depth
    max_depth: usize,

    /// Maximum allowed API calls
    max_api_calls: usize,

    /// Maximum allowed loop count
    max_loops: usize,

    /// Maximum allowed statements
    max_statements: usize,

    /// Allowed SDK operation names (camelCase). When non-empty, SDK mode is active:
    /// known operations are accepted, unknown ones generate UnknownApiCall violations.
    sdk_operations: HashSet<String>,
}

impl Default for JavaScriptValidator {
    fn default() -> Self {
        Self {
            sensitive_paths: vec![
                "/admin".into(),
                "/internal".into(),
                "/debug".into(),
                "/metrics".into(),
                "/health".into(),
            ],
            max_depth: 10,
            max_api_calls: 50,
            max_loops: 10,
            max_statements: 100,
            sdk_operations: HashSet::new(),
        }
    }
}

/// Check if a word is a TypeScript/JSDoc type keyword (not a field name).
fn is_type_keyword(word: &str) -> bool {
    matches!(
        word,
        "string"
            | "number"
            | "boolean"
            | "null"
            | "undefined"
            | "void"
            | "any"
            | "never"
            | "object"
            | "Array"
            | "Promise"
            | "Record"
            | "Map"
            | "Set"
            | "Date"
            | "type"
            | "interface"
    )
}

impl JavaScriptValidator {
    /// Create a new validator with custom settings.
    pub fn new(
        sensitive_paths: Vec<String>,
        max_depth: usize,
        max_api_calls: usize,
        max_loops: usize,
        max_statements: usize,
    ) -> Self {
        Self {
            sensitive_paths,
            max_depth,
            max_api_calls,
            max_loops,
            max_statements,
            sdk_operations: HashSet::new(),
        }
    }

    /// Set the allowed SDK operation names. When non-empty, the validator operates in SDK mode:
    /// known SDK operations are accepted, unknown ones generate `UnknownApiCall` violations.
    pub fn with_sdk_operations(mut self, operations: HashSet<String>) -> Self {
        self.sdk_operations = operations;
        self
    }

    /// Parse @returns annotation from code comments.
    ///
    /// Supports formats:
    /// - `// @returns { type }` (single-line comment)
    /// - `/// @returns { type }` (triple-slash comment)
    /// - `/** @returns { type } */` (JSDoc comment)
    fn parse_returns_annotation(code: &str) -> OutputDeclaration {
        let mut declaration = OutputDeclaration::default();

        // Try to find @returns annotation in comments
        for line in code.lines() {
            let trimmed = line.trim();

            // Check for triple-slash comment: /// @returns { ... }
            // Must check before double-slash since /// starts with //
            if let Some(rest) = trimmed.strip_prefix("///") {
                if let Some(returns_content) = Self::extract_returns_content(rest) {
                    declaration.has_declaration = true;
                    declaration.type_string = Some(returns_content.clone());
                    declaration.declared_fields = Self::extract_fields_from_type(&returns_content);
                    declaration.has_spread_risk = returns_content.contains("...");
                    return declaration;
                }
            }
            // Check for double-slash comment: // @returns { ... }
            else if let Some(rest) = trimmed.strip_prefix("//") {
                if let Some(returns_content) = Self::extract_returns_content(rest) {
                    declaration.has_declaration = true;
                    declaration.type_string = Some(returns_content.clone());
                    declaration.declared_fields = Self::extract_fields_from_type(&returns_content);
                    declaration.has_spread_risk = returns_content.contains("...");
                    return declaration;
                }
            }

            // Check for JSDoc comment start: /** @returns { ... } */ or /** ... @returns ... */
            if trimmed.starts_with("/**") || trimmed.starts_with("*") {
                let content = trimmed
                    .trim_start_matches("/**")
                    .trim_start_matches('*')
                    .trim_end_matches("*/")
                    .trim();

                if let Some(returns_content) = Self::extract_returns_content(content) {
                    declaration.has_declaration = true;
                    declaration.type_string = Some(returns_content.clone());
                    declaration.declared_fields = Self::extract_fields_from_type(&returns_content);
                    declaration.has_spread_risk = returns_content.contains("...");
                    return declaration;
                }
            }
        }

        declaration
    }

    /// Extract the content after @returns.
    fn extract_returns_content(text: &str) -> Option<String> {
        let text = text.trim();

        // Look for @returns or @return
        let returns_pos = text.find("@returns").or_else(|| text.find("@return"))?;

        // Extract everything after the tag
        let after_tag = &text[returns_pos..];
        let content_start = after_tag.find(|c: char| c == '{' || c == '(')?;

        // Find the matching closing bracket
        let chars: Vec<char> = after_tag[content_start..].chars().collect();
        let open_char = chars[0];
        let close_char = if open_char == '{' { '}' } else { ')' };

        let mut depth = 0;
        let mut end_pos = 0;

        for (i, c) in chars.iter().enumerate() {
            if *c == open_char {
                depth += 1;
            } else if *c == close_char {
                depth -= 1;
                if depth == 0 {
                    end_pos = i + 1;
                    break;
                }
            }
        }

        if end_pos > 0 {
            Some(after_tag[content_start..content_start + end_pos].to_string())
        } else {
            // If no closing bracket found, take the rest of the line
            Some(after_tag[content_start..].trim().to_string())
        }
    }

    /// Extract field names from a type declaration string.
    ///
    /// This is a simple parser that extracts identifiers that appear
    /// before colons, which are typically field names in object types.
    fn extract_fields_from_type(type_string: &str) -> HashSet<String> {
        let mut fields = HashSet::new();

        // Simple regex-free parser: find "fieldName:" patterns
        let chars: Vec<char> = type_string.chars().collect();
        let mut current_word = String::new();
        let mut in_word = false;

        for (_i, c) in chars.iter().enumerate() {
            if c.is_alphanumeric() || *c == '_' {
                current_word.push(*c);
                in_word = true;
            } else {
                if in_word && *c == ':' {
                    // This is a field name
                    if !current_word.is_empty()
                        && !is_type_keyword(&current_word)
                        && !current_word.chars().next().unwrap().is_ascii_uppercase()
                    {
                        fields.insert(current_word.clone());
                    }
                }
                current_word.clear();
                in_word = false;
            }
        }

        fields
    }

    /// Check if declared output fields contain any blocked fields.
    pub fn check_output_against_blocklist(
        declaration: &OutputDeclaration,
        blocked_fields: &HashSet<String>,
    ) -> Vec<String> {
        let mut violations = Vec::new();

        for field in &declaration.declared_fields {
            // Check exact match
            if blocked_fields.contains(field) {
                violations.push(format!("Output declares blocked field: {}", field));
                continue;
            }

            // Check wildcard patterns like *.fieldName
            for blocked in blocked_fields {
                if blocked.starts_with("*.") {
                    let pattern = &blocked[2..];
                    if field == pattern {
                        violations.push(format!(
                            "Output declares blocked field pattern: {}",
                            blocked
                        ));
                    }
                }
            }
        }

        violations
    }

    /// Parse and validate JavaScript code.
    pub fn validate(&self, code: &str) -> Result<JavaScriptCodeInfo, ValidationError> {
        // Parse the code using SWC
        let cm: Lrc<SourceMap> = Default::default();

        // Create a source file from the code - BytesStr accepts String
        let fm = cm.new_source_file(
            swc_common::FileName::Custom("code.js".into()).into(),
            code.to_string(),
        );

        let lexer = Lexer::new(
            Syntax::Es(Default::default()),
            EsVersion::Es2022,
            StringInput::from(&*fm),
            None,
        );

        let mut parser = Parser::new_from(lexer);

        let module = parser
            .parse_module()
            .map_err(|e| ValidationError::ParseError {
                message: format!("JavaScript parse error: {:?}", e.into_kind()),
                line: 0,
                column: 0,
            })?;

        // Visit the AST to extract information and check safety
        let mut visitor = SafetyVisitor::new(&cm).with_sdk_operations(self.sdk_operations.clone());
        module.visit_with(&mut visitor);

        let mut info = visitor.into_info();

        // Parse @returns annotation from code comments
        info.output_declaration = Self::parse_returns_annotation(code);

        // Validate constraints
        if info.api_calls.len() > self.max_api_calls {
            return Err(ValidationError::SecurityError {
                message: format!(
                    "Too many API calls: {} (max: {})",
                    info.api_calls.len(),
                    self.max_api_calls
                ),
                issue: SecurityIssueType::HighComplexity,
            });
        }

        if info.max_depth > self.max_depth {
            return Err(ValidationError::SecurityError {
                message: format!(
                    "Code nesting depth {} exceeds maximum {}",
                    info.max_depth, self.max_depth
                ),
                issue: SecurityIssueType::DeepNesting,
            });
        }

        if info.loop_count > self.max_loops {
            return Err(ValidationError::SecurityError {
                message: format!(
                    "Too many loops: {} (max: {})",
                    info.loop_count, self.max_loops
                ),
                issue: SecurityIssueType::HighComplexity,
            });
        }

        if info.statement_count > self.max_statements {
            return Err(ValidationError::SecurityError {
                message: format!(
                    "Too many statements: {} (max: {})",
                    info.statement_count, self.max_statements
                ),
                issue: SecurityIssueType::HighComplexity,
            });
        }

        // Check for safety violations
        if !info.violations.is_empty() {
            let first = &info.violations[0];
            return Err(ValidationError::SecurityError {
                message: first.message.clone(),
                issue: violation_to_security_issue(first.violation_type),
            });
        }

        // Determine if read-only based on API calls
        info.is_read_only = info.api_calls.iter().all(|c| c.method.is_read_only());

        Ok(info)
    }

    /// Perform security analysis on code info.
    pub fn analyze_security(&self, info: &JavaScriptCodeInfo) -> SecurityAnalysis {
        let mut analysis = SecurityAnalysis {
            is_read_only: info.is_read_only,
            tables_accessed: info.endpoints_accessed.clone(),
            fields_accessed: HashSet::new(),
            has_aggregation: false,
            has_subqueries: info.max_depth > 3,
            estimated_complexity: self.estimate_complexity(info),
            potential_issues: Vec::new(),
            estimated_rows: None,
        };

        // Check for sensitive endpoints
        for endpoint in &info.endpoints_accessed {
            let endpoint_lower = endpoint.to_lowercase();
            if self
                .sensitive_paths
                .iter()
                .any(|s| endpoint_lower.contains(&s.to_lowercase()))
            {
                analysis.potential_issues.push(SecurityIssue::new(
                    SecurityIssueType::SensitiveFields,
                    format!("Code accesses potentially sensitive endpoint: {}", endpoint),
                ));
            }
        }

        // Check for dynamic paths (potential injection)
        for call in &info.api_calls {
            if call.is_dynamic_path {
                analysis.potential_issues.push(
                    SecurityIssue::new(
                        SecurityIssueType::DynamicTableName,
                        format!(
                            "API call at line {} uses dynamic path interpolation",
                            call.line
                        ),
                    )
                    .with_location(CodeLocation {
                        line: call.line,
                        column: call.column,
                    }),
                );
            }
        }

        // Check for deep nesting
        if info.max_depth > 5 {
            analysis.potential_issues.push(SecurityIssue::new(
                SecurityIssueType::DeepNesting,
                format!("Code has deep nesting (depth: {})", info.max_depth),
            ));
        }

        // Check for unbounded loops
        if !info.all_loops_bounded && info.loop_count > 0 {
            analysis.potential_issues.push(SecurityIssue::new(
                SecurityIssueType::UnboundedQuery,
                "Code contains for...of loops without .slice() bounds",
            ));
        }

        // Check for high complexity
        if matches!(analysis.estimated_complexity, Complexity::High) {
            analysis.potential_issues.push(SecurityIssue::new(
                SecurityIssueType::HighComplexity,
                "Code has high complexity",
            ));
        }

        analysis
    }

    /// Estimate code complexity.
    fn estimate_complexity(&self, info: &JavaScriptCodeInfo) -> Complexity {
        let api_count = info.api_calls.len();
        let loop_count = info.loop_count;
        let depth = info.max_depth;
        let statement_count = info.statement_count;

        // Simple heuristic
        let complexity_score = api_count * 3 + loop_count * 5 + depth * 2 + statement_count;

        if complexity_score > 100 {
            Complexity::High
        } else if complexity_score > 50 {
            Complexity::Medium
        } else {
            Complexity::Low
        }
    }

    /// Convert code info to CodeType.
    pub fn to_code_type(&self, info: &JavaScriptCodeInfo) -> CodeType {
        if info.is_read_only {
            CodeType::RestGet
        } else {
            CodeType::RestMutation
        }
    }
}

/// AST visitor for extracting info and checking safety.
struct SafetyVisitor {
    source_map: Lrc<SourceMap>,
    api_calls: Vec<ApiCall>,
    violations: Vec<SafetyViolation>,
    variable_names: Vec<String>,
    endpoints_accessed: HashSet<String>,
    methods_used: HashSet<String>,
    uses_async: bool,
    current_depth: usize,
    max_depth: usize,
    loop_count: usize,
    bounded_loops: usize,
    statement_count: usize,
    /// Whether the code uses spread operators in return values (potential field leakage)
    has_spread_in_return: bool,
    /// Whether we're currently inside a return statement
    in_return_context: bool,
    /// Allowed SDK operation names (camelCase). When non-empty, SDK mode validation is active.
    sdk_operations: HashSet<String>,
}

impl SafetyVisitor {
    fn new(source_map: &Lrc<SourceMap>) -> Self {
        Self {
            source_map: source_map.clone(),
            api_calls: Vec::new(),
            violations: Vec::new(),
            variable_names: Vec::new(),
            endpoints_accessed: HashSet::new(),
            methods_used: HashSet::new(),
            uses_async: false,
            current_depth: 0,
            max_depth: 0,
            loop_count: 0,
            bounded_loops: 0,
            statement_count: 0,
            has_spread_in_return: false,
            in_return_context: false,
            sdk_operations: HashSet::new(),
        }
    }

    fn with_sdk_operations(mut self, operations: HashSet<String>) -> Self {
        self.sdk_operations = operations;
        self
    }

    fn into_info(self) -> JavaScriptCodeInfo {
        JavaScriptCodeInfo {
            api_calls: self.api_calls,
            is_read_only: false, // Set later based on API calls
            endpoints_accessed: self.endpoints_accessed,
            methods_used: self.methods_used,
            uses_async: self.uses_async,
            variable_names: self.variable_names,
            max_depth: self.max_depth,
            loop_count: self.loop_count,
            all_loops_bounded: self.loop_count == 0 || self.bounded_loops == self.loop_count,
            violations: self.violations,
            statement_count: self.statement_count,
            output_declaration: OutputDeclaration::default(), // Set later by validator
            has_output_spread_risk: self.has_spread_in_return,
        }
    }

    fn span_to_location(&self, span: Span) -> CodeLocation {
        let loc = self.source_map.lookup_char_pos(span.lo);
        CodeLocation {
            line: loc.line as u32,
            column: loc.col_display as u32,
        }
    }

    fn add_violation(&mut self, violation_type: SafetyViolationType, message: &str, span: Span) {
        self.violations.push(SafetyViolation {
            violation_type,
            message: message.into(),
            location: Some(self.span_to_location(span)),
        });
    }

    fn check_api_call(&mut self, call: &CallExpr) {
        // Check for api.method() pattern
        if let Callee::Expr(expr) = &call.callee {
            if let Expr::Member(member) = &**expr {
                if let Expr::Ident(obj) = &*member.obj {
                    if obj.sym.as_ref() == "api" {
                        if let MemberProp::Ident(method_ident) = &member.prop {
                            let method_name = method_ident.sym.as_ref();

                            if !self.sdk_operations.is_empty() {
                                // SDK mode: validate against allowed operation names
                                if self.sdk_operations.contains(method_name) {
                                    self.methods_used.insert(method_name.to_string());
                                    self.endpoints_accessed
                                        .insert(format!("sdk:{}", method_name));
                                    // No path extraction needed for SDK calls
                                } else {
                                    self.add_violation(
                                        SafetyViolationType::UnknownApiCall,
                                        &format!(
                                            "Unknown SDK operation: api.{}(). Check the code mode schema resource for available operations.",
                                            method_name
                                        ),
                                        call.span,
                                    );
                                }
                                return;
                            }

                            // HTTP mode: validate against known HTTP methods
                            if let Some(method) = HttpMethod::from_str(method_name) {
                                self.methods_used.insert(method_name.to_uppercase());

                                // Extract path from first argument
                                let (path, is_dynamic) = if let Some(arg) = call.args.first() {
                                    self.extract_path(&arg.expr)
                                } else {
                                    ("unknown".into(), false)
                                };

                                self.endpoints_accessed.insert(path.clone());

                                let loc = self.span_to_location(call.span);
                                self.api_calls.push(ApiCall {
                                    method,
                                    path,
                                    is_dynamic_path: is_dynamic,
                                    line: loc.line,
                                    column: loc.column,
                                });
                            } else {
                                self.add_violation(
                                    SafetyViolationType::UnknownApiCall,
                                    &format!("Unknown api method: api.{}()", method_name),
                                    call.span,
                                );
                            }
                        }
                    }
                }
            }
        }
    }

    fn extract_path(&self, expr: &Expr) -> (String, bool) {
        match expr {
            Expr::Lit(Lit::Str(s)) => {
                // Convert Wtf8Atom to String using to_string_lossy (handles WTF-8 encoding)
                (s.value.to_string_lossy().into_owned(), false)
            },
            Expr::Tpl(tpl) => {
                // Template literal - extract static parts
                let mut path = String::new();
                for quasi in &tpl.quasis {
                    // quasi.raw is an Atom (UTF-8), not Wtf8Atom
                    path.push_str(&quasi.raw.to_string());
                    if !tpl.exprs.is_empty() {
                        path.push_str("{...}");
                    }
                }
                (path, !tpl.exprs.is_empty())
            },
            _ => ("dynamic".into(), true),
        }
    }

    fn check_for_bounded(&mut self, for_of: &ForOfStmt) -> bool {
        // Check if the iterable uses .slice()
        if let Expr::Call(call) = &*for_of.right {
            if let Callee::Expr(callee) = &call.callee {
                if let Expr::Member(member) = &**callee {
                    if let MemberProp::Ident(ident) = &member.prop {
                        if ident.sym.as_ref() == "slice" {
                            return true;
                        }
                    }
                }
            }
        }
        false
    }
}

impl Visit for SafetyVisitor {
    // Track depth
    fn visit_block_stmt(&mut self, n: &BlockStmt) {
        self.current_depth += 1;
        self.max_depth = self.max_depth.max(self.current_depth);
        n.visit_children_with(self);
        self.current_depth -= 1;
    }

    // Count statements
    fn visit_stmt(&mut self, n: &Stmt) {
        self.statement_count += 1;
        n.visit_children_with(self);
    }

    // Check for import/export
    fn visit_import_decl(&mut self, n: &ImportDecl) {
        self.add_violation(
            SafetyViolationType::ImportExport,
            "import statements are not allowed",
            n.span,
        );
    }

    fn visit_export_decl(&mut self, n: &ExportDecl) {
        self.add_violation(
            SafetyViolationType::ImportExport,
            "export statements are not allowed",
            n.span,
        );
    }

    fn visit_export_default_decl(&mut self, n: &ExportDefaultDecl) {
        self.add_violation(
            SafetyViolationType::ImportExport,
            "export default is not allowed",
            n.span,
        );
    }

    fn visit_export_default_expr(&mut self, n: &ExportDefaultExpr) {
        self.add_violation(
            SafetyViolationType::ImportExport,
            "export default is not allowed",
            n.span,
        );
    }

    // Check for eval/Function
    fn visit_call_expr(&mut self, n: &CallExpr) {
        // Check for eval() or Function()
        if let Callee::Expr(callee) = &n.callee {
            if let Expr::Ident(ident) = &**callee {
                let name = ident.sym.as_ref();
                if name == "eval" || name == "Function" {
                    self.add_violation(
                        SafetyViolationType::DynamicCodeExecution,
                        &format!("{}() is not allowed", name),
                        n.span,
                    );
                }
            }
        }

        // Check for api.method() calls
        self.check_api_call(n);

        n.visit_children_with(self);
    }

    // Check for while loops
    fn visit_while_stmt(&mut self, n: &WhileStmt) {
        self.add_violation(
            SafetyViolationType::UnboundedLoop,
            "while loops are not allowed (use bounded for...of with .slice())",
            n.span,
        );
        n.visit_children_with(self);
    }

    fn visit_do_while_stmt(&mut self, n: &DoWhileStmt) {
        self.add_violation(
            SafetyViolationType::UnboundedLoop,
            "do-while loops are not allowed (use bounded for...of with .slice())",
            n.span,
        );
        n.visit_children_with(self);
    }

    // Check for...of loops for bounds
    fn visit_for_of_stmt(&mut self, n: &ForOfStmt) {
        self.loop_count += 1;
        if self.check_for_bounded(n) {
            self.bounded_loops += 1;
        }
        n.visit_children_with(self);
    }

    // Check for regular for loops (allow but track)
    fn visit_for_stmt(&mut self, n: &ForStmt) {
        self.loop_count += 1;
        // Regular for loops are bounded by definition
        self.bounded_loops += 1;
        n.visit_children_with(self);
    }

    // Check for function declarations (only arrow functions allowed)
    fn visit_fn_decl(&mut self, n: &FnDecl) {
        self.add_violation(
            SafetyViolationType::FunctionDeclaration,
            "function declarations are not allowed (use arrow functions)",
            n.function.span,
        );
        n.visit_children_with(self);
    }

    // Allow try/catch - it's just control flow and doesn't pose a security risk.
    // This enables scripts to gracefully handle API errors.
    fn visit_try_stmt(&mut self, n: &TryStmt) {
        // Visit children to validate the contents of try/catch blocks
        n.visit_children_with(self);
    }

    // Check for new keyword
    fn visit_new_expr(&mut self, n: &NewExpr) {
        // Allow specific constructors like Date, URL, URLSearchParams
        let allowed = if let Expr::Ident(ident) = &*n.callee {
            matches!(
                ident.sym.as_ref(),
                "Date" | "URL" | "URLSearchParams" | "Map" | "Set" | "Array"
            )
        } else {
            false
        };

        if !allowed {
            self.add_violation(
                SafetyViolationType::NewKeyword,
                "new keyword is only allowed for Date, URL, URLSearchParams, Map, Set, Array",
                n.span,
            );
        }
        n.visit_children_with(self);
    }

    // Check for this keyword
    fn visit_this_expr(&mut self, n: &ThisExpr) {
        self.add_violation(
            SafetyViolationType::ThisKeyword,
            "'this' keyword is not allowed",
            n.span,
        );
    }

    // Check for class declarations
    fn visit_class_decl(&mut self, n: &ClassDecl) {
        self.add_violation(
            SafetyViolationType::ClassDeclaration,
            "class declarations are not allowed",
            n.class.span,
        );
        n.visit_children_with(self);
    }

    // Check for with statement
    fn visit_with_stmt(&mut self, n: &WithStmt) {
        self.add_violation(
            SafetyViolationType::WithStatement,
            "'with' statement is not allowed",
            n.span,
        );
        n.visit_children_with(self);
    }

    // Track async
    fn visit_await_expr(&mut self, n: &AwaitExpr) {
        self.uses_async = true;
        n.visit_children_with(self);
    }

    // Track variable declarations
    fn visit_var_decl(&mut self, n: &VarDecl) {
        for decl in &n.decls {
            if let Pat::Ident(ident) = &decl.name {
                self.variable_names.push(ident.id.sym.to_string());
            }
        }
        n.visit_children_with(self);
    }

    // Check for generators
    fn visit_function(&mut self, n: &Function) {
        if n.is_generator {
            self.add_violation(
                SafetyViolationType::Generator,
                "generator functions are not allowed",
                n.span,
            );
        }
        n.visit_children_with(self);
    }

    // Check for delete operator
    fn visit_unary_expr(&mut self, n: &UnaryExpr) {
        if n.op == UnaryOp::Delete {
            self.add_violation(
                SafetyViolationType::DeleteOperator,
                "'delete' operator is not allowed",
                n.span,
            );
        }
        n.visit_children_with(self);
    }

    // Check for prototype manipulation
    fn visit_member_expr(&mut self, n: &MemberExpr) {
        if let MemberProp::Ident(ident) = &n.prop {
            let name = ident.sym.as_ref();
            if name == "__proto__" || name == "prototype" {
                self.add_violation(
                    SafetyViolationType::PrototypeManipulation,
                    "prototype manipulation is not allowed",
                    n.span,
                );
            }
        }
        n.visit_children_with(self);
    }

    // Track return statements to detect spread operators in output
    fn visit_return_stmt(&mut self, n: &ReturnStmt) {
        self.in_return_context = true;
        n.visit_children_with(self);
        self.in_return_context = false;
    }

    // Detect spread operators in return values (potential field leakage)
    fn visit_spread_element(&mut self, n: &SpreadElement) {
        if self.in_return_context {
            self.has_spread_in_return = true;
        }
        n.visit_children_with(self);
    }
}

/// Convert a safety violation type to a security issue type.
fn violation_to_security_issue(violation: SafetyViolationType) -> SecurityIssueType {
    match violation {
        SafetyViolationType::DynamicCodeExecution => SecurityIssueType::PotentialInjection,
        SafetyViolationType::PrototypeManipulation => SecurityIssueType::PotentialInjection,
        SafetyViolationType::UnboundedLoop | SafetyViolationType::UnboundedForLoop => {
            SecurityIssueType::UnboundedQuery
        },
        _ => SecurityIssueType::HighComplexity,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_simple_api_call() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            const response = await api.get("/users");
            return response.data;
        "#;

        let info = validator.validate(code).unwrap();
        assert!(info.is_read_only);
        assert_eq!(info.api_calls.len(), 1);
        assert_eq!(info.api_calls[0].method, HttpMethod::Get);
        assert!(info.endpoints_accessed.contains("/users"));
    }

    #[test]
    fn test_multiple_api_calls() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            const user = await api.get("/users/123");
            const orders = await api.get(`/users/${user.id}/orders`);
            return { user, orders };
        "#;

        let info = validator.validate(code).unwrap();
        assert!(info.is_read_only);
        assert_eq!(info.api_calls.len(), 2);
        assert!(info.api_calls[1].is_dynamic_path);
    }

    #[test]
    fn test_mutation_detection() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            const result = await api.post("/users", { name: "test" });
            return result;
        "#;

        let info = validator.validate(code).unwrap();
        assert!(!info.is_read_only);
        assert_eq!(info.api_calls[0].method, HttpMethod::Post);
    }

    #[test]
    fn test_reject_eval() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            const result = eval("api.get('/users')");
        "#;

        let result = validator.validate(code);
        assert!(result.is_err());
    }

    #[test]
    fn test_reject_while_loop() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            let i = 0;
            while (i < 10) {
                await api.get("/data");
                i++;
            }
        "#;

        let result = validator.validate(code);
        assert!(result.is_err());
    }

    #[test]
    fn test_allow_bounded_for_of() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            const results = [];
            for (const id of userIds.slice(0, 10)) {
                const user = await api.get(`/users/${id}`);
                results.push(user);
            }
            return results;
        "#;

        let info = validator.validate(code).unwrap();
        assert!(info.all_loops_bounded);
        assert_eq!(info.loop_count, 1);
    }

    #[test]
    fn test_reject_import() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            import axios from 'axios';
            const result = await api.get("/users");
        "#;

        let result = validator.validate(code);
        assert!(result.is_err());
    }

    #[test]
    fn test_allow_arrow_functions() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            const users = await api.get("/users");
            const names = users.data.map(u => u.name);
            return names;
        "#;

        let info = validator.validate(code).unwrap();
        assert!(info.violations.is_empty());
    }

    #[test]
    fn test_reject_function_declaration() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            function fetchUser(id) {
                return api.get(`/users/${id}`);
            }
        "#;

        let result = validator.validate(code);
        assert!(result.is_err());
    }

    #[test]
    fn test_security_analysis_sensitive_endpoint() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            const config = await api.get("/admin/config");
            return config;
        "#;

        let info = validator.validate(code).unwrap();
        let analysis = validator.analyze_security(&info);

        assert!(analysis
            .potential_issues
            .iter()
            .any(|i| matches!(i.issue_type, SecurityIssueType::SensitiveFields)));
    }

    #[test]
    fn test_parse_returns_annotation_triple_slash() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            /// @returns { users: Array<{ id: string, name: string }> }
            const users = await api.get("/users");
            return { users: users.map(u => ({ id: u.id, name: u.name })) };
        "#;

        let info = validator.validate(code).unwrap();
        assert!(info.output_declaration.has_declaration);
        assert!(info.output_declaration.declared_fields.contains("id"));
        assert!(info.output_declaration.declared_fields.contains("name"));
        assert!(info.output_declaration.declared_fields.contains("users"));
    }

    #[test]
    fn test_parse_returns_annotation_double_slash() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            // @returns { products: Array<{ id: string, name: string, price: number }> }
            const products = await api.get("/products");
            return { products: products.map(p => ({ id: p.id, name: p.name, price: p.price })) };
        "#;

        let info = validator.validate(code).unwrap();
        assert!(info.output_declaration.has_declaration);
        assert!(info.output_declaration.declared_fields.contains("id"));
        assert!(info.output_declaration.declared_fields.contains("name"));
        assert!(info.output_declaration.declared_fields.contains("price"));
        assert!(info.output_declaration.declared_fields.contains("products"));
    }

    #[test]
    fn test_parse_returns_annotation_jsdoc() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            /** @returns { user: { id: string, email: string } } */
            const user = await api.get("/users/123");
            return { user: { id: user.id, email: user.email } };
        "#;

        let info = validator.validate(code).unwrap();
        assert!(info.output_declaration.has_declaration);
        assert!(info.output_declaration.declared_fields.contains("id"));
        assert!(info.output_declaration.declared_fields.contains("email"));
        assert!(info.output_declaration.declared_fields.contains("user"));
    }

    #[test]
    fn test_no_returns_annotation() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            const users = await api.get("/users");
            return users;
        "#;

        let info = validator.validate(code).unwrap();
        assert!(!info.output_declaration.has_declaration);
        assert!(info.output_declaration.declared_fields.is_empty());
    }

    #[test]
    fn test_spread_operator_detection() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            const user = await api.get("/users/123");
            return { ...user, computed: "value" };
        "#;

        let info = validator.validate(code).unwrap();
        assert!(info.has_output_spread_risk);
    }

    #[test]
    fn test_no_spread_operator_in_return() {
        let validator = JavaScriptValidator::default();
        let code = r#"
            const user = await api.get("/users/123");
            return { id: user.id, name: user.name };
        "#;

        let info = validator.validate(code).unwrap();
        assert!(!info.has_output_spread_risk);
    }

    #[test]
    fn test_check_output_against_blocklist() {
        let declaration = OutputDeclaration {
            has_declaration: true,
            type_string: Some("{ id: string, ssn: string }".to_string()),
            declared_fields: ["id", "ssn"].iter().map(|s| s.to_string()).collect(),
            has_spread_risk: false,
        };

        let blocked_fields: HashSet<String> =
            ["ssn", "password"].iter().map(|s| s.to_string()).collect();

        let violations =
            JavaScriptValidator::check_output_against_blocklist(&declaration, &blocked_fields);
        assert_eq!(violations.len(), 1);
        assert!(violations[0].contains("ssn"));
    }

    #[test]
    fn test_check_output_against_wildcard_blocklist() {
        let declaration = OutputDeclaration {
            has_declaration: true,
            type_string: Some("{ user: { id: string, salary: number } }".to_string()),
            declared_fields: ["user", "id", "salary"]
                .iter()
                .map(|s| s.to_string())
                .collect(),
            has_spread_risk: false,
        };

        let blocked_fields: HashSet<String> = ["*.salary"].iter().map(|s| s.to_string()).collect();

        let violations =
            JavaScriptValidator::check_output_against_blocklist(&declaration, &blocked_fields);
        assert_eq!(violations.len(), 1);
        assert!(violations[0].contains("salary"));
    }
}