# ---------------------------------------------------------------------------
#  Copyright (c) DevSecCode. All rights reserved.
# ---------------------------------------------------------------------------
"""Data flow analyzer — finds taint paths from sources to sinks by propagating
through the code intelligence call graph.

Operates on the resolved_edges (CALLS, IMPORTS) and the fine-grained
variable_defs / call_arguments / return_expressions tables in the code intel
SQLite store.  Does NOT require re-parsing source; it reuses the data that the
code intel engine has already extracted.

Analysis strategies (in priority order):
1. **Variable-flow intra-procedural taint** — within each function, builds a
   taint map from variable_defs (which variables hold source-call results),
   propagates through variable-to-variable assignments, then checks
   call_arguments to confirm that a tainted variable actually flows into a
   sink call.  This catches ``let x = req.body; db.query(x)`` even when x
   passes through intermediate assignments.
2. **Raw-call intra-function analysis** — classifies raw_calls directly
   against the taint catalog; finds source→sink pairs in the same function
   by line-number ordering (legacy fallback for scopes with no extracted
   variable_defs).
3. **Edge-based BFS** — forward traversal through resolved_edges (CALLS)
   from source symbols to sink symbols for inter-file taint paths.
4. **Return-value inter-procedural propagation** — if a function returns a
   tainted variable (via return_expressions), callers that consume the result
   inherit taint transitively.
"""

from __future__ import annotations

import logging
from collections import deque
from dataclasses import dataclass, field
from typing import Optional

from dsc_api.services.code_intel.cpg.taint_catalog import (
    TaintCatalog,
    TaintMatch,
    TaintRole,
    VulnCategory,
)
from dsc_api.services.code_intel.store import CodeIntelStore

logger = logging.getLogger(__name__)

# Maximum BFS depth to prevent runaway analysis
MAX_TAINT_DEPTH = 12

# Maximum paths to return per analysis run
MAX_PATHS = 500


@dataclass
class TaintNode:
    """A symbol classified with its taint role."""
    symbol_id: int
    symbol_name: str
    qualified_name: str
    file_path: str
    start_line: int
    end_line: int
    role: TaintRole
    match: Optional[TaintMatch] = None


@dataclass
class TaintPath:
    """A complete source-to-sink data flow path."""
    source: TaintNode
    sink: TaintNode
    hops: list[TaintNode] = field(default_factory=list)
    category: Optional[VulnCategory] = None
    cwes: tuple[str, ...] = ()
    risk_score: float = 0.0
    sanitized: bool = False
    sanitizer: Optional[TaintNode] = None

    @property
    def length(self) -> int:
        return len(self.hops) + 2  # source + hops + sink

    @property
    def file_paths(self) -> list[str]:
        paths = [self.source.file_path]
        for hop in self.hops:
            if hop.file_path not in paths:
                paths.append(hop.file_path)
        if self.sink.file_path not in paths:
            paths.append(self.sink.file_path)
        return paths


@dataclass
class DataFlowResult:
    """Result of a data flow analysis run."""
    paths: list[TaintPath]
    sources_found: int
    sinks_found: int
    sanitizers_found: int
    total_symbols_analyzed: int
    truncated: bool = False  # True if MAX_PATHS was hit


# Paths that should never be analyzed for taint (vendored/generated code)
_VENDOR_PATH_SEGMENTS = (
    "node_modules/", "venv/", ".venv/", "env/", ".env/",
    "site-packages/", "dist-packages/", "__pycache__/",
    ".tox/", ".nox/", ".mypy_cache/", "vendor/",
    "google-cloud-sdk/", ".git/",
)


