"""Data flow analysis for root cause identification.

Uses Semgrep taint mode for industrial-strength data flow tracking,
with AST-based fallback when Semgrep is not available.
"""

from __future__ import annotations

import ast
import copy
import json
import re
import subprocess
import tempfile
import hashlib
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from autodebug.detection.transcript_analyzer import Transcript


@dataclass
class CodeLocation:
    """A location in source code."""

    file: str = ""
    line: int = 0
    column: int = 0
    function: str = ""

    def __str__(self) -> str:
        result = f"{self.file}:{self.line}" if self.file else f"line {self.line}"
        if self.function:
            result += f" in {self.function}()"
        return result


@dataclass
class DataOrigin:
    """The origin of a piece of data."""

    source_type: str  # "api_response", "database", "user_input", "config", "computed"
    location: CodeLocation
    raw_value: Optional[Any] = None
    transcript_seq: Optional[int] = None
    confidence: float = 0.5
    description: str = ""

    def __str__(self) -> str:
        return f"{self.source_type} at {self.location}"


@dataclass
class Transformation:
    """A transformation applied to data."""

    operation: str  # "index", "attribute", "call", "slice", etc.
    location: CodeLocation
    description: str = ""


@dataclass
class DataFlowPath:
    """A path tracing data from origin to failure point."""

    variable: str
    origins: List[DataOrigin] = field(default_factory=list)
    transformations: List[Transformation] = field(default_factory=list)
    failure_point: Optional[CodeLocation] = None
    taint_path: List[str] = field(default_factory=list)

    @property
    def primary_origin(self) -> Optional[DataOrigin]:
        """Get the most likely origin."""
        if not self.origins:
            return None
        # Return highest confidence origin
        return max(self.origins, key=lambda o: o.confidence)

    def format_report(self) -> str:
        """Format the data flow path as a readable report."""
        lines = [f"Data Flow for '{self.variable}':", "=" * 40]

        if self.origins:
            lines.append("\nOrigins:")
            for origin in self.origins:
                conf_pct = int(origin.confidence * 100)
                lines.append(f"  [{conf_pct}%] {origin}")
                if origin.description:
                    lines.append(f"       {origin.description}")

        if self.transformations:
            lines.append("\nTransformations:")
            for t in self.transformations:
                lines.append(f"  - {t.operation} at {t.location}")

        if self.failure_point:
            lines.append(f"\nFailure Point: {self.failure_point}")

        if self.taint_path:
            lines.append("\nTaint Trace:")
            for step in self.taint_path:
                lines.append(f"  -> {step}")

        return "\n".join(lines)


