"""
AI_LLM01 – LLM-powered audit of “hot” GenAI calls.

✓ Skips when all tainted inputs are sanitised
✓ Understands helper names  escape_braces(),  sanitize_*,  redact_* …
✓ Sends the *entire surrounding function* to the model for context
"""

from __future__ import annotations

import ast
import json
import logging
import re
import textwrap
from typing import Optional

from lintai.engine.analysis import ProjectAnalyzer
from lintai.core.finding import Finding
from lintai.detectors import register
from lintai.llm import get_client
from lintai.engine.classification import is_ai_related as is_ai_call
from lintai.engine.ast_utils import get_full_attr_name

logger = logging.getLogger(__name__)
_CLIENT = get_client()  # dummy stub if provider missing

# --------------------------------------------------------------------------- #
# constants / patterns                                                        #
# --------------------------------------------------------------------------- #
_CODE_RE = re.compile(r"```(?:json)?\s*(\{.*?})\s*```", re.S | re.I)

_SANITIZERS = {"escape_braces", "sanitize", "redact_secrets"}
_SANITIZER_RE = re.compile(r"(^|\.)\s*(sanitize|escape|redact|clean)\w*$", re.I)

_IGNORE_PREFIXES = (
    "llm detection disabled",
    "openai sdk error",
    "sdk unavailable",
    "clean",  # model replied “issue = clean”
)
_EMITTED: set[tuple[str, int, str]] = set()


# --------------------------------------------------------------------------- #
# helper functions                                                             #
# --------------------------------------------------------------------------- #
import ast
import textwrap


def _snippet(node: ast.AST, src: str, max_lines: int = 60) -> str:
    """
    Return cleaned, dedented source for *node*, trimmed to *max_lines*.

    • Strips module / function / class doc-strings (so the LLM
      does not waste tokens on doc text).
    • Drops single-line comments and blank lines.
    • Leaves triple-quoted strings that are *inside* the code logic
      (those are usually prompts we *do* want the model to see).
    """
    # --- 1. get raw text for the node -----------------------------------
    raw = ast.get_source_segment(src, node) or ""
    if not raw:
        return "<code unavailable>"

    # --- 2. parse & delete doc-strings ----------------------------------
    class _DocstringStripper(ast.NodeTransformer):
        def _strip(self, n: ast.AST):
            if (
                n.body
                and isinstance(n.body[0], ast.Expr)
                and isinstance(n.body[0].value, ast.Constant)
                and isinstance(n.body[0].value.value, str)
            ):
                n.body = n.body[1:]

        def visit_FunctionDef(self, n):
            self.generic_visit(n)
            self._strip(n)
            return n

        def visit_AsyncFunctionDef(self, n):
            self.generic_visit(n)
            self._strip(n)
            return n

        def visit_ClassDef(self, n):
            self.generic_visit(n)
            self._strip(n)
            return n

        def visit_Module(self, n):
            self.generic_visit(n)
            self._strip(n)
            return n

    try:
        tree = ast.parse(raw)
        tree = _DocstringStripper().visit(tree)
        ast.fix_missing_locations(tree)
        cleaned = ast.unparse(tree)
    except Exception:  # fallback – never fail hard
        cleaned = raw

    # --- 3. dedent + drop comments / blanks -----------------------------
    lines = textwrap.dedent(cleaned).splitlines()
    lines = [ln for ln in lines if ln.strip() and not ln.lstrip().startswith("#")]

    if len(lines) > max_lines:
        lines = lines[:max_lines] + ["    # …trimmed…"]

    return "\n".join(lines)


def _path_context(unit, func_node: ast.AST, max_funcs: int = 3) -> str:
    """
    Using the project-wide call-graph gather up to *max_funcs* direct
    callers **and** direct callees of *func_node* and return their source
    text (clipped) as one big string.
    """
    try:
        from lintai.engine import ai_analyzer  # late import → no cycle

        if not ai_analyzer:
            return ""

        qname = unit.qualname(func_node)  # PythonASTUnit helper

        callers = list(ai_analyzer.callers_of(qname))[:max_funcs]

        callers_src = []
        for name in callers:
            try:
                src_unit, node = ai_analyzer.source_of(name)  # (PythonASTUnit, ast.AST)
                callers_src.append(_snippet(node, src_unit.source))
            except Exception as e:
                logger.debug(f"llm_code_audit: skipping caller {name} - no source: {e}")

        san_callees = set()
        san_callee_blocks = []
        for name in ai_analyzer.callees_of(qname):
            if any(fn in name for fn in _SANITIZERS) or _SANITIZER_RE.search(name):
                try:
                    src_unit, node = ai_analyzer.source_of(name)
                    san_callees.add(name)
                    san_callee_blocks.append(_snippet(node, src_unit.source))
                except Exception as e:
                    logger.debug(
                        f"llm_code_audit: skipping sanitizer callee {name} - no source: {e}"
                    )

        other_callees = [
            c for c in ai_analyzer.callees_of(qname) if c not in san_callees
        ][:max_funcs]

        callees_src = san_callee_blocks
        for name in other_callees:
            try:
                src_unit, node = ai_analyzer.source_of(name)
                callees_src.append(_snippet(node, src_unit.source))
            except Exception as e:
                logger.debug(f"llm_code_audit: skipping callee {name} - no source: {e}")

        if not callers_src and not callees_src:
            return ""

        parts = ["\n### CALL-FLOW CONTEXT ###"]
        if callers_src:
            parts.append("\n\n# ⇧  DIRECT CALLERS\n" + "\n\n".join(callers_src))
        if callees_src:
            parts.append("\n\n# ⇩  DIRECT CALLEES\n" + "\n\n".join(callees_src))
        return "\n".join(parts)

    except Exception as exc:
        logger.debug("llm_code_audit: context error – %s", exc)
        return ""