class DataFlowAnalyzer:
    """Taint propagation analyzer over the code intel graph.

    Usage:
        analyzer = DataFlowAnalyzer(store, catalog)
        result = analyzer.analyze(language="typescript")
        for path in result.paths:
            print(f"{path.source.symbol_name} -> ... -> {path.sink.symbol_name} [{path.category}]")
    """

    def __init__(self, store: CodeIntelStore, catalog: Optional[TaintCatalog] = None) -> None:
        self._store = store
        self._catalog = catalog or TaintCatalog()

    def analyze(
        self,
        language: str = "",
        max_depth: int = MAX_TAINT_DEPTH,
        max_paths: int = MAX_PATHS,
        target_categories: Optional[list[VulnCategory]] = None,
        file_filter: Optional[list[str]] = None,
    ) -> DataFlowResult:
        """Run full taint propagation analysis.

        Uses three complementary strategies:
        1. **Variable-flow intra-procedural**: uses variable_defs + call_arguments
           to track ``let x = source(); sink(x)`` through assignments.
           Falls back to line-number heuristic when variable data is absent.
        2. **Edge-based BFS**: traverses resolved_edges (CALLS) from source symbols
           to sink symbols for inter-file taint paths.
        3. **Return-value inter-procedural**: if a function returns a tainted var,
           callers consuming the result inherit taint transitively.
        """
        # Step 1: Classify all symbols
        sources, sinks, sanitizers, propagators, total = self._classify_symbols(
            language, file_filter
        )

        logger.info(
            "Taint classification: %d sources, %d sinks, %d sanitizers, %d propagators out of %d symbols",
            len(sources), len(sinks), len(sanitizers), len(propagators), total,
        )

        # Step 2: Intra-function analysis (variable-flow + line heuristic)
        raw_paths = self._analyze_raw_calls(language, file_filter, target_categories)
        logger.info("Intra-procedural analysis found %d taint paths", len(raw_paths))

        paths: list[TaintPath] = list(raw_paths)
        truncated = False

        # Step 3: Inter-procedural return-value propagation
        return_paths = self._analyze_return_propagation(
            language, file_filter, target_categories
        )
        if return_paths:
            logger.info("Return-value analysis found %d inter-procedural paths", len(return_paths))
            paths.extend(return_paths)

        # Step 4: Edge-based BFS (only if we have both source and sink symbols)
        if sources and sinks:
            filtered_sinks = sinks
            if target_categories:
                cat_set = set(target_categories)
                filtered_sinks = {sid: node for sid, node in sinks.items()
                         if node.match and node.match.entry.category in cat_set}

            sink_ids = set(filtered_sinks.keys())
            sanitizer_ids = set(sanitizers.keys())

            for source_id, source_node in sources.items():
                if len(paths) >= max_paths:
                    truncated = True
                    break

                found = self._bfs_forward(
                    source_id, source_node,
                    sink_ids, filtered_sinks,
                    sanitizer_ids, sanitizers,
                    max_depth, max_paths - len(paths),
                )
                paths.extend(found)

        # Deduplicate in two passes:
        # Pass 1: Exact dedup — same (file, source_line, sink_line)
        seen_exact: set[tuple[str, int, int]] = set()
        deduped: list[TaintPath] = []
        for p in paths:
            key = (p.source.file_path, p.source.start_line, p.sink.start_line)
            if key not in seen_exact:
                seen_exact.add(key)
                deduped.append(p)

        # Pass 2: Structural dedup — same (source_symbol, sink_symbol, hop_symbols, category)
        # across different files. Keeps the first occurrence (highest priority).
        seen_structural: set[tuple] = set()
        structural_deduped: list[TaintPath] = []
        for p in deduped:
            hop_names = tuple(h.symbol_name for h in p.hops)
            struct_key = (p.source.symbol_name, p.sink.symbol_name, hop_names, p.category)
            if struct_key not in seen_structural:
                seen_structural.add(struct_key)
                structural_deduped.append(p)
            else:
                # Still count it for source/sink stats but don't include in results
                pass

        if len(structural_deduped) < len(deduped):
            logger.info("Structural dedup: %d -> %d paths (removed %d duplicates across files)",
                        len(deduped), len(structural_deduped), len(deduped) - len(structural_deduped))
        paths = structural_deduped

        # Combine source/sink counts from all strategies
        seen_sources: set[int] = set(sources.keys())
        seen_sinks: set[int] = set(sinks.keys())
        for p in raw_paths:
            seen_sources.add(p.source.symbol_id)
            seen_sinks.add(p.sink.symbol_id)

        # Step 5: Score and sort paths
        for path in paths:
            path.risk_score = self._score_path(path)

        paths.sort(key=lambda p: -p.risk_score)

        return DataFlowResult(
            paths=paths[:max_paths],
            sources_found=len(seen_sources),
            sinks_found=len(seen_sinks),
            sanitizers_found=len(sanitizers),
            total_symbols_analyzed=total,
            truncated=truncated,
        )

    def find_paths_to_sink(
        self,
        sink_symbol_id: int,
        max_depth: int = MAX_TAINT_DEPTH,
    ) -> list[TaintPath]:
        """Reverse BFS: find all tainted sources that can reach a specific sink."""
        sink_row = self._store.db.execute(
            "SELECT s.*, f.path as file_path FROM symbols s JOIN files f ON f.id = s.file_id WHERE s.id=?",
            (sink_symbol_id,),
        ).fetchone()
        if not sink_row:
            return []

        # Classify the sink
        sink_matches = self._catalog.classify(sink_row["name"])
        sink_match = next((m for m in sink_matches if m.entry.role == TaintRole.SINK), None)
        if not sink_match:
            return []

        sink_node = TaintNode(
            symbol_id=sink_symbol_id,
            symbol_name=sink_row["name"],
            qualified_name=sink_row["qualified_name"],
            file_path=sink_row["file_path"],
            start_line=sink_row["start_line"],
            end_line=sink_row["end_line"],
            role=TaintRole.SINK,
            match=sink_match,
        )

        # Reverse BFS through incoming CALLS edges
        paths: list[TaintPath] = []
        queue: deque[tuple[int, list[int]]] = deque([(sink_symbol_id, [])])
        visited: set[int] = set()

        while queue:
            current_id, trail = queue.popleft()
            if current_id in visited or len(trail) > max_depth:
                continue
            visited.add(current_id)

            # Check if current is a source
            row = self._store.db.execute(
                "SELECT s.*, f.path as file_path FROM symbols s JOIN files f ON f.id = s.file_id WHERE s.id=?",
                (current_id,),
            ).fetchone()
            if row:
                matches = self._catalog.classify(row["name"])
                source_match = next((m for m in matches if m.entry.role == TaintRole.SOURCE), None)
                if source_match and current_id != sink_symbol_id:
                    source_node = TaintNode(
                        symbol_id=current_id, symbol_name=row["name"],
                        qualified_name=row["qualified_name"], file_path=row["file_path"],
                        start_line=row["start_line"], end_line=row["end_line"],
                        role=TaintRole.SOURCE, match=source_match,
                    )
                    hops = self._resolve_trail(trail)
                    path = TaintPath(
                        source=source_node, sink=sink_node, hops=hops,
                        category=sink_match.entry.category, cwes=sink_match.entry.cwes,
                    )
                    path.risk_score = self._score_path(path)
                    paths.append(path)

            # Follow incoming edges (reverse direction)
            callers = self._store.db.execute(
                "SELECT source_symbol_id FROM resolved_edges WHERE target_symbol_id=? AND kind IN ('CALLS', 'IMPORTS')",
                (current_id,),
            ).fetchall()
            for r in callers:
                caller_id = r["source_symbol_id"]
                if caller_id not in visited:
                    queue.append((caller_id, trail + [current_id]))

        paths.sort(key=lambda p: -p.risk_score)
        return paths

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    # ------------------------------------------------------------------
    # Variable-flow intra-procedural taint analysis
    # ------------------------------------------------------------------

    def _analyze_raw_calls(
        self,
        language: str,
        file_filter: Optional[list[str]],
        target_categories: Optional[list[VulnCategory]],
    ) -> list[TaintPath]:
        """Find intra-function taint paths using variable-flow data when
        available, falling back to line-number heuristics otherwise.

        This is the primary intra-procedural analysis. It uses three tables:
        - ``raw_calls``:  every call expression, grouped by enclosing function
        - ``variable_defs``:  ``let x = source()`` records x as tainted via
          source_call_id, and ``let y = x`` records y as tainted via
          source_variable pointing to x
        - ``call_arguments``:  ``sink(x)`` records that variable x is the 0th
          argument of the sink call, so we can confirm data actually flows in
        """
        paths: list[TaintPath] = []
        cat_set = set(target_categories) if target_categories else None

        # ---- Load raw calls grouped by caller scope ----
        query = """
            SELECT c.id, c.file_id, c.caller_symbol_id, c.callee_name, c.receiver, c.line, c.col,
                   s.name as caller_name, s.qualified_name as caller_qn,
                   s.start_line as caller_start, s.end_line as caller_end,
                   f.path as file_path, f.language
            FROM raw_calls c
            LEFT JOIN symbols s ON s.id = c.caller_symbol_id
            LEFT JOIN files f ON f.id = c.file_id
        """
        params: list = []
        conditions: list[str] = []
        if language:
            conditions.append("f.language = ?")
            params.append(language)
        # Exclude vendored/generated paths at the SQL level
        for seg in _VENDOR_PATH_SEGMENTS:
            conditions.append("f.path NOT LIKE ?")
            params.append(f"%{seg}%")
        if conditions:
            query += " WHERE " + " AND ".join(conditions)
        query += " ORDER BY c.caller_symbol_id, c.line"

        rows = self._store.db.execute(query, params).fetchall()

        calls_by_caller: dict[int, list] = {}
        call_by_id: dict[int, dict] = {}
        for row in rows:
            caller_id = row["caller_symbol_id"]
            if caller_id is None:
                continue
            if file_filter and row["file_path"] not in file_filter:
                continue
            calls_by_caller.setdefault(caller_id, []).append(row)
            call_by_id[row["id"]] = row

        # ---- Load variable_defs grouped by scope (only for relevant scopes) ----
        # Limit to scopes that actually appear in our filtered raw_calls
        scope_ids = set(calls_by_caller.keys())
        vardefs_by_scope: dict[int, list] = {}
        if scope_ids:
            # Batch query — SQLite has a limit of ~999 params, so chunk if needed
            scope_list = list(scope_ids)
            for chunk_start in range(0, len(scope_list), 900):
                chunk = scope_list[chunk_start:chunk_start + 900]
                placeholders = ",".join("?" * len(chunk))
                vd_rows = self._store.db.execute(
                    f"SELECT * FROM variable_defs WHERE scope_symbol_id IN ({placeholders})",
                    chunk,
                ).fetchall()
                for vd in vd_rows:
                    sid = vd["scope_symbol_id"]
                    vardefs_by_scope.setdefault(sid, []).append(vd)

        # ---- Load call_arguments only for calls in our scope ----
        all_call_ids = [row["id"] for rows in calls_by_caller.values() for row in rows]
        callargs_by_call: dict[int, list] = {}
        if all_call_ids:
            for chunk_start in range(0, len(all_call_ids), 900):
                chunk = all_call_ids[chunk_start:chunk_start + 900]
                placeholders = ",".join("?" * len(chunk))
                ca_rows = self._store.db.execute(
                    f"SELECT * FROM call_arguments WHERE call_id IN ({placeholders}) ORDER BY call_id, argument_index",
                    chunk,
                ).fetchall()
                for ca in ca_rows:
                    callargs_by_call.setdefault(ca["call_id"], []).append(ca)

        # ---- Per-function analysis ----
        for caller_id, calls in calls_by_caller.items():
            if not calls:
                continue

            # Classify each call as source/sink/sanitizer
            source_calls: list[tuple[dict, TaintMatch]] = []
            sink_calls: list[tuple[dict, TaintMatch]] = []
            sanitizer_calls: list[tuple[dict, TaintMatch]] = []

            for call in calls:
                callee = call["callee_name"]
                receiver = call["receiver"]
                lang = call["language"] or language

                names_to_try = [callee]
                if receiver:
                    names_to_try.insert(0, f"{receiver}.{callee}")

                for name in names_to_try:
                    matches = self._catalog.classify(name, lang)
                    if not matches:
                        continue
                    best = matches[0]
                    if best.entry.role == TaintRole.SOURCE:
                        source_calls.append((call, best))
                    elif best.entry.role == TaintRole.SINK:
                        if cat_set is None or (best.entry.category and best.entry.category in cat_set):
                            sink_calls.append((call, best))
                    elif best.entry.role == TaintRole.SANITIZER:
                        sanitizer_calls.append((call, best))
                    break

            if not source_calls or not sink_calls:
                continue

            caller_info = calls[0]
            file_path = caller_info["file_path"]
            caller_name = caller_info["caller_name"] or f"function@{caller_id}"
            caller_qn = caller_info["caller_qn"] or caller_name

            # Build the variable taint map for this scope
            scope_vardefs = vardefs_by_scope.get(caller_id, [])
            has_varflow = len(scope_vardefs) > 0

            if has_varflow:
                # -- Strategy 1: variable-flow taint propagation --
                found = self._analyze_scope_varflow(
                    caller_id, caller_info, file_path, caller_name, caller_qn,
                    source_calls, sink_calls, sanitizer_calls,
                    scope_vardefs, callargs_by_call, call_by_id,
                )
                paths.extend(found)
            else:
                # -- Strategy 2 (fallback): line-number heuristic --
                found = self._analyze_scope_line_heuristic(
                    caller_id, caller_info, file_path, caller_name, caller_qn,
                    source_calls, sink_calls, sanitizer_calls,
                )
                paths.extend(found)

        return paths

    def _analyze_scope_varflow(
        self,
        caller_id: int,
        caller_info: dict,
        file_path: str,
        caller_name: str,
        caller_qn: str,
        source_calls: list[tuple[dict, TaintMatch]],
        sink_calls: list[tuple[dict, TaintMatch]],
        sanitizer_calls: list[tuple[dict, TaintMatch]],
        scope_vardefs: list,
        callargs_by_call: dict[int, list],
        call_by_id: dict[int, dict],
    ) -> list[TaintPath]:
        """Variable-flow-aware intra-procedural taint analysis.

        Builds a taint map: variable_name -> (source_call, source_match, line).
        Then propagates through variable-to-variable assignments.
        Finally checks call_arguments on sink calls to confirm a tainted var
        actually flows into the sink.
        """
        paths: list[TaintPath] = []

        # Step 1: Build source_call_id -> (call_row, match) lookup
        source_call_ids: dict[int, tuple[dict, TaintMatch]] = {}
        for call, match in source_calls:
            source_call_ids[call["id"]] = (call, match)

        # Step 2: Build initial taint set from variable_defs whose
        # source_call_id points to a classified source call.
        # tainted_vars: var_name -> (originating_source_call, source_match, def_line)
        tainted_vars: dict[str, tuple[dict, TaintMatch, int]] = {}
        # Also track parameter taint (e.g., function params named "req", "body")
        for vd in scope_vardefs:
            src_call_id = vd["source_call_id"]
            if src_call_id and src_call_id in source_call_ids:
                src_call, src_match = source_call_ids[src_call_id]
                tainted_vars[vd["variable_name"]] = (src_call, src_match, vd["line"])
            elif vd["source_type"] == "parameter":
                # Check if the parameter name itself matches a source pattern
                var_name = vd["variable_name"]
                matches = self._catalog.classify(var_name)
                if matches and matches[0].entry.role == TaintRole.SOURCE:
                    # Synthetic source call for parameter-sourced taint
                    param_call = {
                        "id": -(vd["id"]),
                        "callee_name": var_name,
                        "receiver": None,
                        "line": vd["line"],
                        "file_path": file_path,
                    }
                    tainted_vars[var_name] = (param_call, matches[0], vd["line"])

        # Step 3: Propagate taint through variable-to-variable and
        # property-access assignments. Sort defs by line so we propagate
        # in source order.
        sorted_defs = sorted(scope_vardefs, key=lambda v: v["line"])
        changed = True
        max_iters = 10  # prevent infinite loops in pathological code
        iteration = 0
        while changed and iteration < max_iters:
            changed = False
            iteration += 1
            for vd in sorted_defs:
                vname = vd["variable_name"]
                if vname in tainted_vars:
                    continue
                src_type = vd["source_type"]
                if src_type == "identifier":
                    src_var = vd["source_variable"]
                    if src_var and src_var in tainted_vars:
                        tainted_vars[vname] = tainted_vars[src_var]
                        changed = True
                elif src_type == "property_access":
                    # e.g., const name = req.body.name -- check if base ("req") is tainted
                    prop = vd["source_property"] or ""
                    base = prop.split(".")[0] if "." in prop else prop
                    if base in tainted_vars:
                        tainted_vars[vname] = tainted_vars[base]
                        changed = True
                elif src_type == "destructuring":
                    # e.g., const { name, email } = req.body
                    prop = vd["source_property"] or ""
                    base = prop.split(".")[0] if "." in prop else prop
                    if base in tainted_vars:
                        tainted_vars[vname] = tainted_vars[base]
                        changed = True

        if not tainted_vars:
            return paths

        # Step 4: Check which sanitizer calls process tainted vars
        # sanitized_vars: var_name -> line of sanitization
        sanitized_vars: dict[str, int] = {}
        sanitizer_lines: dict[int, tuple[dict, TaintMatch]] = {}
        for san_call, san_match in sanitizer_calls:
            san_args = callargs_by_call.get(san_call["id"], [])
            for arg in san_args:
                if arg["kind"] == "identifier" and arg["variable_name"] in tainted_vars:
                    sanitized_vars[arg["variable_name"]] = san_call["line"]
                    sanitizer_lines[san_call["line"]] = (san_call, san_match)
            # Also check if the result of the sanitizer overwrites a tainted var
            for vd in sorted_defs:
                if vd["source_call_id"] == san_call["id"] and vd["variable_name"] in tainted_vars:
                    sanitized_vars[vd["variable_name"]] = san_call["line"]
                    sanitizer_lines[san_call["line"]] = (san_call, san_match)

        # Step 5: For each sink call, check its arguments for tainted variables
        for snk_call, snk_match in sink_calls:
            snk_args = callargs_by_call.get(snk_call["id"], [])
            matched_tainted_arg = False
            taint_source: Optional[tuple[dict, TaintMatch, int]] = None

            for arg in snk_args:
                arg_var = arg["variable_name"]
                if arg["kind"] == "identifier" and arg_var and arg_var in tainted_vars:
                    matched_tainted_arg = True
                    taint_source = tainted_vars[arg_var]
                    break
                elif arg["kind"] == "property_access":
                    # e.g., sink(obj.prop) -- check if obj is tainted
                    chain = arg["property_chain"] or ""
                    base = chain.split(".")[0] if "." in chain else (arg_var or "")
                    if base in tainted_vars:
                        matched_tainted_arg = True
                        taint_source = tainted_vars[base]
                        break
                elif arg["kind"] == "call_result" and arg["nested_call_id"]:
                    # e.g., sink(source()) -- check if the nested call is a source
                    nested_id = arg["nested_call_id"]
                    if nested_id in source_call_ids:
                        matched_tainted_arg = True
                        taint_source = (*source_call_ids[nested_id], arg["line"])
                        break

            if not matched_tainted_arg:
                # Fallback: if no call_arguments were extracted but source
                # appears before sink in the same scope, use line heuristic
                for src_call, src_match in source_calls:
                    if src_call["line"] < snk_call["line"]:
                        matched_tainted_arg = True
                        taint_source = (src_call, src_match, src_call["line"])
                        break

            if not matched_tainted_arg or not taint_source:
                continue

            src_call, src_match, src_line = taint_source

            # Check sanitization
            sanitized = False
            sanitizer_node = None
            # Any sanitizer between source line and sink line?
            for san_call, san_match in sanitizer_calls:
                if src_line <= san_call["line"] <= snk_call["line"]:
                    sanitized = True
                    san_name = san_call["callee_name"]
                    if san_call["receiver"]:
                        san_name = f"{san_call['receiver']}.{san_name}"
                    sanitizer_node = TaintNode(
                        symbol_id=-san_call["id"],
                        symbol_name=san_name,
                        qualified_name=san_name,
                        file_path=file_path,
                        start_line=san_call["line"],
                        end_line=san_call["line"],
                        role=TaintRole.SANITIZER,
                        match=san_match,
                    )
                    break

            # Build the path nodes
            src_name = src_call["callee_name"]
            if src_call["receiver"]:
                src_name = f"{src_call['receiver']}.{src_name}"
            source_node = TaintNode(
                symbol_id=-abs(src_call["id"]),
                symbol_name=src_name,
                qualified_name=src_name,
                file_path=file_path,
                start_line=src_call["line"],
                end_line=src_call["line"],
                role=TaintRole.SOURCE,
                match=src_match,
            )

            snk_name = snk_call["callee_name"]
            if snk_call["receiver"]:
                snk_name = f"{snk_call['receiver']}.{snk_name}"
            sink_node = TaintNode(
                symbol_id=-snk_call["id"],
                symbol_name=snk_name,
                qualified_name=snk_name,
                file_path=file_path,
                start_line=snk_call["line"],
                end_line=snk_call["line"],
                role=TaintRole.SINK,
                match=snk_match,
            )

            hop = TaintNode(
                symbol_id=caller_id,
                symbol_name=caller_name,
                qualified_name=caller_qn,
                file_path=file_path,
                start_line=caller_info["caller_start"] or src_call["line"],
                end_line=caller_info["caller_end"] or snk_call["line"],
                role=TaintRole.PROPAGATOR,
            )

            paths.append(TaintPath(
                source=source_node,
                sink=sink_node,
                hops=[hop],
                category=snk_match.entry.category,
                cwes=snk_match.entry.cwes,
                sanitized=sanitized,
                sanitizer=sanitizer_node,
            ))

        return paths

    def _analyze_scope_line_heuristic(
        self,
        caller_id: int,
        caller_info: dict,
        file_path: str,
        caller_name: str,
        caller_qn: str,
        source_calls: list[tuple[dict, TaintMatch]],
        sink_calls: list[tuple[dict, TaintMatch]],
        sanitizer_calls: list[tuple[dict, TaintMatch]],
    ) -> list[TaintPath]:
        """Fallback: match source→sink by line ordering within a function.

        Used when no variable_defs have been extracted for the scope (e.g.,
        Python files before the py_extractor is upgraded).
        """
        paths: list[TaintPath] = []

        for src_call, src_match in source_calls:
            for snk_call, snk_match in sink_calls:
                if src_call["line"] > snk_call["line"]:
                    continue

                sanitized = False
                sanitizer_node = None
                for san_call, san_match in sanitizer_calls:
                    if src_call["line"] <= san_call["line"] <= snk_call["line"]:
                        sanitized = True
                        san_name = san_call["callee_name"]
                        if san_call["receiver"]:
                            san_name = f"{san_call['receiver']}.{san_name}"
                        sanitizer_node = TaintNode(
                            symbol_id=-san_call["id"],
                            symbol_name=san_name,
                            qualified_name=san_name,
                            file_path=file_path,
                            start_line=san_call["line"],
                            end_line=san_call["line"],
                            role=TaintRole.SANITIZER,
                            match=san_match,
                        )
                        break

                src_name = src_call["callee_name"]
                if src_call["receiver"]:
                    src_name = f"{src_call['receiver']}.{src_name}"
                source_node = TaintNode(
                    symbol_id=-src_call["id"],
                    symbol_name=src_name,
                    qualified_name=src_name,
                    file_path=file_path,
                    start_line=src_call["line"],
                    end_line=src_call["line"],
                    role=TaintRole.SOURCE,
                    match=src_match,
                )

                snk_name = snk_call["callee_name"]
                if snk_call["receiver"]:
                    snk_name = f"{snk_call['receiver']}.{snk_name}"
                sink_node = TaintNode(
                    symbol_id=-snk_call["id"],
                    symbol_name=snk_name,
                    qualified_name=snk_name,
                    file_path=file_path,
                    start_line=snk_call["line"],
                    end_line=snk_call["line"],
                    role=TaintRole.SINK,
                    match=snk_match,
                )

                hop = TaintNode(
                    symbol_id=caller_id,
                    symbol_name=caller_name,
                    qualified_name=caller_qn,
                    file_path=file_path,
                    start_line=caller_info["caller_start"] or src_call["line"],
                    end_line=caller_info["caller_end"] or snk_call["line"],
                    role=TaintRole.PROPAGATOR,
                )

                paths.append(TaintPath(
                    source=source_node,
                    sink=sink_node,
                    hops=[hop],
                    category=snk_match.entry.category,
                    cwes=snk_match.entry.cwes,
                    sanitized=sanitized,
                    sanitizer=sanitizer_node,
                ))

        return paths

    # ------------------------------------------------------------------
    # Inter-procedural return-value taint propagation
    # ------------------------------------------------------------------

    def _analyze_return_propagation(
        self,
        language: str,
        file_filter: Optional[list[str]],
        target_categories: Optional[list[VulnCategory]],
    ) -> list[TaintPath]:
        """Find inter-procedural taint via return values.

        If function ``getInput()`` returns a tainted variable (tracked by
        return_expressions), and another function calls ``getInput()`` and
        passes its result to a sink, that's a two-hop taint path.

        Algorithm:
        1. Find return_expressions where kind='identifier' and the variable
           was tainted in its scope (source_call in variable_defs).
        2. For those functions, find callers via raw_calls or resolved_edges.
        3. Check if the caller stores the result in a variable, and that
           variable flows into a sink.
        """
        paths: list[TaintPath] = []
        cat_set = set(target_categories) if target_categories else None

        # Find functions that return tainted variables
        ret_query = """
            SELECT r.*, s.name as scope_name, s.qualified_name as scope_qn,
                   s.start_line as scope_start, s.end_line as scope_end,
                   f.path as file_path, f.language
            FROM return_expressions r
            JOIN symbols s ON s.id = r.scope_symbol_id
            JOIN files f ON f.id = r.file_id
            WHERE r.kind = 'identifier' AND r.variable_name IS NOT NULL
        """
        ret_params: list = []
        if language:
            ret_query += " AND f.language = ?"
            ret_params.append(language)
        for seg in _VENDOR_PATH_SEGMENTS:
            ret_query += " AND f.path NOT LIKE ?"
            ret_params.append(f"%{seg}%")

        ret_rows = self._store.db.execute(ret_query, ret_params).fetchall()
        if not ret_rows:
            return paths

        # For each return expression, check if the returned variable is tainted
        tainted_returns: dict[int, tuple[dict, TaintMatch]] = {}  # scope_symbol_id -> (source_call_info, match)

        for ret in ret_rows:
            if file_filter and ret["file_path"] not in file_filter:
                continue
            scope_id = ret["scope_symbol_id"]
            var_name = ret["variable_name"]

            # Check variable_defs for this scope: is var_name assigned from a source?
            vd_rows = self._store.db.execute(
                """SELECT vd.*, rc.callee_name, rc.receiver, rc.line as call_line
                   FROM variable_defs vd
                   LEFT JOIN raw_calls rc ON rc.id = vd.source_call_id
                   WHERE vd.scope_symbol_id = ? AND vd.variable_name = ?""",
                (scope_id, var_name),
            ).fetchall()

            for vd in vd_rows:
                if vd["source_type"] != "call_result" or not vd["callee_name"]:
                    continue
                callee = vd["callee_name"]
                receiver = vd["receiver"]
                name = f"{receiver}.{callee}" if receiver else callee
                lang = ret["language"] or language
                matches = self._catalog.classify(name, lang)
                if not matches:
                    matches = self._catalog.classify(callee, lang)
                if matches and matches[0].entry.role == TaintRole.SOURCE:
                    tainted_returns[scope_id] = (
                        {
                            "id": vd["source_call_id"] or -(vd["id"]),
                            "callee_name": callee,
                            "receiver": receiver,
                            "line": vd["call_line"] or vd["line"],
                            "file_path": ret["file_path"],
                        },
                        matches[0],
                    )
                    break

        if not tainted_returns:
            return paths

        # Now find callers of these tainted functions, and check if the
        # caller passes the result to a sink
        for scope_id, (src_info, src_match) in tainted_returns.items():
            # Get the function name/qualified_name
            scope_row = self._store.db.execute(
                "SELECT s.*, f.path as file_path FROM symbols s JOIN files f ON f.id = s.file_id WHERE s.id = ?",
                (scope_id,),
            ).fetchone()
            if not scope_row:
                continue

            fn_name = scope_row["name"]

            # Find calls to this function
            caller_calls = self._store.db.execute(
                """SELECT c.id, c.file_id, c.caller_symbol_id, c.callee_name, c.receiver, c.line,
                          f.path as file_path, f.language,
                          cs.name as caller_name, cs.qualified_name as caller_qn,
                          cs.start_line as caller_start, cs.end_line as caller_end
                   FROM raw_calls c
                   LEFT JOIN files f ON f.id = c.file_id
                   LEFT JOIN symbols cs ON cs.id = c.caller_symbol_id
                   WHERE c.callee_name = ?""",
                (fn_name,),
            ).fetchall()

            for caller_call in caller_calls:
                if file_filter and caller_call["file_path"] not in file_filter:
                    continue
                caller_scope = caller_call["caller_symbol_id"]
                if not caller_scope:
                    continue

                # Check if any variable in the caller is assigned from this call
                vd_consumer = self._store.db.execute(
                    """SELECT * FROM variable_defs
                       WHERE scope_symbol_id = ? AND source_call_id = ?""",
                    (caller_scope, caller_call["id"]),
                ).fetchall()

                # Get the tainted variable name(s) in the caller scope
                tainted_in_caller: set[str] = set()
                for vd in vd_consumer:
                    tainted_in_caller.add(vd["variable_name"])

                if not tainted_in_caller:
                    # Could also be used directly: sink(getInput())
                    # Check call_arguments for nested call
                    pass

                # Find sink calls in the caller scope
                sink_calls_in_caller = self._store.db.execute(
                    """SELECT c.id, c.callee_name, c.receiver, c.line, f.language
                       FROM raw_calls c
                       LEFT JOIN files f ON f.id = c.file_id
                       WHERE c.caller_symbol_id = ? AND c.line > ?""",
                    (caller_scope, caller_call["line"]),
                ).fetchall()

                for snk_call in sink_calls_in_caller:
                    callee = snk_call["callee_name"]
                    receiver = snk_call["receiver"]
                    lang = snk_call["language"] or language
                    name = f"{receiver}.{callee}" if receiver else callee

                    snk_matches = self._catalog.classify(name, lang)
                    if not snk_matches:
                        snk_matches = self._catalog.classify(callee, lang)
                    if not snk_matches or snk_matches[0].entry.role != TaintRole.SINK:
                        continue
                    snk_match = snk_matches[0]
                    if cat_set and snk_match.entry.category not in cat_set:
                        continue

                    # Confirm the tainted variable flows into the sink's arguments
                    confirmed = False
                    snk_args = self._store.db.execute(
                        "SELECT * FROM call_arguments WHERE call_id = ?",
                        (snk_call["id"],),
                    ).fetchall()

                    for arg in snk_args:
                        if arg["kind"] == "identifier" and arg["variable_name"] in tainted_in_caller:
                            confirmed = True
                            break

                    if not confirmed and not tainted_in_caller:
                        # Fallback: line ordering if no var data
                        confirmed = True

                    if not confirmed:
                        continue

                    # Build the 2-hop inter-procedural path
                    src_name = src_info["callee_name"]
                    if src_info["receiver"]:
                        src_name = f"{src_info['receiver']}.{src_name}"
                    source_node = TaintNode(
                        symbol_id=-abs(src_info["id"]),
                        symbol_name=src_name,
                        qualified_name=src_name,
                        file_path=src_info["file_path"],
                        start_line=src_info["line"],
                        end_line=src_info["line"],
                        role=TaintRole.SOURCE,
                        match=src_match,
                    )

                    snk_name = name
                    sink_node = TaintNode(
                        symbol_id=-snk_call["id"],
                        symbol_name=snk_name,
                        qualified_name=snk_name,
                        file_path=caller_call["file_path"],
                        start_line=snk_call["line"],
                        end_line=snk_call["line"],
                        role=TaintRole.SINK,
                        match=snk_match,
                    )

                    # Two hops: the tainted function, then the caller function
                    hop_fn = TaintNode(
                        symbol_id=scope_id,
                        symbol_name=scope_row["name"],
                        qualified_name=scope_row["qualified_name"],
                        file_path=scope_row["file_path"],
                        start_line=scope_row["start_line"],
                        end_line=scope_row["end_line"],
                        role=TaintRole.PROPAGATOR,
                    )
                    hop_caller = TaintNode(
                        symbol_id=caller_scope,
                        symbol_name=caller_call["caller_name"] or f"function@{caller_scope}",
                        qualified_name=caller_call["caller_qn"] or caller_call["caller_name"] or "",
                        file_path=caller_call["file_path"],
                        start_line=caller_call["caller_start"] or caller_call["line"],
                        end_line=caller_call["caller_end"] or snk_call["line"],
                        role=TaintRole.PROPAGATOR,
                    )

                    paths.append(TaintPath(
                        source=source_node,
                        sink=sink_node,
                        hops=[hop_fn, hop_caller],
                        category=snk_match.entry.category,
                        cwes=snk_match.entry.cwes,
                    ))

        return paths

    def _classify_symbols(
        self, language: str, file_filter: Optional[list[str]],
    ) -> tuple[
        dict[int, TaintNode],  # sources
        dict[int, TaintNode],  # sinks
        dict[int, TaintNode],  # sanitizers
        dict[int, TaintNode],  # propagators
        int,                   # total symbols
    ]:
        """Classify all symbols in the store by taint role.

        For small indexes (<50K symbols) we load all symbols and classify in
        Python.  For very large indexes we apply a coarse SQL LIKE pre-filter
        using the key segment of each catalog pattern.
        """
        sources: dict[int, TaintNode] = {}
        sinks: dict[int, TaintNode] = {}
        sanitizers: dict[int, TaintNode] = {}
        propagators: dict[int, TaintNode] = {}

        base = "SELECT s.*, f.path as file_path, f.language FROM symbols s JOIN files f ON f.id = s.file_id WHERE s.kind != 'module'"
        params: list = []

        if language:
            base += " AND f.language = ?"
            params.append(language)

        # Exclude vendored paths
        for seg in _VENDOR_PATH_SEGMENTS:
            base += " AND f.path NOT LIKE ?"
            params.append(f"%{seg}%")

        # Check total count — only apply LIKE pre-filter for large indexes
        count_sql = f"SELECT COUNT(*) FROM symbols s JOIN files f ON f.id = s.file_id WHERE s.kind != 'module'"
        count_params: list = []
        if language:
            count_sql += " AND f.language = ?"
            count_params.append(language)
        sym_count = self._store.db.execute(count_sql, count_params).fetchone()[0]

        if sym_count > 50_000:
            # Extract key non-wildcard segments from catalog patterns for LIKE
            like_clauses: list[str] = []
            seen: set[str] = set()
            for entry in self._catalog._entries:
                # Extract the last meaningful segment: "*.request.query.*" -> "query"
                parts = [p for p in entry.pattern.split(".") if p != "*"]
                if parts:
                    key = parts[-1]
                    if key not in seen:
                        seen.add(key)
                        like_clauses.append(f"%{key}%")

            if like_clauses:
                # Match against both name and qualified_name
                name_clauses = " OR ".join(["s.name LIKE ?"] * len(like_clauses))
                qn_clauses = " OR ".join(["s.qualified_name LIKE ?"] * len(like_clauses))
                base += f" AND (({name_clauses}) OR ({qn_clauses}))"
                params.extend(like_clauses)
                params.extend(like_clauses)

        rows = self._store.db.execute(base, params).fetchall()
        total = len(rows)

        for row in rows:
            if file_filter and row["file_path"] not in file_filter:
                continue

            sym_name = row["name"]
            lang = row["language"] or language
            matches = self._catalog.classify(sym_name, lang)

            if not matches:
                continue

            # Take the highest-priority match
            best = matches[0]
            node = TaintNode(
                symbol_id=row["id"],
                symbol_name=sym_name,
                qualified_name=row["qualified_name"],
                file_path=row["file_path"],
                start_line=row["start_line"],
                end_line=row["end_line"],
                role=best.entry.role,
                match=best,
            )

            if best.entry.role == TaintRole.SOURCE:
                sources[row["id"]] = node
            elif best.entry.role == TaintRole.SINK:
                sinks[row["id"]] = node
            elif best.entry.role == TaintRole.SANITIZER:
                sanitizers[row["id"]] = node
            elif best.entry.role == TaintRole.PROPAGATOR:
                propagators[row["id"]] = node

        # Also classify call expressions from raw_calls table.
        # The symbols table has declarations; raw_calls has actual call sites
        # like fetch(), eval(), db.collection() that the catalog targets.
        #
        # Performance: extract catalog callee keywords and pre-filter at SQL
        # level so we only load rows that could possibly match (turns 700K+
        # rows into a few thousand on large workspaces).
        exact_callees: set[str] = set()
        like_callees: list[str] = []
        for entry in self._catalog._entries:
            parts = [p for p in entry.pattern.split(".") if p != "*"]
            if not parts:
                continue
            last = parts[-1]
            if "*" not in entry.pattern:
                # Exact pattern like "eval" or "JSON.parse" -> last segment
                exact_callees.add(last)
            else:
                # Wildcard pattern -> LIKE match on callee_name
                if last not in exact_callees:
                    exact_callees.add(last)

        call_sql = """
            SELECT c.id, c.callee_name, c.receiver, c.caller_symbol_id,
                   c.line, f.path as file_path, f.language
            FROM raw_calls c
            JOIN files f ON f.id = c.file_id
        """
        call_conditions: list[str] = []
        call_params: list = []
        if language:
            call_conditions.append("f.language = ?")
            call_params.append(language)
        for seg in _VENDOR_PATH_SEGMENTS:
            call_conditions.append("f.path NOT LIKE ?")
            call_params.append(f"%{seg}%")

        # Callee name pre-filter: only load calls whose callee_name matches
        # a catalog keyword (exact IN or LIKE for partial matches)
        if exact_callees:
            callee_list = sorted(exact_callees)
            # Chunk to stay within SQLite param limits
            chunks = [callee_list[i:i + 400] for i in range(0, len(callee_list), 400)]
            in_clauses = []
            for chunk in chunks:
                placeholders = ",".join("?" * len(chunk))
                in_clauses.append(f"c.callee_name IN ({placeholders})")
                call_params.extend(chunk)
            call_conditions.append(f"({' OR '.join(in_clauses)})")

        if call_conditions:
            call_sql += " WHERE " + " AND ".join(call_conditions)

        call_rows = self._store.db.execute(call_sql, call_params).fetchall()
        total += len(call_rows)

        for crow in call_rows:
            if file_filter and crow["file_path"] not in file_filter:
                continue

            # Build qualified call name: "receiver.callee" or just "callee"
            callee = crow["callee_name"]
            receiver = crow["receiver"]
            if receiver:
                call_name = f"{receiver}.{callee}"
            else:
                call_name = callee

            lang = crow["language"] or language
            matches = self._catalog.classify(call_name, lang)
            if not matches:
                # Try just the callee name (e.g. "fetch", "eval")
                matches = self._catalog.classify(callee, lang)

            if not matches:
                continue

            best = matches[0]
            # Use the caller symbol as the node ID for graph connectivity
            # (the call site belongs to its enclosing function)
            node_id = crow["caller_symbol_id"] or -(crow["id"])
            # Skip if this symbol already classified from the symbols table
            if node_id in sources or node_id in sinks or node_id in sanitizers:
                continue

            node = TaintNode(
                symbol_id=node_id,
                symbol_name=call_name,
                qualified_name=call_name,
                file_path=crow["file_path"],
                start_line=crow["line"],
                end_line=crow["line"],
                role=best.entry.role,
                match=best,
            )

            if best.entry.role == TaintRole.SOURCE:
                sources[node_id] = node
            elif best.entry.role == TaintRole.SINK:
                sinks[node_id] = node
            elif best.entry.role == TaintRole.SANITIZER:
                sanitizers[node_id] = node
            elif best.entry.role == TaintRole.PROPAGATOR:
                propagators[node_id] = node

        return sources, sinks, sanitizers, propagators, total

    def _bfs_forward(
        self,
        source_id: int,
        source_node: TaintNode,
        sink_ids: set[int],
        sinks: dict[int, TaintNode],
        sanitizer_ids: set[int],
        sanitizers: dict[int, TaintNode],
        max_depth: int,
        max_paths: int,
    ) -> list[TaintPath]:
        """Forward BFS from a source, looking for sinks."""
        paths: list[TaintPath] = []

        # Queue: (current_symbol_id, trail_of_ids, sanitized_by)
        queue: deque[tuple[int, list[int], Optional[int]]] = deque(
            [(source_id, [], None)]
        )
        visited: set[int] = set()

        while queue and len(paths) < max_paths:
            current_id, trail, sanitized_by = queue.popleft()

            if current_id in visited or len(trail) > max_depth:
                continue
            visited.add(current_id)

            # Check if we hit a sanitizer
            if current_id in sanitizer_ids and current_id != source_id:
                sanitized_by = current_id

            # Check if we hit a sink
            if current_id in sink_ids and current_id != source_id:
                sink_node = sinks[current_id]
                hops = self._resolve_trail(trail)

                path = TaintPath(
                    source=source_node,
                    sink=sink_node,
                    hops=hops,
                    category=sink_node.match.entry.category if sink_node.match else None,
                    cwes=sink_node.match.entry.cwes if sink_node.match else (),
                    sanitized=sanitized_by is not None,
                    sanitizer=sanitizers.get(sanitized_by) if sanitized_by else None,
                )
                paths.append(path)
                continue  # Don't traverse past sinks

            # Follow outgoing CALLS edges, excluding vendored targets
            callees = self._store.db.execute(
                """SELECT e.target_symbol_id
                   FROM resolved_edges e
                   JOIN symbols s ON s.id = e.target_symbol_id
                   JOIN files f ON f.id = s.file_id
                   WHERE e.source_symbol_id = ? AND e.kind IN ('CALLS', 'IMPORTS')
                     AND f.path NOT LIKE '%node_modules/%'
                     AND f.path NOT LIKE '%venv/%'
                     AND f.path NOT LIKE '%.venv/%'
                     AND f.path NOT LIKE '%site-packages/%'
                     AND f.path NOT LIKE '%vendor/%'
                     AND f.path NOT LIKE '%google-cloud-sdk/%'
                     AND f.path NOT LIKE '%.git/%'""",
                (current_id,),
            ).fetchall()
            for r in callees:
                callee_id = r["target_symbol_id"]
                if callee_id not in visited:
                    queue.append((callee_id, trail + [current_id], sanitized_by))

        return paths

    def _resolve_trail(self, symbol_ids: list[int]) -> list[TaintNode]:
        """Convert a list of symbol IDs to TaintNode objects."""
        nodes: list[TaintNode] = []
        for sid in symbol_ids:
            row = self._store.db.execute(
                "SELECT s.*, f.path as file_path FROM symbols s JOIN files f ON f.id = s.file_id WHERE s.id=?",
                (sid,),
            ).fetchone()
            if row:
                matches = self._catalog.classify(row["name"])
                best_match = matches[0] if matches else None
                role = best_match.entry.role if best_match else TaintRole.PROPAGATOR
                nodes.append(TaintNode(
                    symbol_id=sid,
                    symbol_name=row["name"],
                    qualified_name=row["qualified_name"],
                    file_path=row["file_path"],
                    start_line=row["start_line"],
                    end_line=row["end_line"],
                    role=role,
                    match=best_match,
                ))
        return nodes

    @staticmethod
    def _score_path(path: TaintPath) -> float:
        """Score a taint path for risk prioritization.

        Higher = more dangerous. Factors:
        - Shorter paths are higher risk (more direct data flow)
        - Sanitized paths are lower risk
        - Critical categories score higher
        - Higher source/sink confidence = higher risk
        """
        category_weights = {
            VulnCategory.SQL_INJECTION: 10.0,
            VulnCategory.COMMAND_INJECTION: 10.0,
            VulnCategory.CODE_INJECTION: 10.0,
            VulnCategory.DESERIALIZATION: 9.0,
            VulnCategory.SSRF: 8.5,
            VulnCategory.XXE: 8.0,
            VulnCategory.SSTI: 8.0,
            VulnCategory.XSS: 7.5,
            VulnCategory.PATH_TRAVERSAL: 7.0,
            VulnCategory.FILE_WRITE: 7.0,
            VulnCategory.FILE_DELETE: 7.0,
            VulnCategory.NOSQL_INJECTION: 7.0,
            VulnCategory.LDAP_INJECTION: 7.0,
            VulnCategory.AUTH_BYPASS: 9.0,
            VulnCategory.DYNAMIC_IMPORT: 6.5,
            VulnCategory.EMAIL_INJECTION: 5.5,
            VulnCategory.OPEN_REDIRECT: 5.0,
            VulnCategory.LOG_INJECTION: 3.0,
        }

        base = category_weights.get(path.category, 5.0)

        # Shorter = more dangerous (inverse relationship, capped)
        length_factor = 1.0 / max(path.length - 1, 1)

        # Sanitized paths are much lower risk
        sanitizer_factor = 0.1 if path.sanitized else 1.0

        # Confidence from source and sink matches
        source_conf = path.source.match.confidence if path.source.match else 0.5
        sink_conf = path.sink.match.confidence if path.sink.match else 0.5
        confidence_factor = (source_conf + sink_conf) / 2.0

        return base * length_factor * sanitizer_factor * confidence_factor