# Semgrep taint rules for different data sources
TAINT_RULES_YAML = """
rules:
  # --- User input to dangerous sinks (actual security risks) ---
  - id: user-input-to-code-execution
    mode: taint
    pattern-sources:
      - patterns:
          - pattern: request.json
      - patterns:
          - pattern: request.form[$KEY]
      - patterns:
          - pattern: request.args.get(...)
      - patterns:
          - pattern: input(...)
      - patterns:
          - pattern: await $REQUEST.json()
      - patterns:
          - pattern: await $REQUEST.body()
      - patterns:
          - pattern: $REQUEST.query_params.get(...)
    pattern-sinks:
      - patterns:
          - pattern: eval($X)
      - patterns:
          - pattern: exec($X)
      - patterns:
          - pattern: subprocess.run($X, ...)
      - patterns:
          - pattern: subprocess.call($X, ...)
      - patterns:
          - pattern: subprocess.Popen($X, ...)
      - patterns:
          - pattern: os.system($X)
      - patterns:
          - pattern: os.popen($X)
    message: "User input flows to code execution sink"
    languages: [python]
    severity: ERROR

  - id: user-input-to-sql
    mode: taint
    pattern-sources:
      - patterns:
          - pattern: request.json
      - patterns:
          - pattern: request.form[$KEY]
      - patterns:
          - pattern: request.args.get(...)
      - patterns:
          - pattern: await $REQUEST.json()
      - patterns:
          - pattern: await $REQUEST.body()
      - patterns:
          - pattern: $REQUEST.query_params.get(...)
    pattern-sinks:
      - patterns:
          - pattern: $CURSOR.execute($QUERY)
      - patterns:
          - pattern: $CONN.execute($QUERY)
      - patterns:
          - pattern: f"...{$X}..."
    message: "User input flows to SQL query (potential injection)"
    languages: [python]
    severity: ERROR

  - id: user-input-to-file-ops
    mode: taint
    pattern-sources:
      - patterns:
          - pattern: request.json
      - patterns:
          - pattern: request.form[$KEY]
      - patterns:
          - pattern: request.args.get(...)
      - patterns:
          - pattern: await $REQUEST.json()
      - patterns:
          - pattern: await $REQUEST.body()
      - patterns:
          - pattern: $REQUEST.query_params.get(...)
    pattern-sinks:
      - patterns:
          - pattern: open($PATH, ...)
      - patterns:
          - pattern: Path($PATH)
      - patterns:
          - pattern: os.path.join(..., $PATH, ...)
    message: "User input flows to file operation (potential path traversal)"
    languages: [python]
    severity: WARNING

  - id: user-input-to-deserialization
    mode: taint
    pattern-sources:
      - patterns:
          - pattern: request.json
      - patterns:
          - pattern: await $REQUEST.json()
      - patterns:
          - pattern: await $REQUEST.body()
    pattern-sinks:
      - patterns:
          - pattern: pickle.loads($X)
      - patterns:
          - pattern: pickle.load($X)
      - patterns:
          - pattern: yaml.load($X)
      - patterns:
          - pattern: yaml.unsafe_load($X)
    message: "User input flows to unsafe deserialization"
    languages: [python]
    severity: ERROR

  # --- API/external data to dangerous sinks ---
  - id: api-response-to-code-execution
    mode: taint
    pattern-sources:
      - patterns:
          - pattern: requests.get(...)
      - patterns:
          - pattern: requests.post(...)
      - patterns:
          - pattern: httpx.get(...)
      - patterns:
          - pattern: httpx.post(...)
      - patterns:
          - pattern: $RESPONSE.json()
      - patterns:
          - pattern: httpx.AsyncClient().get(...)
      - patterns:
          - pattern: httpx.AsyncClient().post(...)
    pattern-sinks:
      - patterns:
          - pattern: eval($X)
      - patterns:
          - pattern: exec($X)
      - patterns:
          - pattern: subprocess.run($X, ...)
      - patterns:
          - pattern: subprocess.Popen($X, ...)
      - patterns:
          - pattern: pickle.loads($X)
      - patterns:
          - pattern: yaml.load($X)
    message: "API response data flows to code execution or deserialization sink"
    languages: [python]
    severity: WARNING

  # --- Config to dangerous sinks ---
  - id: config-to-code-execution
    mode: taint
    pattern-sources:
      - patterns:
          - pattern: os.getenv(...)
      - patterns:
          - pattern: os.environ[$KEY]
      - patterns:
          - pattern: os.environ.get(...)
    pattern-sinks:
      - patterns:
          - pattern: eval($X)
      - patterns:
          - pattern: exec($X)
      - patterns:
          - pattern: subprocess.run($X, shell=True, ...)
      - patterns:
          - pattern: subprocess.call($X, shell=True, ...)
    message: "Config value flows to code execution with shell=True"
    languages: [python]
    severity: WARNING
"""


