"""Pipeline orchestrator — scan phase and fix phase are separate commands."""

from __future__ import annotations

import json
import logging
import os
from pathlib import Path

from rich.console import Console

from secagent.agents.fixer_agent import generate_fix
from secagent.agents.planner_agent import plan_scan
from secagent.agents.validator_agent import validate_fix
from secagent.agents.verifier_agent import verify_vulnerability
from secagent.llm.grok_client import GrokClient, _DEFAULT_FAST_MODEL, _DEFAULT_MODEL
from secagent.models import (
    PipelineResult,
    ScanResult,
    Severity,
    SkippedFix,
    TokenUsage,
    Vulnerability,
    VerificationStatus,
)
from secagent.parsers.ai_parser import AIParser
from secagent.parsers.builtin_parser import BuiltinParser
from secagent.parsers.gitleaks_parser import GitleaksParser
from secagent.report.fix_report import generate_report
from secagent.report.vuln_report import generate_vuln_report
from secagent.scanners.ai_scanner import AIScanner
from secagent.scanners.builtin_scanner import BuiltinSecretScanner
from secagent.scanners.gitleaks_scanner import GitleaksScanner
from secagent.utils.patch_applier import apply_patch

logger = logging.getLogger(__name__)
console = Console()

# Registry — easy to extend with new scanners
_SCANNER_MAP = {
    "ai-scanner": (AIScanner, AIParser),
    "builtin-secrets": (BuiltinSecretScanner, BuiltinParser),
    "gitleaks": (GitleaksScanner, GitleaksParser),
}


def _make_clients() -> tuple[GrokClient, GrokClient]:
    """Return (reasoning_client, fast_client) for the pipeline.

    reasoning_client — grok-4-1-fast-reasoning (ai-scanner, fixer, validator)
    fast_client      — grok-4-1-fast-non-reasoning (planner, verifier)

    Both models can be overridden via env vars:
      GROK_MODEL       → reasoning model
      GROK_FAST_MODEL  → fast non-reasoning model
    """
    reasoning_model = os.getenv("GROK_MODEL", _DEFAULT_MODEL)
    fast_model = os.getenv("GROK_FAST_MODEL", _DEFAULT_FAST_MODEL)
    return GrokClient(model=reasoning_model), GrokClient(model=fast_model)


# ── Vulnerability persistence helpers ────────────────────────────────────────

def _vulns_file(target: Path) -> Path:
    """Return path to the vulnerability cache file for *target*."""
    if target.is_file():
        base_dir = target.parent
        name = target.name
    else:
        base_dir = target
        name = target.name
    data_dir = base_dir / ".secagent" / "data"
    data_dir.mkdir(parents=True, exist_ok=True)
    return data_dir / f"{name}.vulns.json"


def _save_vulns(target: Path, vulns: list[Vulnerability]) -> Path:
    """Serialize *vulns* to a JSON cache file next to *target*."""
    path = _vulns_file(target)
    data = {
        "target": str(target),
        "vulnerabilities": [
            {
                "scanner": v.scanner,
                "file": v.file,
                "line": v.line,
                "severity": v.severity.value,
                "vuln_type": v.vuln_type,
                "description": v.description,
                "raw": v.raw,
                "verification": v.verification.value if v.verification else None,
            }
            for v in vulns
        ],
    }
    path.write_text(json.dumps(data, indent=2), encoding="utf-8")
    return path


def _load_vulns(target: Path) -> list[Vulnerability] | None:
    """Load cached vulnerabilities for *target*, or return None if not found."""
    path = _vulns_file(target)
    if not path.exists():
        return None
    data = json.loads(path.read_text(encoding="utf-8"))
    vulns = []
    for v in data["vulnerabilities"]:
        vuln = Vulnerability(
            scanner=v["scanner"],
            file=v["file"],
            line=v["line"],
            severity=Severity(v["severity"]),
            vuln_type=v["vuln_type"],
            description=v["description"],
            raw=v.get("raw", {}),
            verification=VerificationStatus(v["verification"]) if v.get("verification") else None,
        )
        vulns.append(vuln)
    return vulns


# ── Shared scan stages (1-4) ──────────────────────────────────────────────────