def _is_sanitizer(name: str) -> bool:
    return name in _SANITIZERS or bool(_SANITIZER_RE.match(name))


def _json_fragment(reply: str) -> Optional[str]:
    """Extract the `{…}` JSON object from an LLM reply."""
    m = _CODE_RE.search(reply)
    if m:
        return m.group(1).strip()
    reply = reply.strip()
    return reply if reply.startswith("{") and reply.endswith("}") else None


def _call_has_sanitised_args(node: ast.Call) -> bool:
    """
    Quick heuristic: return True iff *every* positional / keyword arg is either
    not user-controlled **or** directly wrapped in a recognised sanitiser.
    """

    def _sanitised(arg: ast.AST) -> bool:
        if isinstance(arg, ast.Call):
            fn = arg.func
            fn_name = (
                fn.attr
                if isinstance(fn, ast.Attribute)
                else fn.id if isinstance(fn, ast.Name) else ""
            )
            return _is_sanitizer(fn_name)
        return False

    args_iter = list(node.args) + [kw.value for kw in node.keywords]
    return all(_sanitised(a) for a in args_iter)


def _get_enclosing_function_source(call: ast.Call, source: str, tree: ast.AST) -> str:
    """
    Return the source text for the smallest `FunctionDef` that fully encloses
    *call*. Falls back to the call snippet when no enclosing function exists.
    """
    lineno = call.lineno
    best: Optional[ast.FunctionDef] = None

    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef):
            # Requires Python 3.8+ for end_lineno
            start, end = node.lineno, getattr(node, "end_lineno", node.lineno)
            if start <= lineno <= end:
                # Choose the *innermost* function (greatest start line)
                if best is None or start > best.lineno:
                    best = node

    target = best if best is not None else call
    return ast.get_source_segment(source, target) or "<code unavailable>"


_SEEN_FUNCS: set[tuple[str, int]] = set()


def _debug_ancestry(node):
    chain = []
    while node:
        chain.append(f"{type(node).__name__}:{getattr(node,'lineno', '?')}")
        node = getattr(node, "parent", None)
    return " -> ".join(chain)


def _is_trivial_wrapper(
    fn: ast.FunctionDef | ast.AsyncFunctionDef,
    *,
    src: str,
) -> bool:
    """
    True **only when** the *cleaned* source of *fn* (with doc-strings,
    comments and blank lines stripped) contains a **single** executable
    statement that is one of:

      • a bare call        →   ``foo(bar)``
      • ``return``         →   ``return``
      • ``return <Call>``  →   ``return do()``.

    Anything longer – even an extra ``print()`` – is considered
    non-trivial and must be audited.
    """
    try:
        # Use the same cleaner that feeds the LLM so comments/doc-strings
        # never influence the result.
        cleaned = _snippet(fn, src, max_lines=999_999)  # keep full body
        tree = ast.parse(cleaned)
        body = (
            tree.body[0].body
            if isinstance(tree.body[0], (ast.FunctionDef, ast.AsyncFunctionDef))
            else tree.body
        )

        if len(body) != 1:
            return False

        stmt = body[0]
        if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call):
            return True
        if isinstance(stmt, ast.Return):
            return stmt.value is None or isinstance(stmt.value, ast.Call)
        return False
    except Exception:
        # On any parsing hiccup fall back to *not* skipping
        return False


