//! The conservative source scanner.
//!
//! Heuristic and deterministic: it does not build a full AST or a real taint
//! graph. It finds MCP tool-handler registrations and risky sink calls by
//! line-oriented pattern, then attributes each sink to the nearest preceding
//! handler. That is enough to point a reviewer at a source location and a sink
//! category, which is all a static lead should claim.

use super::model::{Language, ProofFinding, ProofReport, SinkCategory, PROOF_SCHEMA_VERSION};

/// A resolved handler: the tool name, the line it was declared on, the
/// parameter identifiers that are taint sources, and the 1-based line its body
/// ends on (the line before the next handler, or the file end).
struct Handler {
    name: String,
    line: usize,
    params: Vec<String>,
    end_line: usize,
}

/// Scan a single source file's text into findings and resolved handler names.
///
/// Each sink is attributed to its enclosing handler and run through a
/// lightweight intra-function taint heuristic: a handler parameter (a taint
/// source) that reaches the sink within the body, directly or through a local
/// assignment, marks the finding `tainted` and lifts its confidence to `high`;
/// a sink in a handler body that no parameter reaches is downgraded to `low`.
pub fn scan_source(path: &str, language: Language, text: &str) -> (Vec<ProofFinding>, Vec<String>) {
    let lines: Vec<&str> = text.lines().collect();
    let handlers = resolve_handlers(language, &lines);
    let mut findings = Vec::new();
    for (idx, raw) in lines.iter().enumerate() {
        let line_no = idx + 1;
        if let Some(category) = sink_category(language, raw) {
            let handler = enclosing_handler(&handlers, line_no);
            let tool = handler.map(|h| h.name.clone());
            let tainted = handler
                .map(|h| taint_reaches(&lines, h, line_no))
                .unwrap_or(false);
            findings.push(ProofFinding {
                tool,
                category,
                severity: category.severity(),
                language,
                file: path.to_string(),
                line: line_no,
                snippet: raw.trim().to_string(),
                confidence: confidence_for(category, handler.is_some(), tainted).to_string(),
                tainted,
                advisory: true,
                confirmed: false,
            });
        }
    }
    let names = handlers.into_iter().map(|h| h.name).collect();
    (findings, names)
}

/// Assemble a report from per-file scan results, sorted for stable output.
pub fn build_report(
    root: &str,
    mut findings: Vec<ProofFinding>,
    mut handlers: Vec<String>,
) -> ProofReport {
    findings.sort_by(|a, b| {
        a.file
            .cmp(&b.file)
            .then(a.line.cmp(&b.line))
            .then(a.category.cmp(&b.category))
    });
    handlers.sort();
    handlers.dedup();
    ProofReport {
        schema_version: PROOF_SCHEMA_VERSION.to_string(),
        root: root.to_string(),
        handlers,
        findings,
    }
}

/// Resolve tool-handler declarations for a language, with each handler's
/// parameter taint sources and body extent.
fn resolve_handlers(language: Language, lines: &[&str]) -> Vec<Handler> {
    let mut starts: Vec<(usize, String, Vec<String>)> = Vec::new();
    for (idx, raw) in lines.iter().enumerate() {
        if let Some(name) = handler_name(language, raw) {
            let params = extract_params(language, lines, idx);
            starts.push((idx + 1, name, params));
        }
    }
    let total = lines.len();
    let mut out = Vec::with_capacity(starts.len());
    for (i, (line, name, params)) in starts.iter().enumerate() {
        // The body runs until the next handler declaration, or the file end.
        let end_line = starts.get(i + 1).map(|(l, _, _)| l - 1).unwrap_or(total);
        out.push(Handler {
            name: name.clone(),
            line: *line,
            params: params.clone(),
            end_line,
        });
    }
    out
}