class SemgrepTaintAnalyzer:
    """Use Semgrep taint mode for data flow analysis.

    Semgrep's taint mode provides industrial-strength tracking of
    data from sources (API, database, config) to sinks (operations).
    """

    def __init__(
        self,
        custom_rules: Optional[str] = None,
        *,
        rules_path: Optional[str] = None,
        cache_max_entries: int = 64,
    ):
        """Initialize the analyzer.

        Args:
            custom_rules: Custom YAML rules (uses default if None)
        """
        # Allow rule override via explicit path or env var (no code edits required per project).
        # Path should point to a YAML file in Semgrep config format.
        env_rules_path = (
            os.getenv("HIKAFLOW_SEMGREP_TAINT_RULES_PATH")
            or os.getenv("HIKAFLOW_TAINT_RULES_PATH")
        )
        effective_path = rules_path or env_rules_path
        if effective_path:
            try:
                self.rules = Path(effective_path).read_text(encoding="utf-8")
            except Exception:
                self.rules = custom_rules or TAINT_RULES_YAML
        else:
            self.rules = custom_rules or TAINT_RULES_YAML
        self._semgrep_available: Optional[bool] = None
        self._cache: Dict[str, List[DataFlowPath]] = {}
        self._cache_max_entries = max(0, int(cache_max_entries))

    @property
    def available(self) -> bool:
        """Check if Semgrep is available."""
        if self._semgrep_available is None:
            self._semgrep_available = self._check_semgrep()
        return self._semgrep_available

    def _check_semgrep(self) -> bool:
        """Check if Semgrep CLI is installed."""
        try:
            result = subprocess.run(
                ["semgrep", "--version"],
                capture_output=True,
                timeout=10,
            )
            return result.returncode == 0
        except (FileNotFoundError, subprocess.TimeoutExpired):
            return False

    def analyze(
        self,
        source_file: str,
        failure_line: int,
        context_lines: int = 20,
    ) -> List[DataFlowPath]:
        """Analyze data flow using Semgrep taint mode.

        Args:
            source_file: Path to source file
            failure_line: Line where failure occurred
            context_lines: Lines around failure to consider

        Returns:
            List of data flow paths
        """
        if not self.available:
            return []

        cache_key = self._cache_key(source_file, failure_line, context_lines)
        if cache_key and cache_key in self._cache:
            return copy.deepcopy(self._cache[cache_key])

        # Use TemporaryDirectory as context manager for safe cleanup
        # (no orphaned files on SIGKILL, matching trunk_analyzer.py pattern)
        with tempfile.TemporaryDirectory() as tmpdir:
            rules_file = str(Path(tmpdir) / "taint_rules.yaml")
            with open(rules_file, "w") as f:
                f.write(self.rules)

            try:
                # Run Semgrep with taint mode
                result = subprocess.run(
                    [
                        "semgrep", "scan",
                        "--config", rules_file,
                        "--json",
                        "--dataflow-traces",
                        source_file,
                    ],
                    capture_output=True,
                    text=True,
                    timeout=60,
                )

                if result.returncode not in [0, 1]:  # 1 = findings
                    return []

                findings = json.loads(result.stdout) if result.stdout.strip() else {}
                paths = self._parse_taint_results(findings, failure_line, context_lines)
                if cache_key:
                    self._cache_put(cache_key, paths)
                return copy.deepcopy(paths)

            except (subprocess.TimeoutExpired, json.JSONDecodeError):
                return []

    def _cache_key(self, source_file: str, failure_line: int, context_lines: int) -> str:
        try:
            p = Path(source_file)
            content = p.read_bytes()
            h = hashlib.sha256()
            h.update(content)
            h.update(b"\0")
            h.update(str(failure_line).encode("utf-8"))
            h.update(b"\0")
            h.update(str(context_lines).encode("utf-8"))
            h.update(b"\0")
            h.update(self.rules.encode("utf-8"))
            return h.hexdigest()
        except Exception:
            return ""

    def _cache_put(self, key: str, value: List[DataFlowPath]) -> None:
        if not key or self._cache_max_entries <= 0:
            return
        if len(self._cache) >= self._cache_max_entries:
            # Drop an arbitrary old key (insertion-ordered dict in py3.7+).
            try:
                self._cache.pop(next(iter(self._cache)))
            except Exception:
                self._cache.clear()
        self._cache[key] = copy.deepcopy(value)

    def _parse_taint_results(
        self,
        findings: Dict,
        failure_line: int,
        context_lines: int,
    ) -> List[DataFlowPath]:
        """Parse Semgrep taint results into DataFlowPath objects."""
        paths = []

        for result in findings.get("results", []):
            # Filter to results near failure line
            result_line = result.get("end", {}).get("line", 0)
            if abs(result_line - failure_line) > context_lines:
                continue

            # Extract taint trace
            taint_trace = []
            trace_data = result.get("extra", {}).get("dataflow_trace", {})
            for trace in trace_data.get("taint_source", []):
                loc = trace.get("location", {})
                trace_str = f"{loc.get('path', '')}:{loc.get('start', {}).get('line', 0)}"
                taint_trace.append(trace_str)

            # Determine origin type from rule
            rule_id = result.get("check_id", "")
            origin_type = self._origin_type_from_rule(rule_id)

            # Extract variable name from metavars
            metavars = result.get("extra", {}).get("metavars", {})
            var_name = "unknown"
            for key in ["$DATA", "$OBJ", "$X"]:
                if key in metavars:
                    var_name = metavars[key].get("abstract_content", "unknown")
                    break

            paths.append(
                DataFlowPath(
                    variable=var_name,
                    origins=[
                        DataOrigin(
                            source_type=origin_type,
                            location=CodeLocation(
                                file=result.get("path", ""),
                                line=result.get("start", {}).get("line", 0),
                            ),
                            confidence=0.9,
                            description=result.get("extra", {}).get("message", ""),
                        )
                    ],
                    transformations=[],
                    failure_point=CodeLocation(
                        file=result.get("path", ""),
                        line=result_line,
                    ),
                    taint_path=taint_trace,
                )
            )

        return paths

    def _origin_type_from_rule(self, rule_id: str) -> str:
        """Determine origin type from rule ID."""
        rule_id = rule_id.lower()
        if "api" in rule_id or "http" in rule_id:
            return "api_response"
        elif "database" in rule_id or "db" in rule_id or "sql" in rule_id:
            return "database"
        elif "config" in rule_id or "env" in rule_id:
            return "config"
        elif "user" in rule_id or "input" in rule_id:
            return "user_input"
        return "unknown"