@register("AI_DETECTOR01", scope="node", node_types=(ast.Call,))
def llm_audit(unit):
    call = unit._current

    call_name = ""
    if isinstance(call.func, ast.Name):
        call_name = call.func.id
    elif isinstance(call.func, ast.Attribute):
        call_name = get_full_attr_name(call.func)

    if not call_name or not is_ai_call(call_name):
        return

    if getattr(_CLIENT, "is_dummy", False):
        logger.debug(
            "llm_code_audit: dummy client – skipping call at %s:%s",
            unit.path,
            call.lineno,
        )
        return

    logger.debug("llm_code_audit: call at %s:%s", unit.path, call.lineno)

    # climb to the nearest function *or* lambda
    func_node = call
    while func_node and not isinstance(
        func_node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)
    ):
        func_node = getattr(func_node, "parent", None)

    if func_node is None:  # top-level call → treat the Module as key
        key = (str(unit.path), "module", call.lineno)
    else:
        key = (str(unit.path), func_node.lineno)

    if key in _SEEN_FUNCS:
        logger.debug("Skipping duplicate call in scope %s", key)
        return
    _SEEN_FUNCS.add(key)

    logger.debug(
        "Visiting call at %s:%s - ancestry: %s",
        unit.path,
        call.lineno,
        _debug_ancestry(call),
    )

    # Short-circuit when everything is already routed through sanitisers
    if unit.is_user_tainted(call) and _call_has_sanitised_args(call):
        return

    if isinstance(
        func_node, (ast.FunctionDef, ast.AsyncFunctionDef)
    ) and _is_trivial_wrapper(func_node, src=unit.source):
        logger.debug("llm_code_audit: trivial wrapper – skipped")
        return

    # Obtain the module AST (attribute name differs across versions)
    tree = getattr(unit, "tree", getattr(unit, "_tree", None))
    if tree is None:
        logger.debug("llm_code_audit: AST missing – skipped")
        return

    func_src = _get_enclosing_function_source(call, unit.source, tree)
    flow_src = _path_context(unit, func_node)

    prompt = textwrap.dedent(
        f"""
        ## Context
        You are a security expert reviewing Python source code for **OWASP Top-10 for LLM Applications** risks:
        1. LLM01:2025 Prompt Injection
        2. LLM02:2025 Sensitive Information Disclosure
        3. LLM03:2025 Supply Chain
        4. LLM04:2025 Data and Model Poisoning
        5. LLM05:2025 Improper Output Handling
        6. LLM06:2025 Excessive Agency
        7. LLM07:2025 System Prompt Leakage
        8. LLM08:2025 Vector and Embedding Weaknesses
        9. LLM09:2025 Misinformation
        10. LLM10:2025 Unbounded Consumption

        You will receive:

        • **LOCAL FUNCTION**  – the code to audit
        • **CALL-FLOW CONTEXT** – snippets of its immediate callers / callees (for reference only)

        ### NON-NEGOTIABLE RULES
        1. If LOCAL FUNCTION is merely a thin wrapper (no extra logic), reply with `{{"issue": "clean"}}`.
        2. Report **only** vulnerabilities that are **inside LOCAL FUNCTION itself**.
        3. Ignore risks that exist *solely* in CALL-FLOW CONTEXT.

        ### TASK
        Return **exactly one line of JSON** with keys:<br>
        `"issue" · "sev" · "fix" · "owasp" · ("mitre" optional)`.<br>
        Use `"issue": "clean"` when no problem is present.

        ### LOCAL FUNCTION ###
        ```python
        {func_src}
        ```

        {flow_src}
    """
    ).strip()

    logger.debug(
        f"Detecting issues in {unit.path} {call.lineno} with LLM prompt:\n{prompt}\n\n"
    )

    try:
        reply = _CLIENT.ask(prompt, max_tokens=180)
    except Exception as exc:
        logger.error("llm_code_audit: provider error %s – skipped", exc)
        return

    payload = _json_fragment(reply)
    if not payload:
        logger.debug("llm_code_audit: non-JSON reply – skipped")
        return

    try:
        data = json.loads(payload)
    except json.JSONDecodeError:
        logger.debug("llm_code_audit: bad JSON – skipped")
        return

    issue = str(data.get("issue", "")).lower()
    if any(issue.startswith(p) for p in _IGNORE_PREFIXES):
        logger.debug("llm_code_audit: benign / clean – skipped")
        return

    dedup = (str(unit.path), call.lineno, str(data.get("owasp", "Axx")))
    if dedup in _EMITTED:
        return
    _EMITTED.add(dedup)

    logger.debug(
        "llm_code_audit: found %s (%s) in %s:%s",
        data.get("issue"),
        data.get("sev"),
        unit.path,
        call.lineno,
    )
    yield Finding(
        detector_id="AI_DETECTOR01",
        owasp_id=data.get("owasp", "Axx"),
        mitre=data.get("mitre", []),
        severity=str(data.get("sev", "info")).lower(),
        message=f"LLM audit: {data['issue']}",
        location=unit.path,
        line=call.lineno,
        fix=data.get("fix", ""),
    )