/// Extract a handler's parameter identifiers (taint sources). Takes the binding
/// side of each comma-separated parameter (the text before any `: type`
/// annotation), so `cmd: String` yields `cmd` and a destructured `{ cmd, path }`
/// yields both. Skips `self`/`cls` and language keywords. Best-effort and
/// line-local: a parameter list spanning lines yields only the declaration
/// line's tokens.
///
/// In Rust the handler is matched on a `#[tool(...)]` attribute whose parens
/// hold the macro args, not the function parameters, so the declaration line is
/// the nearby `fn NAME(...)`. The lookahead finds it within a few lines.
fn extract_params(language: Language, lines: &[&str], idx: usize) -> Vec<String> {
    let decl = decl_line(language, lines, idx);
    let Some(group) = param_group(language, decl) else {
        return Vec::new();
    };
    let mut out = Vec::new();
    for part in split_top_level(&group) {
        let binding = part.split(':').next().unwrap_or("");
        for token in identifiers(binding) {
            if !PARAM_SKIP.contains(&token.as_str()) && !out.contains(&token) {
                out.push(token);
            }
        }
    }
    out
}

/// The line that carries a handler's parameter list. For Rust, when the match
/// landed on a `#[tool(...)]` attribute, scan forward (a few lines) for the
/// `fn NAME(...)` declaration whose parens hold the real parameters. For every
/// other case the matched line itself is the declaration.
fn decl_line<'a>(language: Language, lines: &[&'a str], idx: usize) -> &'a str {
    let matched = lines[idx];
    if matches!(language, Language::Rust) && !matched.contains("fn ") {
        for line in lines.iter().skip(idx).take(4) {
            if line.contains("fn ") && line.contains('(') {
                return line;
            }
        }
    }
    matched
}

/// The text inside the parameter parentheses for a handler declaration. For an
/// arrow function the parens immediately before `=>` are used; otherwise the
/// first balanced `(...)` group on the line.
fn param_group(language: Language, line: &str) -> Option<String> {
    if matches!(language, Language::TypeScript) {
        if let Some(arrow) = line.find("=>") {
            return balanced_group_before(&line[..arrow]);
        }
    }
    balanced_group_after(line, 0)
}

/// The contents of the first balanced `(...)` at or after `from`.
fn balanced_group_after(s: &str, from: usize) -> Option<String> {
    let bytes = s.as_bytes();
    let open = s[from..].find('(')? + from;
    let mut depth = 0i32;
    for (i, &b) in bytes.iter().enumerate().skip(open) {
        match b {
            b'(' => depth += 1,
            b')' => {
                depth -= 1;
                if depth == 0 {
                    return Some(s[open + 1..i].to_string());
                }
            }
            _ => {}
        }
    }
    None
}

/// The contents of the last balanced `(...)` group that closes within `s`.
fn balanced_group_before(s: &str) -> Option<String> {
    let close = s.rfind(')')?;
    let bytes = s.as_bytes();
    let mut depth = 0i32;
    let mut i = close as isize;
    while i >= 0 {
        match bytes[i as usize] {
            b')' => depth += 1,
            b'(' => {
                depth -= 1;
                if depth == 0 {
                    return Some(s[i as usize + 1..close].to_string());
                }
            }
            _ => {}
        }
        i -= 1;
    }
    None
}

/// Split a parameter list on top-level commas (ignoring commas nested in
/// brackets), so `a, { b, c }, d` splits into three parts.
fn split_top_level(group: &str) -> Vec<String> {
    let mut parts = Vec::new();
    let mut depth = 0i32;
    let mut current = String::new();
    for c in group.chars() {
        match c {
            '(' | '[' | '{' | '<' => depth += 1,
            ')' | ']' | '}' | '>' => depth -= 1,
            ',' if depth == 0 => {
                parts.push(std::mem::take(&mut current));
                continue;
            }
            _ => {}
        }
        current.push(c);
    }
    if !current.trim().is_empty() {
        parts.push(current);
    }
    parts
}

/// Every identifier-like token in a string (a run of `[A-Za-z0-9_]` that does
/// not start with a digit).
fn identifiers(s: &str) -> Vec<String> {
    let mut out = Vec::new();
    let mut current = String::new();
    for c in s.chars() {
        if c.is_alphanumeric() || c == '_' {
            current.push(c);
        } else if !current.is_empty() {
            out.push(std::mem::take(&mut current));
        }
    }
    if !current.is_empty() {
        out.push(current);
    }
    out.retain(|t| !t.chars().next().map(|c| c.is_ascii_digit()).unwrap_or(true));
    out
}