def _run_scan_stages(
    target: Path,
    reasoning_client: GrokClient,
    fast_client: GrokClient,
) -> tuple[list[Vulnerability], dict[str, TokenUsage], dict[str, str]]:
    """Run stages 1-4: planning, scanning, parsing, deduplication.

    Returns:
        (unique_vulns, stage_usage, stage_models)
    """
    stage_usage: dict[str, TokenUsage] = {}
    stage_models: dict[str, str] = {}

    # ── 1. Planning — fast non-reasoning model ───────────────────────
    console.print("\n[bold cyan]⚙  Planning scan…[/]")
    fast_client.reset_usage()
    scanner_names = plan_scan(target, fast_client)
    stage_usage["planner"] = fast_client.reset_usage()
    stage_models["planner"] = fast_client.model
    console.print(f"   Scanners selected: [green]{', '.join(scanner_names)}[/]\n")

    # ── 2 & 3. Scanning + Parsing ─────────────────────────────────────
    console.print("[bold cyan]🔍 Running scanners…[/]")
    all_vulns: list[Vulnerability] = []
    scanner_usage = TokenUsage()

    for name in scanner_names:
        if name not in _SCANNER_MAP:
            logger.warning("Unknown scanner '%s' — skipping.", name)
            continue

        scanner_cls, parser_cls = _SCANNER_MAP[name]
        # AIScanner uses the reasoning client; other scanners have no LLM calls
        scanner = scanner_cls(client=reasoning_client) if name == "ai-scanner" else scanner_cls()
        parser = parser_cls()

        console.print(f"   [{name}] ", end="")
        reasoning_client.reset_usage()
        raw = scanner.safe_run(target)
        scanner_usage.merge(reasoning_client.reset_usage())
        vulns = parser.parse(raw)
        all_vulns.extend(vulns)
        console.print(f"[green]{len(vulns)} finding(s)[/]")

    stage_usage["scanner"] = scanner_usage
    stage_models["scanner"] = reasoning_client.model

    if not all_vulns:
        return [], stage_usage, stage_models

    # ── 4. Deduplicate & merge ─────────────────────────────────────────
    seen: dict[tuple[str, int], Vulnerability] = {}
    scanner_sets: dict[tuple[str, int], set[str]] = {}
    dep_vulns: list[Vulnerability] = []

    for vuln in all_vulns:
        if vuln.line <= 0:
            dep_vulns.append(vuln)
            continue

        norm_file = str(Path(vuln.file).resolve())
        key = (norm_file, vuln.line)
        if key in seen:
            existing = seen[key]
            existing.description = existing.description + " | " + vuln.description
            scanner_sets[key].add(vuln.scanner)
            existing.scanner = "+".join(sorted(scanner_sets[key]))
            console.print(f"   [dim]Merged duplicate: {vuln.id} → {existing.id}[/]")
        else:
            seen[key] = vuln
            scanner_sets[key] = {vuln.scanner}

    unique_vulns = list(seen.values()) + dep_vulns
    console.print(
        f"   [dim]{len(all_vulns)} raw finding(s) → {len(unique_vulns)} unique[/]\n"
    )
    return unique_vulns, stage_usage, stage_models


# ── Public API ────────────────────────────────────────────────────────────────

def run_scan(target: Path) -> ScanResult:
    """Run stages 1-4 only: plan, scan, parse, deduplicate.

    Saves results to ``.secagent/<target>.vulns.json`` so ``run_fix`` can
    pick them up later without re-scanning.
    """
    from secagent.report.cost_report import update_cost_report

    reasoning_client, fast_client = _make_clients()
    result = ScanResult(target=str(target))

    unique_vulns, stage_usage, stage_models = _run_scan_stages(target, reasoning_client, fast_client)
    result.total_findings = len(unique_vulns)
    result.vulnerabilities = unique_vulns
    result.stage_usage = stage_usage
    result.stage_models = stage_models

    secagent_dir = _vulns_file(target).parent.parent  # .secagent/

    if not unique_vulns:
        console.print("\n[bold green]✅ No vulnerabilities found![/]\n")
    else:
        scan_file = _save_vulns(target, unique_vulns)
        result.scan_file = str(scan_file)
        console.print(f"[bold green]✅ Found {len(unique_vulns)} vulnerability(s). Results saved to:[/] {scan_file}\n")
        console.print("[dim]Run [bold]secagent fix <path>[/] to verify and fix them.[/]\n")

    generate_vuln_report(unique_vulns, str(target), secagent_dir)
    update_cost_report(stage_usage, stage_models, "scan", str(target), secagent_dir)
    return result


