"""Data flow primitive detection against parsed Python ASTs.

Identifies entry points (sources of external/untrusted data), sinks (dangerous
consumption points), and sanitizers (trust-boundary crossings that clean data).
These are graph-building primitives used to construct the knowledge graph; they
are intentionally neutral — no severity or compliance mapping happens here.
"""

from __future__ import annotations

import ast
from dataclasses import dataclass
from pathlib import Path

from sanicode.scanner._sql_detect import is_formatted_sql

# ---------------------------------------------------------------------------
# Dataclasses
# ---------------------------------------------------------------------------


@dataclass
class EntryPointInfo:
    """A location where external data enters the program."""

    name: str           # e.g. "handle_login", "os.environ.get"
    kind: str           # "http_handler", "cli_arg", "env_var", "stdin", "file_read"
    file: Path
    line: int
    function: str | None   # enclosing function name, or None at module scope


@dataclass
class SinkInfo:
    """A location where data is consumed in a potentially dangerous way."""

    name: str           # e.g. "cursor.execute", "eval"
    kind: str           # "sql", "command", "eval", "template", "file_write", "network"
    file: Path
    line: int
    function: str | None
    cwe_id: int | None


@dataclass
class SanitizerInfo:
    """A location where data is cleaned or escaped before use."""

    name: str           # e.g. "html.escape", "shlex.quote"
    kind: str           # "html_escape", "shell_quote", "parameterized_query", etc.
    file: Path
    line: int
    function: str | None


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------


def _dotted_name(node: ast.expr) -> str:
    """Return a dotted name string for a Name or Attribute node, or ''."""
    if isinstance(node, ast.Name):
        return node.id
    if isinstance(node, ast.Attribute):
        return f"{_dotted_name(node.value)}.{node.attr}"
    return ""


def _enclosing_function(node: ast.AST, parents: list[ast.AST]) -> str | None:
    """Walk the parent chain to find the nearest enclosing FunctionDef name."""
    for parent in reversed(parents):
        if isinstance(parent, ast.FunctionDef | ast.AsyncFunctionDef):
            return parent.name
    return None


# ---------------------------------------------------------------------------
# Entry point detection
# ---------------------------------------------------------------------------

# Decorator attribute names that signal HTTP handler registration.
_HTTP_HANDLER_DECORATORS = frozenset(
    {
        "route",
        "get",
        "post",
        "put",
        "delete",
        "patch",
        "head",
        "options",
    }
)

# Qualified names whose mere presence is an env-var entry point.
_ENV_VAR_CALLS = frozenset({"os.environ.get", "os.getenv"})