class DataFlowAnalyzer:
    """Main data flow analyzer combining Semgrep + AST fallback."""

    def __init__(
        self,
        source_code: str,
        transcript: Optional[Transcript] = None,
        file_path: str = "",
    ):
        """Initialize the analyzer.

        Args:
            source_code: Python source code
            transcript: Optional I/O transcript for enrichment
            file_path: Path to source file
        """
        self.source = source_code
        self.transcript = transcript
        self.file_path = file_path
        self.semgrep = SemgrepTaintAnalyzer()

    def trace_variable(
        self,
        var_name: str,
        failure_line: int,
    ) -> DataFlowPath:
        """Trace a variable back to its origins.

        Args:
            var_name: Variable name to trace
            failure_line: Line where failure occurred

        Returns:
            DataFlowPath with origins and transformations
        """
        # First try Semgrep taint analysis
        if self.semgrep.available:
            with tempfile.TemporaryDirectory() as tmpdir:
                source_file = str(Path(tmpdir) / "source.py")
                with open(source_file, "w") as f:
                    f.write(self.source)

                semgrep_paths = self.semgrep.analyze(source_file, failure_line)
                if semgrep_paths:
                    # Enrich with transcript data
                    for path in semgrep_paths:
                        self._enrich_with_transcript(path)
                        # Normalize file path/line to caller-provided context
                        if self.file_path:
                            if path.failure_point:
                                path.failure_point.file = self.file_path
                                if failure_line:
                                    path.failure_point.line = failure_line
                            for origin in path.origins:
                                if origin.location:
                                    origin.location.file = self.file_path
                    return semgrep_paths[0]  # Return best match

        # Fallback to AST-based analysis
        return self._ast_trace(var_name, failure_line)

    def trace_exception_variable(
        self,
        exception_type: str,
        message: str,
        line: int,
    ) -> DataFlowPath:
        """Trace the variable involved in an exception.

        Args:
            exception_type: Type of exception (e.g., "KeyError")
            message: Exception message
            line: Line where exception occurred

        Returns:
            DataFlowPath for the problematic variable
        """
        var_name = self._extract_variable_from_exception(exception_type, message)
        if var_name:
            return self.trace_variable(var_name, line)

        return DataFlowPath(
            variable="unknown",
            origins=[],
            transformations=[],
            failure_point=CodeLocation(file=self.file_path, line=line),
            taint_path=[],
        )

    def _enrich_with_transcript(self, path: DataFlowPath) -> None:
        """Enrich data flow path with transcript data."""
        if not self.transcript:
            return

        for origin in path.origins:
            if origin.source_type == "api_response":
                # Find matching HTTP record
                for http in self.transcript.http_records:
                    if http.seq:
                        origin.transcript_seq = http.seq
                        if http.response_body:
                            origin.raw_value = http.response_body[:500]
                        break

            elif origin.source_type == "database":
                # Find matching SQL record
                for sql in self.transcript.sql_records:
                    if sql.seq:
                        origin.transcript_seq = sql.seq
                        if sql.rows:
                            origin.raw_value = sql.rows[:5]
                        break

    def _find_enclosing_function(self, tree: ast.Module, line: int) -> Optional[ast.AST]:
        """Find the innermost function/class containing the given line."""
        best = None
        for node in ast.walk(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                if hasattr(node, "lineno") and hasattr(node, "end_lineno"):
                    if node.lineno <= line <= (node.end_lineno or node.lineno):
                        if best is None or (node.end_lineno - node.lineno) < (best.end_lineno - best.lineno):
                            best = node
        return best

    def _ast_trace(self, var_name: str, failure_line: int) -> DataFlowPath:
        """Fallback AST-based tracing.

        Scopes search to the enclosing function and handles multiple
        assignment forms (Assign, AugAssign, AnnAssign, NamedExpr, for, with).

        Args:
            var_name: Variable to trace
            failure_line: Line where failure occurred

        Returns:
            DataFlowPath from AST analysis
        """
        path = DataFlowPath(
            variable=var_name,
            origins=[],
            transformations=[],
            failure_point=CodeLocation(file=self.file_path, line=failure_line),
            taint_path=[],
        )

        try:
            tree = ast.parse(self.source)
        except SyntaxError:
            return path

        # Scope search to enclosing function (not the entire file)
        scope = self._find_enclosing_function(tree, failure_line) or tree

        # Find all assignments to this variable before the failure line
        assignments = []
        for node in ast.walk(scope):
            if not hasattr(node, "lineno") or node.lineno >= failure_line:
                continue
            # Skip nested function/class scopes
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                if node is not scope:
                    continue

            # ast.Assign: x = expr
            if isinstance(node, ast.Assign):
                for target in node.targets:
                    if isinstance(target, ast.Name) and target.id == var_name:
                        assignments.append((node.lineno, node))
                        break
                    # Tuple unpacking: a, b = expr
                    elif isinstance(target, ast.Tuple):
                        for elt in target.elts:
                            if isinstance(elt, ast.Name) and elt.id == var_name:
                                assignments.append((node.lineno, node))
                                break
            # ast.AugAssign: x += expr
            elif isinstance(node, ast.AugAssign):
                if isinstance(node.target, ast.Name) and node.target.id == var_name:
                    assignments.append((node.lineno, node))
            # ast.AnnAssign: x: int = expr
            elif isinstance(node, ast.AnnAssign):
                if isinstance(node.target, ast.Name) and node.target.id == var_name and node.value:
                    assignments.append((node.lineno, node))
            # For loop: for x in expr
            elif isinstance(node, ast.For):
                if isinstance(node.target, ast.Name) and node.target.id == var_name:
                    assignments.append((node.lineno, node))
            # With statement: with expr as x
            elif isinstance(node, ast.With):
                for item in node.items:
                    if item.optional_vars and isinstance(item.optional_vars, ast.Name):
                        if item.optional_vars.id == var_name:
                            assignments.append((node.lineno, node))
            # NamedExpr (walrus operator): x := expr
            elif isinstance(node, ast.NamedExpr):
                if isinstance(node.target, ast.Name) and node.target.id == var_name:
                    assignments.append((node.lineno, node))

        # Process assignments (most recent first)
        assignments.sort(key=lambda x: x[0], reverse=True)
        for lineno, node in assignments[:3]:  # Check last 3 assignments
            origin = self._determine_origin(node)
            if origin:
                origin.location.file = self.file_path
                origin.location.line = lineno
                path.origins.append(origin)

        # Alias chain tracking: if x = y (simple Name-to-Name assignment),
        # transitively trace through aliases to find the original source.
        if not path.origins:
            alias_target = self._resolve_alias(var_name, failure_line, scope)
            if alias_target and alias_target != var_name:
                alias_path = self._ast_trace(alias_target, failure_line)
                path.origins.extend(alias_path.origins)
                if alias_path.origins:
                    path.taint_path.append(f"{var_name} -> {alias_target} (alias)")

        return path

    def _resolve_alias(
        self,
        var_name: str,
        before_line: int,
        scope: ast.AST,
    ) -> Optional[str]:
        """Resolve a simple Name-to-Name alias chain.

        When ``x = y`` is encountered, this returns ``y`` so that the
        caller can transitively trace through aliases.

        Args:
            var_name: Variable name to look up
            before_line: Only consider assignments before this line
            scope: AST scope to search within

        Returns:
            The aliased variable name, or None if no alias found.
        """
        for node in ast.walk(scope):
            if not isinstance(node, ast.Assign):
                continue
            if not hasattr(node, "lineno") or node.lineno >= before_line:
                continue
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id == var_name:
                    # Check if value is a simple Name (alias)
                    if isinstance(node.value, ast.Name):
                        return node.value.id
        return None

    def _determine_origin(self, node: ast.AST) -> Optional[DataOrigin]:
        """Determine if assignment is from an external source.

        Handles Assign, AugAssign, AnnAssign, For, and With nodes.

        Args:
            node: AST assignment node

        Returns:
            DataOrigin if external source detected
        """
        # Extract the value expression from different assignment forms
        if isinstance(node, ast.Assign):
            value = node.value
        elif isinstance(node, ast.AugAssign):
            value = node.value
        elif isinstance(node, ast.AnnAssign) and node.value:
            value = node.value
        elif isinstance(node, ast.For):
            value = node.iter
        elif isinstance(node, ast.With) and node.items:
            value = node.items[0].context_expr
        elif isinstance(node, ast.NamedExpr):
            value = node.value
        else:
            return None

        # Handle subscript access: data["key"] or data[0]
        if isinstance(value, ast.Subscript):
            return DataOrigin(
                source_type="computed",
                location=CodeLocation(),
                confidence=0.4,
                description="Value from subscript access",
            )

        # Handle attribute access: obj.attr
        if isinstance(value, ast.Attribute):
            attr_name = self._get_attr_chain(value).lower()
            if any(p in attr_name for p in ["response", "result", "data", "json"]):
                return DataOrigin(
                    source_type="api_response",
                    location=CodeLocation(),
                    confidence=0.6,
                    description="Value from attribute access",
                )
            return None

        if not isinstance(value, ast.Call):
            return None

        func_name = self._get_call_name(value)
        func_name_lower = func_name.lower()

        # Use fully-qualified call name for matching to reduce false positives.
        # For example, "dict.get" should NOT match the ".get" API pattern,
        # but "requests.get" should.

        # API response patterns -- require module prefix to avoid matching dict.get etc.
        _API_PREFIXES = ("requests.", "httpx.", "aiohttp.", "urllib.")
        _API_METHODS = (".json",)
        if any(func_name_lower.startswith(p) for p in _API_PREFIXES):
            return DataOrigin(
                source_type="api_response",
                location=CodeLocation(),
                confidence=0.8,
                description=f"Value from API call ({func_name})",
            )
        if any(func_name_lower.endswith(m) for m in _API_METHODS):
            return DataOrigin(
                source_type="api_response",
                location=CodeLocation(),
                confidence=0.6,
                description=f"Value from .json() call ({func_name})",
            )

        # Supabase patterns (project-specific: api/db.py uses supabase client)
        _SUPABASE_PATTERNS = ("supabase.", ".table.", ".select.", ".insert.", ".update.", ".delete.")
        if any(p in func_name_lower for p in _SUPABASE_PATTERNS):
            return DataOrigin(
                source_type="database",
                location=CodeLocation(),
                confidence=0.8,
                description=f"Value from Supabase query ({func_name})",
            )

        # Database patterns -- use terminal method names to reduce false positives
        _DB_TERMINAL_METHODS = ("fetchone", "fetchall", "first", "one", "one_or_none", "scalar")
        _DB_PREFIXES = ("execute",)
        if any(func_name_lower.endswith(m) for m in _DB_TERMINAL_METHODS):
            return DataOrigin(
                source_type="database",
                location=CodeLocation(),
                confidence=0.8,
                description=f"Value from database query ({func_name})",
            )
        if any(func_name_lower.endswith(p) for p in _DB_PREFIXES):
            return DataOrigin(
                source_type="database",
                location=CodeLocation(),
                confidence=0.6,
                description=f"Value from database execute ({func_name})",
            )

        # Config patterns
        _CONFIG_METHODS = ("getenv", "environ.get")
        _CONFIG_PREFIXES = ("os.getenv", "os.environ", "config.", "settings.")
        if any(func_name_lower.endswith(m) for m in _CONFIG_METHODS) or any(
            func_name_lower.startswith(p) for p in _CONFIG_PREFIXES
        ):
            return DataOrigin(
                source_type="config",
                location=CodeLocation(),
                confidence=0.7,
                description=f"Value from configuration ({func_name})",
            )

        # User input patterns
        _INPUT_PATTERNS = ("input", "request.json", "request.form", "request.args")
        if any(func_name_lower == p or func_name_lower.startswith(p + ".") for p in _INPUT_PATTERNS):
            return DataOrigin(
                source_type="user_input",
                location=CodeLocation(),
                confidence=0.8,
                description=f"Value from user input ({func_name})",
            )
        # FastAPI request body patterns
        if "request.json" in func_name_lower or "await request.json" in func_name_lower:
            return DataOrigin(
                source_type="user_input",
                location=CodeLocation(),
                confidence=0.8,
                description=f"Value from request body ({func_name})",
            )

        return None

    def _get_attr_chain(self, node: ast.Attribute) -> str:
        """Extract full attribute chain (e.g., 'response.data.items')."""
        parts = []
        current = node
        while isinstance(current, ast.Attribute):
            parts.append(current.attr)
            current = current.value
        if isinstance(current, ast.Name):
            parts.append(current.id)
        return ".".join(reversed(parts))

    def _get_call_name(self, node: ast.Call) -> str:
        """Extract function name from Call node."""
        if isinstance(node.func, ast.Name):
            return node.func.id
        elif isinstance(node.func, ast.Attribute):
            parts = []
            current = node.func
            while isinstance(current, ast.Attribute):
                parts.append(current.attr)
                current = current.value
            if isinstance(current, ast.Name):
                parts.append(current.id)
            return ".".join(reversed(parts))
        return ""

    def _extract_variable_from_exception(
        self,
        exc_type: str,
        message: str,
    ) -> Optional[str]:
        """Extract variable name from exception message.

        Args:
            exc_type: Exception type
            message: Exception message

        Returns:
            Variable name if extractable
        """
        patterns = {
            "KeyError": [
                r"KeyError: ['\"](\w+)['\"]",
                r"KeyError: (\w+)",
            ],
            "AttributeError": [
                r"'(\w+)' object has no attribute",
                r"object has no attribute '(\w+)'",
            ],
            "TypeError": [
                r"'(\w+)' object",
            ],
            "NameError": [
                r"name '(\w+)' is not defined",
            ],
            "IndexError": [
                r"list index out of range",  # Variable not in message
            ],
        }

        for pattern in patterns.get(exc_type, []):
            match = re.search(pattern, message)
            if match:
                return match.group(1)

        return None


# Convenience function
def trace_data_flow(
    source_code: str,
    exception_type: str,
    exception_message: str,
    failure_line: int,
    transcript: Optional[Transcript] = None,
) -> DataFlowPath:
    """Convenience function to trace data flow from exception.

    Args:
        source_code: Python source code
        exception_type: Type of exception
        exception_message: Exception message
        failure_line: Line where failure occurred
        transcript: Optional I/O transcript

    Returns:
        DataFlowPath for the exception
    """
    analyzer = DataFlowAnalyzer(source_code, transcript)
    return analyzer.trace_exception_variable(
        exception_type,
        exception_message,
        failure_line,
    )