def run_fix(target: Path) -> PipelineResult:
    """Run the fix pipeline (stages 5-9) against *target*.

    If a ``.secagent/<target>.vulns.json`` cache exists it is loaded directly.
    Otherwise a fresh scan (stages 1-4) is run first.
    """
    from secagent.report.cost_report import update_cost_report

    reasoning_client, fast_client = _make_clients()
    result = PipelineResult(target=str(target))
    stage_usage: dict[str, TokenUsage] = {}
    stage_models: dict[str, str] = {}

    # Try to load cached scan results
    cached = _load_vulns(target)
    if cached is not None:
        console.print(f"\n[bold cyan]📂 Loading scan results from cache…[/]")
        unique_vulns = cached
        console.print(f"   Loaded [green]{len(unique_vulns)}[/] finding(s)\n")
    else:
        console.print("\n[bold yellow]⚠  No scan results found — running scan first…[/]")
        unique_vulns, scan_stage_usage, scan_stage_models = _run_scan_stages(
            target, reasoning_client, fast_client
        )
        stage_usage.update(scan_stage_usage)
        stage_models.update(scan_stage_models)
        if unique_vulns:
            _save_vulns(target, unique_vulns)

    result.total_findings = len(unique_vulns)
    result.vulnerabilities = unique_vulns

    secagent_dir = _vulns_file(target).parent.parent  # .secagent/

    if not unique_vulns:
        console.print("\n[bold green]✅ No vulnerabilities found![/]\n")
        report_path = generate_report(result, output_dir=secagent_dir)
        result.report_path = str(report_path)
        result.stage_usage = stage_usage
        result.stage_models = stage_models
        update_cost_report(stage_usage, stage_models, "fix", str(target), secagent_dir)
        return result

    # ── 5. Verification — fast non-reasoning model ───────────────────
    console.print(f"[bold cyan]🔎 Verifying {len(unique_vulns)} finding(s)…[/]")
    tp_vulns: list[Vulnerability] = []
    verifier_usage = TokenUsage()

    for vuln in unique_vulns:
        fast_client.reset_usage()
        status = verify_vulnerability(vuln, fast_client)
        verifier_usage.merge(fast_client.reset_usage())
        vuln.verification = status
        if status == VerificationStatus.TP:
            tp_vulns.append(vuln)
            result.true_positives += 1
        else:
            result.false_positives += 1
        console.print(
            f"   {vuln.id} → [{'red' if status == VerificationStatus.TP else 'yellow'}]{status.value}[/]"
        )

    stage_usage["verifier"] = verifier_usage
    stage_models["verifier"] = fast_client.model

    if not tp_vulns:
        console.print("\n[bold green]✅ All findings are false positives![/]\n")
        report_path = generate_report(result, output_dir=secagent_dir)
        result.report_path = str(report_path)
        result.stage_usage = stage_usage
        result.stage_models = stage_models
        update_cost_report(stage_usage, stage_models, "fix", str(target), secagent_dir)
        return result

    # ── 6+7. Fixing + Validation — reasoning model ───────────────────
    console.print(f"\n[bold cyan]🛠  Generating fixes for {len(tp_vulns)} issue(s)…[/]")
    fixer_usage = TokenUsage()
    validator_usage = TokenUsage()

    for vuln in tp_vulns:
        # Fixer — reasoning
        reasoning_client.reset_usage()
        fix = generate_fix(vuln, reasoning_client)
        fixer_usage.merge(reasoning_client.reset_usage())

        if fix is None:
            skip_reason = "no line info" if vuln.line <= 0 else "no fix generated"
            console.print(f"   {vuln.id} → [yellow]skipped ({skip_reason})[/]")
            result.skipped_fixes.append(SkippedFix(vulnerability=vuln, reason=skip_reason))
            continue

        # Validator — reasoning
        reasoning_client.reset_usage()
        validation = validate_fix(fix, reasoning_client)
        validator_usage.merge(reasoning_client.reset_usage())
        result.validations.append(validation)

        if not validation.valid:
            console.print(f"   {vuln.id} → [red]fix invalid: {validation.reason}[/]")
            result.skipped_fixes.append(SkippedFix(vulnerability=vuln, reason=f"validation failed: {validation.reason}"))
            continue

        # ── 8. Apply patch ───────────────────────────────────────────
        applied = apply_patch(fix.file, fix.original_snippet, fix.patched_snippet, create_backup=False)
        if applied:
            result.fixes.append(fix)
            result.fixes_applied += 1
            result.fixes_validated += 1
            console.print(f"   {vuln.id} → [green]fixed ✓[/]")
        else:
            console.print(f"   {vuln.id} → [red]patch failed[/]")
            result.skipped_fixes.append(SkippedFix(vulnerability=vuln, reason="patch failed — snippet not found in file"))

    stage_usage["fixer"] = fixer_usage
    stage_usage["validator"] = validator_usage
    stage_models["fixer"] = reasoning_client.model
    stage_models["validator"] = reasoning_client.model
    result.stage_usage = stage_usage
    result.stage_models = stage_models

    # ── 9. Report ────────────────────────────────────────────────────
    console.print("\n[bold cyan]📄 Generating report…[/]")
    generate_vuln_report(result.vulnerabilities, str(target), secagent_dir)
    report_path = generate_report(result, output_dir=secagent_dir)
    result.report_path = str(report_path)
    console.print(f"   Report: [underline]{report_path}[/]\n")

    cost_path = update_cost_report(stage_usage, stage_models, "fix", str(target), secagent_dir)
    console.print(f"   Cost log: [underline]{cost_path}[/]\n")

    return result