class _EntryPointVisitor(ast.NodeVisitor):
    def __init__(self, file_path: Path) -> None:
        self.file = file_path
        self.results: list[EntryPointInfo] = []
        self._parents: list[ast.AST] = []

    def _add(self, name: str, kind: str, node: ast.AST) -> None:
        self.results.append(
            EntryPointInfo(
                name=name,
                kind=kind,
                file=self.file,
                line=getattr(node, "lineno", 0),
                function=_enclosing_function(node, self._parents),
            )
        )

    def _visit_children(self, node: ast.AST) -> None:
        self._parents.append(node)
        self.generic_visit(node)
        self._parents.pop()

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:  # noqa: N802
        self._check_function(node)
        self._visit_children(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:  # noqa: N802
        self._check_function(node)
        self._visit_children(node)

    def _check_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
        # Check for HTTP handler decorators: @app.route, @router.get, etc.
        for decorator in node.decorator_list:
            if isinstance(decorator, ast.Call):
                func = decorator.func
            else:
                func = decorator

            if isinstance(func, ast.Attribute) and func.attr in _HTTP_HANDLER_DECORATORS:
                self._add(node.name, "http_handler", node)
                break

        # Function parameter named 'request' is a common HTTP entry point.
        param_names = [arg.arg for arg in node.args.args]
        if "request" in param_names:
            # Only add if not already captured via decorator.
            already_captured = any(
                ep.name == node.name and ep.kind == "http_handler" for ep in self.results
            )
            if not already_captured:
                self._add(node.name, "http_handler", node)

    def visit_Call(self, node: ast.Call) -> None:  # noqa: N802
        func = node.func
        qual = _dotted_name(func) if isinstance(func, ast.Attribute | ast.Name) else ""

        # os.environ.get / os.getenv
        if qual in _ENV_VAR_CALLS:
            self._add(qual, "env_var", node)

        # input() / sys.stdin
        elif qual == "input":
            self._add("input", "stdin", node)

        # open() for reading
        elif qual == "open":
            if self._open_is_read(node):
                self._add("open", "file_read", node)

        # Path(...).read_text()
        elif isinstance(func, ast.Attribute) and func.attr == "read_text":
            self._add("Path.read_text", "file_read", node)

        self._parents.append(node)
        self.generic_visit(node)
        self._parents.pop()

    def visit_Subscript(self, node: ast.Subscript) -> None:  # noqa: N802
        """Detect os.environ["KEY"] subscript access."""
        value = node.value
        if _dotted_name(value) == "os.environ":
            self._add("os.environ", "env_var", node)
        self._parents.append(node)
        self.generic_visit(node)
        self._parents.pop()

    def visit_Attribute(self, node: ast.Attribute) -> None:  # noqa: N802
        """Detect sys.argv and sys.stdin attribute access."""
        qual = _dotted_name(node)
        if qual == "sys.argv":
            self._add("sys.argv", "cli_arg", node)
        elif qual == "sys.stdin":
            self._add("sys.stdin", "stdin", node)
        self._parents.append(node)
        self.generic_visit(node)
        self._parents.pop()

    @staticmethod
    def _open_is_read(node: ast.Call) -> bool:
        """Return True when open() call has no mode arg, or mode is 'r'/'rb'."""
        # No arguments at all — default is read.
        if len(node.args) < 2:
            for kw in node.keywords:
                if kw.arg == "mode" and isinstance(kw.value, ast.Constant):
                    return kw.value.value in ("r", "rb", "rt")
            return True
        # Second positional arg is the mode.
        mode_node = node.args[1]
        if isinstance(mode_node, ast.Constant):
            return mode_node.value in ("r", "rb", "rt")
        return False


def detect_entry_points(tree: ast.Module, file_path: Path) -> list[EntryPointInfo]:
    """Detect locations where external data enters the program.

    Args:
        tree: Parsed ast.Module for the file.
        file_path: Source file path (used to populate EntryPointInfo.file).

    Returns:
        List of EntryPointInfo objects (may be empty).
    """
    visitor = _EntryPointVisitor(file_path)
    visitor.visit(tree)
    return visitor.results


# ---------------------------------------------------------------------------
# Sink detection
# ---------------------------------------------------------------------------

# Subprocess functions that become sinks only when shell=True.
_SUBPROCESS_FUNCS = frozenset({"call", "run", "Popen", "check_output", "check_call"})

# Fully-qualified names that are unconditional network sinks.
_NETWORK_SINKS = frozenset(
    {
        "requests.get",
        "requests.post",
        "requests.put",
        "requests.delete",
        "requests.patch",
        "requests.head",
        "httpx.get",
        "httpx.post",
        "httpx.put",
        "httpx.delete",
        "httpx.patch",
        "urllib.request.urlopen",
    }
)

# Template rendering calls → CWE-79.
_TEMPLATE_RENDER_FUNCS = frozenset({"render_template", "render"})


class _SinkVisitor(ast.NodeVisitor):
    def __init__(self, file_path: Path) -> None:
        self.file = file_path
        self.results: list[SinkInfo] = []
        self._parents: list[ast.AST] = []

    def _add(self, name: str, kind: str, node: ast.AST, cwe_id: int | None) -> None:
        self.results.append(
            SinkInfo(
                name=name,
                kind=kind,
                file=self.file,
                line=getattr(node, "lineno", 0),
                function=_enclosing_function(node, self._parents),
                cwe_id=cwe_id,
            )
        )

    def _visit_children(self, node: ast.AST) -> None:
        self._parents.append(node)
        self.generic_visit(node)
        self._parents.pop()

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:  # noqa: N802
        self._visit_children(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:  # noqa: N802
        self._visit_children(node)

    def visit_Call(self, node: ast.Call) -> None:  # noqa: N802
        func = node.func
        qual = _dotted_name(func) if isinstance(func, ast.Attribute | ast.Name) else ""

        if isinstance(func, ast.Name):
            name = func.id
            if name == "eval":
                self._add("eval", "eval", node, 94)
            elif name == "exec":
                self._add("exec", "eval", node, 94)
            elif name == "__import__":
                self._add("__import__", "eval", node, 94)
            elif name in _TEMPLATE_RENDER_FUNCS:
                self._add(name, "template", node, 79)
            elif name == "open" and self._open_is_write(node):
                self._add("open", "file_write", node, 73)

        elif isinstance(func, ast.Attribute):
            attr = func.attr

            if qual == "os.system":
                self._add("os.system", "command", node, 78)

            elif attr in _SUBPROCESS_FUNCS and self._has_shell_true(node):
                self._add(f"subprocess.{attr}", "command", node, 78)

            elif attr == "execute" and is_formatted_sql(node):
                self._add("cursor.execute", "sql", node, 89)

            elif qual == "pickle.loads":
                self._add("pickle.loads", "eval", node, 502)

            elif qual == "yaml.load" and not self._has_loader_kwarg(node):
                self._add("yaml.load", "eval", node, 502)

            elif attr in ("render_template",) or qual.endswith(".render"):
                self._add(qual or attr, "template", node, 79)

            elif qual in _NETWORK_SINKS:
                self._add(qual, "network", node, 918)

            elif attr == "write_text":
                self._add("Path.write_text", "file_write", node, 73)

            elif attr == "open" and self._open_is_write(node):
                self._add("open", "file_write", node, 73)

        self._parents.append(node)
        self.generic_visit(node)
        self._parents.pop()

    @staticmethod
    def _has_shell_true(node: ast.Call) -> bool:
        return any(
            kw.arg == "shell" and isinstance(kw.value, ast.Constant) and kw.value.value
            for kw in node.keywords
        )

    @staticmethod
    def _has_loader_kwarg(node: ast.Call) -> bool:
        return any(kw.arg == "Loader" for kw in node.keywords)

    @staticmethod
    def _open_is_write(node: ast.Call) -> bool:
        """Return True when open() is called in a write mode."""
        write_modes = {"w", "wb", "wt", "a", "ab", "at", "x", "xb", "xt"}
        if len(node.args) >= 2:
            mode_node = node.args[1]
            if isinstance(mode_node, ast.Constant):
                return mode_node.value in write_modes
        for kw in node.keywords:
            if kw.arg == "mode" and isinstance(kw.value, ast.Constant):
                return kw.value.value in write_modes
        return False


def detect_sinks(tree: ast.Module, file_path: Path) -> list[SinkInfo]:
    """Detect locations where data is consumed in a potentially dangerous way.

    Args:
        tree: Parsed ast.Module for the file.
        file_path: Source file path.

    Returns:
        List of SinkInfo objects (may be empty).
    """
    visitor = _SinkVisitor(file_path)
    visitor.visit(tree)
    return visitor.results


# ---------------------------------------------------------------------------
# Sanitizer detection
# ---------------------------------------------------------------------------

# Qualified call names that indicate HTML escaping.
_HTML_ESCAPE_CALLS = frozenset(
    {"html.escape", "markupsafe.escape", "bleach.clean", "escape", "strip_tags"}
)

# URL encoding calls.
_URL_ENCODE_CALLS = frozenset({"urllib.parse.quote", "urllib.parse.urlencode"})


class _SanitizerVisitor(ast.NodeVisitor):
    def __init__(self, file_path: Path) -> None:
        self.file = file_path
        self.results: list[SanitizerInfo] = []
        self._parents: list[ast.AST] = []

    def _add(self, name: str, kind: str, node: ast.AST) -> None:
        self.results.append(
            SanitizerInfo(
                name=name,
                kind=kind,
                file=self.file,
                line=getattr(node, "lineno", 0),
                function=_enclosing_function(node, self._parents),
            )
        )

    def _visit_children(self, node: ast.AST) -> None:
        self._parents.append(node)
        self.generic_visit(node)
        self._parents.pop()

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:  # noqa: N802
        self._visit_children(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:  # noqa: N802
        self._visit_children(node)

    def visit_Call(self, node: ast.Call) -> None:  # noqa: N802
        func = node.func
        qual = _dotted_name(func) if isinstance(func, ast.Attribute | ast.Name) else ""

        if qual in _HTML_ESCAPE_CALLS:
            self._add(qual, "html_escape", node)

        elif qual == "shlex.quote":
            self._add("shlex.quote", "shell_quote", node)

        elif qual in _URL_ENCODE_CALLS:
            self._add(qual, "url_encode", node)

        elif isinstance(func, ast.Attribute) and func.attr == "execute":
            # cursor.execute(query, params) — parameterized query when a second
            # argument (tuple or list) is provided.
            if len(node.args) >= 2 and isinstance(node.args[1], ast.Tuple | ast.List):
                self._add("cursor.execute", "parameterized_query", node)
            else:
                # Check for keyword-style params argument.
                for kw in node.keywords:
                    if kw.arg == "params" and isinstance(kw.value, ast.Tuple | ast.List):
                        self._add("cursor.execute", "parameterized_query", node)
                        break

        self._parents.append(node)
        self.generic_visit(node)
        self._parents.pop()


def detect_sanitizers(tree: ast.Module, file_path: Path) -> list[SanitizerInfo]:
    """Detect locations where data is cleaned or escaped before use.

    Args:
        tree: Parsed ast.Module for the file.
        file_path: Source file path.

    Returns:
        List of SanitizerInfo objects (may be empty).
    """
    visitor = _SanitizerVisitor(file_path)
    visitor.visit(tree)
    return visitor.results