/// Whether `token` appears in `line` as a whole identifier (not as a fragment
/// of a longer identifier like `url` inside `urllib`).
fn contains_token(line: &str, token: &str) -> bool {
    let mut from = 0;
    while let Some(rel) = line[from..].find(token) {
        let start = from + rel;
        let end = start + token.len();
        let before_ok = start == 0
            || !line[..start]
                .chars()
                .next_back()
                .map(|c| c.is_alphanumeric() || c == '_')
                .unwrap_or(false);
        let after_ok = !line[end..]
            .chars()
            .next()
            .map(|c| c.is_alphanumeric() || c == '_')
            .unwrap_or(false);
        if before_ok && after_ok {
            return true;
        }
        from = start + token.len();
    }
    false
}

/// Whether a taint source reaches the sink at `sink_line` within the handler
/// body. Starts from the handler's parameters and propagates through local
/// assignments line by line: a line `lhs = rhs` whose right side mentions a
/// tainted token adds `lhs` to the tainted set. The sink is tainted when any
/// tainted token appears on its line.
fn taint_reaches(lines: &[&str], handler: &Handler, sink_line: usize) -> bool {
    if handler.params.is_empty() {
        return false;
    }
    let mut tainted: Vec<String> = handler.params.clone();
    // Body lines are 1-based; iterate from just after the declaration up to and
    // including the sink line, bounded by the handler body.
    let start = handler.line + 1;
    let end = sink_line.min(handler.end_line);
    for line_no in start..=end {
        let Some(raw) = lines.get(line_no - 1) else {
            continue;
        };
        if line_no == sink_line {
            return tainted.iter().any(|t| contains_token(raw, t));
        }
        propagate_assignment(raw, &mut tainted);
    }
    false
}

/// If `line` is an assignment whose right side mentions a tainted token, add the
/// assigned identifier(s) on the left to the tainted set (a one-line dataflow
/// hop). Recognizes `let x =`, `const x =`, `var x =`, and bare `x =`.
fn propagate_assignment(line: &str, tainted: &mut Vec<String>) {
    let trimmed = line.trim();
    let Some(eq) = trimmed.find('=') else {
        return;
    };
    // Skip comparisons and arrows: `==`, `=>` (char after), and `!=`/`<=`/`>=`
    // (char before).
    let bytes = trimmed.as_bytes();
    if matches!(bytes.get(eq + 1), Some(b'=') | Some(b'>')) {
        return;
    }
    if eq > 0 && matches!(bytes.get(eq - 1), Some(b'!') | Some(b'<') | Some(b'>')) {
        return;
    }
    let (lhs, rhs) = trimmed.split_at(eq);
    let rhs = &rhs[1..];
    if !tainted.iter().any(|t| contains_token(rhs, t)) {
        return;
    }
    let lhs = lhs
        .trim_end()
        .trim_start_matches("let ")
        .trim_start_matches("const ")
        .trim_start_matches("var ");
    for ident in identifiers(lhs) {
        if !tainted.contains(&ident) {
            tainted.push(ident);
        }
    }
}

/// Identifiers that are never taint sources even if they appear in a parameter
/// position.
const PARAM_SKIP: &[&str] = &["self", "cls", "mut", "async", "function", "await"];

/// Extract a tool name from a handler-registration line, if present.
fn handler_name(language: Language, line: &str) -> Option<String> {
    let trimmed = line.trim();
    match language {
        // server.tool("name", ...), registerTool("name"), addTool("name")
        Language::TypeScript => {
            for marker in [".tool(", "registerTool(", "addTool(", "server.tool("] {
                if let Some(name) = quoted_after(trimmed, marker) {
                    return Some(name);
                }
            }
            None
        }
        // @mcp.tool("name") or @mcp.tool() ... def name(  (handled by caller line)
        Language::Python => {
            if trimmed.contains(".tool(") || trimmed.contains("@tool") {
                if let Some(name) = quoted_after(trimmed, ".tool(") {
                    return Some(name);
                }
            }
            // FastMCP style: a decorator with no name uses the function name.
            if let Some(rest) = trimmed.strip_prefix("def ") {
                if let Some(open) = rest.find('(') {
                    // Only treat as a handler when it looks tool-like.
                    let name = rest[..open].trim();
                    if !name.is_empty() && trimmed.ends_with(':') {
                        return Some(name.to_string());
                    }
                }
            }
            None
        }
        // #[tool(name = "x")] or .tool("x") builder, or "x" => match arm
        Language::Rust => {
            for marker in ["#[tool", ".tool(", "name = "] {
                if let Some(name) = quoted_after(trimmed, marker) {
                    return Some(name);
                }
            }
            None
        }
    }
}

