"""LLM-based code review with focused specialist prompts.

Runs 6 focused prompts in parallel via OpenRouter, each targeting a specific
category of issues. One additional sweep prompt catches anything the specialists
miss. All calls run concurrently — wall time ≈ slowest single call (~5-10s).

Cost: ~$0.015-0.02 per file via Gemini 2.5 Flash.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import re
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

REVIEW_MODEL = os.getenv("HIKAFLOW_REVIEW_MODEL", "google/gemini-2.5-flash")

_FINDING_RE = re.compile(
    r"(?:LINE\s+)?(\d+)\s*:\s*\[?(\w{2,6})\]?\s*[:\-|]?\s*(.+)",
    re.IGNORECASE,
)


def _normalize_severity(raw: str) -> str | None:
    """Normalize mangled severity tags. Returns None if not a severity."""
    s = raw.upper().rstrip("|:- ")
    if s.startswith("HI") or s in ("HH",):
        return "HIGH"
    if s.startswith("ME") or s in ("MM",):
        return "MEDIUM"
    if s.startswith("LO"):
        return "LOW"
    return None


def _get_api_key() -> str:
    key = os.environ.get("OPENROUTER_API_KEY", "")
    if not key:
        try:
            creds_path = Path.home() / ".hikaflow" / "credentials.json"
            if creds_path.exists():
                creds = json.loads(creds_path.read_text())
                key = creds.get("openrouter_api_key", "")
        except Exception:
            pass
    return key


# ---------------------------------------------------------------------------
# Specialist prompts
# ---------------------------------------------------------------------------

@dataclass
class ASTFlags:
    """Feature flags from AST analysis for routing LLM specialists."""
    has_async: bool = False
    has_subprocess: bool = False
    has_file_ops: bool = False
    has_auth_patterns: bool = False
    has_except_handlers: bool = False
    has_globals: bool = False

    @classmethod
    def from_source(cls, source: str) -> "ASTFlags":
        """Quick feature detection without full AST parse."""
        s = source.lower()
        return cls(
            has_async="async def " in source,
            has_subprocess="subprocess" in s or any(
                kw in s for kw in ("shell", "command", "exec(", "eval(", "os.system", "os.popen", "shlex")
            ),
            has_file_ops="open(" in s or "os.path" in s or "pathlib" in s,
            has_auth_patterns=any(kw in s for kw in ("password", "auth", "token", "permission", "policy", "approval")),
            has_except_handlers="except " in source,
            has_globals=bool(re.search(r'^[A-Z_]{2,}\s*=', source, re.MULTILINE)),
        )


def _build_cross_file_context(file_path: str) -> str:
    """Deterministically build cross-file context: who imports/calls this module.

    Scans sibling source files for imports and call sites referencing this module.
    Returns a short context block for injection into specialist prompts.
    """
    import ast as _ast
    from pathlib import Path

    target = Path(file_path)
    module_name = target.stem  # e.g. "sessions" from "api/sessions.py"
    parent_dir = target.parent

    if not parent_dir.exists():
        return ""

    callers: list[str] = []
    # Scan siblings + parent for imports of this module
    search_dirs = [parent_dir]
    if parent_dir.parent.exists() and parent_dir.parent != parent_dir:
        search_dirs.append(parent_dir.parent)

    for sdir in search_dirs:
        for py_file in sdir.glob("*.py"):
            if py_file.name == target.name or py_file.name.startswith("test_"):
                continue
            try:
                tree = _ast.parse(py_file.read_text(errors="replace"))
            except Exception:
                continue
            imports_target = False
            for node in _ast.walk(tree):
                if isinstance(node, _ast.ImportFrom) and node.module and module_name in node.module:
                    imported_names = [a.name for a in node.names]
                    imports_target = True
                    callers.append(f"  - `{py_file.name}` imports: {', '.join(imported_names)}")
                    break
                elif isinstance(node, _ast.ImportFrom) and node.names:
                    # Check if module_name is in the imported names (e.g. "from api import sessions")
                    for alias in node.names:
                        if alias.name == module_name:
                            imports_target = True
                            callers.append(f"  - `{py_file.name}` imports `{module_name}` (uses as module)")
                            break
                    if imports_target:
                        break
                elif isinstance(node, _ast.Import):
                    for alias in node.names:
                        if module_name in alias.name:
                            imports_target = True
                            callers.append(f"  - `{py_file.name}` imports `{alias.name}`")
                            break
            if len(callers) >= 8:
                break
        if len(callers) >= 8:
            break

    if not callers:
        return ""

    return (
        f"\nCROSS-FILE CONTEXT (who uses this module):\n"
        + "\n".join(callers[:8])
        + "\nConsider how callers pass data into these functions when assessing severity.\n\n"
    )


def _build_prompts(source: str, file_path: str, flags: ASTFlags) -> dict[str, str]:
    """Build specialist prompts with line-numbered source and AST-guided routing."""
    # Line-numbered source — LLM sees exact line numbers
    numbered = "\n".join(f"{i+1}: {line}" for i, line in enumerate(source.splitlines()))
    code_block = f"```\n{numbered}\n```"

    # Cross-file context — deterministically extracted
    cross_ctx = _build_cross_file_context(file_path)

    fmt = (
        "RULES:\n"
        "- Use the EXACT line numbers shown in the code. Do NOT guess or approximate.\n"
        "- Only report bugs that cause incorrect runtime behavior, security vulnerabilities, or crashes.\n"
        "- Do NOT report style issues, missing docs, or theoretical concerns.\n"
        "- Do NOT report the same issue twice.\n"
        "- Maximum 5 findings. Quality over quantity.\n\n"
        "Format each finding as:\n"
        "LINE <num>: [HIGH/MEDIUM/LOW] <one-sentence description of the bug and its impact>\n"
        'If no real issues, say "No issues found."'
    )

    # All possible specialists — routing decides which actually run
    all_specialists = {
        "security_injection": (
            f"You are a security specialist focused ONLY on injection vulnerabilities.\n"
            f"Review this code for:\n"
            f"- Shell injection via unsanitized input in commands\n"
            f"- SQL injection, path traversal, SSRF\n"
            f"- Bypasses of safety checks (quoting, escaping, operator splitting)\n"
            f"- Ways to smuggle dangerous operations past filters\n\n"
            f"Only report issues where you can describe a concrete attack vector.\n\n"
            f"File: `{file_path}`\n{cross_ctx}{code_block}\n\n{fmt}"
        ),
        "security_auth": (
            f"You are a security specialist focused ONLY on authorization bypass.\n"
            f"Review this code for:\n"
            f"- Ways to escalate from a restricted mode to a permissive one\n"
            f"- Cache poisoning that grants broader approval than intended\n"
            f"- Operations that bypass policy checks entirely\n"
            f"- Modes that accidentally allow dangerous operations\n"
            f"- Default-allow patterns that miss new/unknown operations\n\n"
            f"Only report issues where you can describe how the bypass works.\n\n"
            f"File: `{file_path}`\n{cross_ctx}{code_block}\n\n{fmt}"
        ),
        "logic_bugs": (
            f"You are a logic bug specialist focused ONLY on correctness issues.\n"
            f"Review this code for:\n"
            f"- Edge cases that produce wrong results (empty input, None, boundary values)\n"
            f"- Conditions that don't match their intent\n"
            f"- Functions that can crash on valid input\n"
            f"- Regex patterns that match too much or too little\n"
            f"- Off-by-one errors, wrong operator, inverted logic\n\n"
            f"Only report bugs where the code produces incorrect behavior.\n\n"
            f"File: `{file_path}`\n{cross_ctx}{code_block}\n\n{fmt}"
        ),
        "error_handling": (
            f"You are an error handling specialist focused ONLY on failure modes.\n"
            f"Review this code for:\n"
            f"- Exceptions that crash with unhelpful errors\n"
            f"- Missing None/empty checks before attribute access\n"
            f"- Swallowed exceptions that hide real failures\n"
            f"- Error paths that leave the system in an inconsistent state\n"
            f"- Operations that silently return wrong results on failure\n\n"
            f"Only report real failure scenarios, not theoretical ones.\n\n"
            f"File: `{file_path}`\n{cross_ctx}{code_block}\n\n{fmt}"
        ),
        "design_flaws": (
            f"You are a design review specialist.\n"
            f"Review this code for:\n"
            f"- Security model gaps (operations that should require approval but don't)\n"
            f"- Missing entries in allowlists/denylists\n"
            f"- Policy decisions inconsistent between modes\n"
            f"- Semantics that could lead to unexpected behavior over time\n\n"
            f"Only report issues that affect real-world security or correctness.\n\n"
            f"File: `{file_path}`\n{cross_ctx}{code_block}\n\n{fmt}"
        ),
        "concurrency": (
            f"You are a concurrency and async specialist.\n"
            f"Review this code for:\n"
            f"- Race conditions in shared state\n"
            f"- Missing awaits on coroutines\n"
            f"- Thread-safety issues with mutable globals\n"
            f"- Deadlock or livelock potential\n"
            f"- Resource leaks (unclosed files, connections, locks)\n\n"
            f"Only report issues where concurrent access causes real bugs.\n\n"
            f"File: `{file_path}`\n{cross_ctx}{code_block}\n\n{fmt}"
        ),
    }

    # AST-guided routing — skip irrelevant specialists
    selected: dict[str, str] = {}

    # Always run: logic_bugs (every file can have logic bugs)
    selected["logic_bugs"] = all_specialists["logic_bugs"]

    # Run if file has subprocess/shell/os calls
    if flags.has_subprocess or flags.has_file_ops:
        selected["security_injection"] = all_specialists["security_injection"]

    # Run if file has auth/permission/policy patterns
    if flags.has_auth_patterns:
        selected["security_auth"] = all_specialists["security_auth"]
        selected["design_flaws"] = all_specialists["design_flaws"]

    # Run if file has exception handlers
    if flags.has_except_handlers:
        selected["error_handling"] = all_specialists["error_handling"]

    # Run if file has async or global mutable state
    if flags.has_async or flags.has_globals:
        selected["concurrency"] = all_specialists["concurrency"]

    # Always run design_flaws if not already selected (catches general design issues)
    if "design_flaws" not in selected:
        selected["design_flaws"] = all_specialists["design_flaws"]

    return selected


# ---------------------------------------------------------------------------
# LLM calling
# ---------------------------------------------------------------------------

@dataclass
class LLMFinding:
    line: int
    severity: str  # HIGH, MEDIUM, LOW
    message: str
    source: str    # which specialist found it


async def _call_openrouter(
    prompt: str,
    label: str,
    api_key: str,
    model: str = REVIEW_MODEL,
    timeout: float = 30.0,
) -> tuple[list[LLMFinding], float]:
    """Make one OpenRouter call, parse findings from response."""
    import httpx

    t0 = time.time()
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.post(
                "https://openrouter.ai/api/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": model,
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": 1500,
                    "temperature": 0.1,
                },
            )
            data = resp.json()
    except Exception as e:
        logger.warning("LLM review call failed (%s): %s", label, e)
        return [], time.time() - t0

    elapsed = time.time() - t0

    if "error" in data:
        logger.warning("LLM review error (%s): %s", label, data["error"])
        return [], elapsed

    try:
        content = data["choices"][0]["message"]["content"]
    except (KeyError, IndexError):
        return [], elapsed

    # Parse findings
    findings: list[LLMFinding] = []
    for line_text in content.splitlines():
        line_text = line_text.strip()
        if not line_text:
            continue
        m = _FINDING_RE.match(line_text)
        if m:
            line_num = int(m.group(1))
            severity = _normalize_severity(m.group(2))
            if severity is None:
                continue  # matched a word that isn't a severity
            message = m.group(3).strip().rstrip(".")
            if line_num < 1:
                continue
            findings.append(LLMFinding(
                line=line_num, severity=severity,
                message=message, source=label,
            ))

    return findings, elapsed


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

async def run_llm_review(
    source: str,
    file_path: str,
) -> tuple[list[LLMFinding], dict[str, float]]:
    """Run all LLM specialist prompts in parallel.

    Returns (findings, timing) where timing maps specialist name to seconds.
    """
    api_key = _get_api_key()
    if not api_key:
        logger.warning("No OpenRouter API key — skipping LLM review")
        return [], {}

    flags = ASTFlags.from_source(source)
    prompts = _build_prompts(source, file_path, flags)
    timing: dict[str, float] = {}

    async def _run(name: str, prompt: str) -> list[LLMFinding]:
        findings, elapsed = await _call_openrouter(prompt, name, api_key)
        timing[name] = elapsed
        return findings

    results = await asyncio.gather(*[_run(n, p) for n, p in prompts.items()])

    all_findings: list[LLMFinding] = []
    for finding_list in results:
        all_findings.extend(finding_list)

    if not all_findings:
        return all_findings, timing

    # Pre-dedup: collapse identical patterns before verification
    total_lines = source.count("\n") + 1
    all_findings = _pre_dedup(all_findings, total_lines)

    # Verification pass — filter false positives and merge duplicates
    verified, v_elapsed = await _verify_findings(all_findings, source, file_path, api_key)
    timing["_verify"] = v_elapsed

    return verified, timing


def _finding_core(msg: str) -> str:
    """Normalize a finding message for dedup — strip names, quotes, take first 50 chars."""
    s = msg[:80].lower()
    s = re.sub(r'`[^`]*`', '``', s)
    s = re.sub(r'"[^"]*"', '""', s)
    s = re.sub(r'\b(the\s+)\w+(\s+endpoint|\s+function|\s+method)', r'\1_\2', s)
    s = re.sub(r'(in\s+)\w+(\s*,|\s+mode|\s+the)', r'\1_\2', s)
    s = re.sub(r'line\s+\d+', 'line N', s)
    return s[:50].strip()


def _pre_dedup(findings: list[LLMFinding], total_lines: int) -> list[LLMFinding]:
    """Fast pre-dedup before the LLM verifier. Caps per-pattern and removes out-of-range."""
    seen: dict[str, int] = {}
    proximity: list[tuple[int, str]] = []
    result: list[LLMFinding] = []

    for f in findings:
        if f.line < 1 or f.line > total_lines:
            continue
        core = _finding_core(f.message)
        # Cap at 2 per pattern
        count = seen.get(core, 0)
        if count >= 2:
            continue
        # Proximity: same core within 15 lines
        near = any(abs(pl - f.line) <= 15 and pc == core for pl, pc in proximity)
        if near:
            continue
        seen[core] = count + 1
        proximity.append((f.line, core))
        result.append(f)

    return result


async def _verify_findings(
    findings: list[LLMFinding],
    source: str,
    file_path: str,
    api_key: str,
) -> tuple[list[LLMFinding], float]:
    """Verify findings against actual code. Removes false positives and duplicates.

    One LLM call: sends each finding + the actual code at the cited line,
    asks "is this real? is it a duplicate?"
    """
    lines = source.splitlines()
    total_lines = len(lines)

    # Build context for each finding: the cited line ± 5 lines
    finding_entries: list[str] = []
    for i, f in enumerate(findings):
        line_idx = f.line - 1
        if line_idx < 0 or line_idx >= total_lines:
            ctx = "(line out of range)"
        else:
            start = max(0, line_idx - 3)
            end = min(total_lines, line_idx + 4)
            ctx_lines = []
            for j in range(start, end):
                marker = ">>>" if j == line_idx else "   "
                ctx_lines.append(f"{marker} {j+1}: {lines[j]}")
            ctx = "\n".join(ctx_lines)
        finding_entries.append(
            f"FINDING {i+1}: LINE {f.line} [{f.severity}] ({f.source})\n"
            f"  Claim: {f.message}\n"
            f"  Code:\n{ctx}\n"
        )

    prompt = (
        f"You are a code review verifier for `{file_path}`. Below are findings from multiple "
        f"specialists. For each finding, I've included the actual code at the cited line.\n\n"
        f"Your job:\n"
        f"1. REJECT findings where the claim is WRONG — the code actually handles the case, "
        f"or the described bug cannot happen. Check the code carefully.\n"
        f"2. REJECT DUPLICATES aggressively — if multiple findings describe the SAME underlying "
        f"pattern/bug repeated across different functions or endpoints, keep ONLY the single "
        f"best example. For instance, if 5 endpoints all have the same fallback pattern, "
        f"that's ONE finding, not 5.\n"
        f"3. KEEP findings that describe real, distinct bugs.\n\n"
        + "\n".join(finding_entries)
        + "\n\nRespond with ONLY the finding numbers to KEEP, one per line:\n"
        f"KEEP <number>\n"
        f"Example:\nKEEP 1\nKEEP 3\nKEEP 7\n"
    )

    result_findings, elapsed = await _call_openrouter(
        prompt, "_verify", api_key, timeout=20.0,
    )

    # Parse KEEP responses
    keep_re = re.compile(r"KEEP\s+(\d+)", re.IGNORECASE)
    keep_indices: set[int] = set()
    # _call_openrouter returns LLMFinding objects, but here we need raw text
    # So call directly
    import httpx
    t0 = time.time()
    try:
        async with httpx.AsyncClient(timeout=45.0) as client:
            resp = await client.post(
                "https://openrouter.ai/api/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": REVIEW_MODEL,
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": 500,
                    "temperature": 0.0,
                },
            )
            data = resp.json()
            elapsed = time.time() - t0
            content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
    except Exception:
        # On failure, return all findings unfiltered
        return findings, time.time() - t0

    for line_text in content.splitlines():
        m = keep_re.search(line_text)
        if m:
            idx = int(m.group(1)) - 1  # 1-indexed → 0-indexed
            if 0 <= idx < len(findings):
                keep_indices.add(idx)

    # If verifier kept nothing (possible LLM failure), return all
    if not keep_indices:
        return findings, elapsed

    verified = [f for i, f in enumerate(findings) if i in keep_indices]
    return verified, elapsed