/// The first double-quoted token following a marker substring.
fn quoted_after(line: &str, marker: &str) -> Option<String> {
    let after = &line[line.find(marker)? + marker.len()..];
    let start = after.find('"')? + 1;
    let rest = &after[start..];
    let end = rest.find('"')?;
    let name = &rest[..end];
    if name.is_empty() {
        None
    } else {
        Some(name.to_string())
    }
}

/// Classify a line as a risky sink, if it matches a known pattern.
fn sink_category(language: Language, line: &str) -> Option<SinkCategory> {
    let l = line.to_ascii_lowercase();
    // Order matters: most dangerous first so a line is attributed once.
    if matches_any(&l, SHELL_PATTERNS) {
        return Some(SinkCategory::Shell);
    }
    if matches_any(&l, EVAL_PATTERNS) {
        return Some(SinkCategory::DynamicEval);
    }
    if matches_any(&l, CRED_PATTERNS) {
        return Some(SinkCategory::CredentialAccess);
    }
    if matches_any(&l, NET_PATTERNS) {
        return Some(SinkCategory::NetworkEgress);
    }
    if matches_any(&l, FS_WRITE_PATTERNS) {
        return Some(SinkCategory::FilesystemWrite);
    }
    if matches_any(&l, FS_READ_PATTERNS) {
        return Some(SinkCategory::FilesystemRead);
    }
    let _ = language;
    None
}

/// Whether any pattern is a substring of the (lowercased) line.
fn matches_any(line: &str, patterns: &[&str]) -> bool {
    patterns.iter().any(|p| line.contains(p))
}

/// The handler whose body encloses a sink line: the nearest handler declared at
/// or before the sink whose body extends to it.
fn enclosing_handler(handlers: &[Handler], sink_line: usize) -> Option<&Handler> {
    handlers
        .iter()
        .filter(|h| h.line <= sink_line && h.end_line >= sink_line)
        .max_by_key(|h| h.line)
}

/// Heuristic confidence. A tainted source-to-sink flow is `high`. A sink inside
/// a handler body that no parameter reaches is downgraded to `low`. A sink with
/// no resolved handler keeps the category default (shell and eval read as
/// `medium`, the rest `low`), since proximity is all we can claim.
fn confidence_for(category: SinkCategory, has_handler: bool, tainted: bool) -> &'static str {
    if tainted {
        return "high";
    }
    if has_handler {
        return "low";
    }
    match category {
        SinkCategory::Shell | SinkCategory::DynamicEval => "medium",
        _ => "low",
    }
}

const SHELL_PATTERNS: &[&str] = &[
    "exec(",
    "execsync(",
    "spawn(",
    "child_process",
    "subprocess.",
    "os.system",
    "command::new",
    "std::process::command",
    "popen(",
];
const EVAL_PATTERNS: &[&str] = &[
    "eval(",
    "new function(",
    "__import__(",
    "importlib",
    "vm.run",
];
const CRED_PATTERNS: &[&str] = &["process.env", "os.environ", "std::env::var", "getenv("];
const NET_PATTERNS: &[&str] = &[
    "fetch(",
    "axios",
    "requests.",
    "reqwest",
    "urllib",
    "http.get",
    "httpx.",
];
const FS_WRITE_PATTERNS: &[&str] = &[
    "writefile",
    "write_file",
    "std::fs::write",
    "fs.write",
    "open(",
    ".write(",
];
const FS_READ_PATTERNS: &[&str] = &[
    "readfile",
    "read_file",
    "std::fs::read",
    "fs.read",
    "readdir",
    "read_to_string",
];
