"""CLI entry point."""

from __future__ import annotations

import asyncio
import json
from pathlib import Path

import typer
import yaml
from rich.console import Console
from rich.table import Table

from llm_pentest.analyzer import load_hooks_from_file
from llm_pentest.judge import LLMJudge
from llm_pentest.loader import PROFILES, load_chains, load_payloads, load_profile, load_target
from llm_pentest.models import SEVERITY_ORDER, AttackCategory, ScanResult
from llm_pentest.reporter import diff_results, write_html, write_json, write_sarif
from llm_pentest.scanner import run_scan

app = typer.Typer(name="llm-pentest", help="Security testing toolkit for LLM-based systems.")
console = Console()


@app.command()
def scan(
    config: Path | None = typer.Option(None, "--config", "-c", help="Target config YAML"),
    profile: str | None = typer.Option(None, "--profile", "-p", help="Target profile: openai, anthropic, ollama"),
    model: str | None = typer.Option(None, "--model", "-m", help="Override model name for profile"),
    category: str | None = typer.Option(None, "--category", help="Attack category filter"),
    payload_id: str | None = typer.Option(None, "--payload", help="Single payload ID"),
    payloads_dir: Path | None = typer.Option(None, "--payloads-dir", help="Custom payloads directory"),
    output: Path = typer.Option("results.json", "--output", "-o", help="Output file"),
    fmt: str = typer.Option("json", "--format", "-f", help="Output format: json, html, sarif"),
    dry_run: bool = typer.Option(False, "--dry-run", help="Show payloads without sending"),
    baseline: Path | None = typer.Option(None, "--baseline", "-b", help="Baseline results JSON for comparison"),
    fail_on: str = typer.Option("info", "--fail-on", help="Min severity to fail: info, low, medium, high, critical"),
    concurrency: int = typer.Option(1, "--concurrency", "-n", help="Parallel requests"),
    hooks: Path | None = typer.Option(None, "--hooks", help="Python file with detect_* functions"),
    judge_endpoint: str | None = typer.Option(None, "--judge-endpoint", help="LLM judge API endpoint"),
    judge_key: str | None = typer.Option(None, "--judge-key", envvar="LLM_JUDGE_API_KEY", help="LLM judge API key"),
    judge_model: str = typer.Option("gpt-4o-mini", "--judge-model", help="LLM judge model name"),
) -> None:
    """Run attacks against a target LLM system."""
    # Load target from config or profile
    if profile:
        target = load_profile(profile, model=model)
    elif config:
        target = load_target(config)
    else:
        console.print("[red]Provide --config or --profile[/]")
        raise typer.Exit(1)

    cat = AttackCategory(category) if category else None
    all_payloads = load_payloads(directory=payloads_dir, category=cat, payload_id=payload_id)
    chains = load_chains(directory=payloads_dir)

    if not all_payloads:
        console.print("[red]No payloads found.[/]")
        raise typer.Exit(1)

    if hooks:
        load_hooks_from_file(hooks)
        console.print(f"Hooks: [bold]{hooks}[/]")

    judge = None
    if judge_endpoint and judge_key:
        judge = LLMJudge(endpoint=judge_endpoint, api_key=judge_key, model=judge_model)
        console.print(f"Judge: [bold]{judge_model}[/] via {judge_endpoint}")

    console.print(f"Target: [bold]{target.name}[/] — {target.endpoint}")
    console.print(f"Payloads: [bold]{len(all_payloads)}[/] loaded" +
                  (f" + {len(chains)} chains" if chains else ""))
    if concurrency > 1:
        console.print(f"Concurrency: [bold]{concurrency}[/]")
    if dry_run:
        console.print("[yellow]DRY RUN — no requests will be sent[/]")

    result = asyncio.run(run_scan(
        target, all_payloads, dry_run=dry_run, judge=judge,
        concurrency=concurrency, chains=chains,
    ))

    # Baseline comparison
    diff = None
    if baseline and baseline.exists():
        base_result = ScanResult(**json.loads(baseline.read_text()))
        diff = diff_results(base_result, result)
        _print_diff(diff)

    _print_summary(result)

    # Write report
    payloads_by_id = {p.id: p for p in all_payloads}
    if fmt == "html":
        p = write_html(result, output.with_suffix(".html"), diff=diff)
    elif fmt == "sarif":
        p = write_sarif(result, output.with_suffix(".sarif"), payloads_by_id=payloads_by_id)
    else:
        p = write_json(result, output)
    console.print(f"\nReport saved to [bold]{p}[/]")

    # Fail based on severity threshold
    threshold = SEVERITY_ORDER.get(fail_on, 0)
    failing = [r for r in result.results
                if r.success and SEVERITY_ORDER.get(r.severity.value, 0) >= threshold]
    if failing:
        raise typer.Exit(1)


@app.command()
def payloads(
    category: str | None = typer.Option(None, "--category", help="Filter by category"),
    payloads_dir: Path | None = typer.Option(None, "--payloads-dir", help="Custom payloads directory"),
) -> None:
    """List available attack payloads."""
    cat = AttackCategory(category) if category else None
    items = load_payloads(directory=payloads_dir, category=cat)

    table = Table(title="Available Payloads")
    table.add_column("ID", style="cyan")
    table.add_column("Name")
    table.add_column("Category", style="magenta")
    table.add_column("Severity")
    table.add_column("OWASP")
    table.add_column("Prompts", justify="right")

    for p in items:
        sev_color = {"critical": "red", "high": "red", "medium": "yellow", "low": "green"}.get(
            p.severity.value, "white"
        )
        total = len(p.prompts) + len(p.turns)
        owasp = ", ".join(p.owasp) if p.owasp else ""
        table.add_row(p.id, p.name, p.category.value,
                      f"[{sev_color}]{p.severity.value}[/]", owasp, str(total))

    console.print(table)

    # Show chains
    chains = load_chains(directory=payloads_dir)
    if chains:
        console.print()
        ct = Table(title="Attack Chains")
        ct.add_column("ID", style="cyan")
        ct.add_column("Name")
        ct.add_column("Steps", justify="right")
        for c in chains:
            ct.add_row(c.id, c.name, str(len(c.steps)))
        console.print(ct)


@app.command()
def report(
    input_file: Path = typer.Option(..., "--input", "-i", help="Results JSON file"),
    fmt: str = typer.Option("html", "--format", "-f", help="Output format: json, html, sarif"),
    output: Path | None = typer.Option(None, "--output", "-o", help="Output file path"),
    baseline: Path | None = typer.Option(None, "--baseline", "-b", help="Baseline results JSON"),
    payloads_dir: Path | None = typer.Option(None, "--payloads-dir", help="Payloads dir (enriches SARIF)"),
) -> None:
    """Generate a report from scan results."""
    data = json.loads(input_file.read_text())
    result = ScanResult(**data)

    diff = None
    if baseline and baseline.exists():
        base_result = ScanResult(**json.loads(baseline.read_text()))
        diff = diff_results(base_result, result)
        _print_diff(diff)

    payloads_by_id = {p.id: p for p in load_payloads(directory=payloads_dir)}

    ext = {"html": ".html", "sarif": ".sarif"}.get(fmt, ".json")
    if output is None:
        output = input_file.with_suffix(ext)

    if fmt == "html":
        p = write_html(result, output, diff=diff)
    elif fmt == "sarif":
        p = write_sarif(result, output, payloads_by_id=payloads_by_id)
    else:
        p = write_json(result, output)
    console.print(f"Report saved to [bold]{p}[/]")


@app.command()
def diff(
    baseline_file: Path = typer.Option(..., "--baseline", "-b", help="Baseline results JSON"),
    current_file: Path = typer.Option(..., "--current", help="Current results JSON"),
) -> None:
    """Compare two scan results."""
    base = ScanResult(**json.loads(baseline_file.read_text()))
    current = ScanResult(**json.loads(current_file.read_text()))
    d = diff_results(base, current)
    _print_diff(d)


@app.command()
def init() -> None:
    """Interactive setup — generates a target config file."""
    console.print("[bold]LLM Pentest — Target Setup[/]\n")

    # Profile or custom?
    console.print("Available profiles: " + ", ".join(PROFILES.keys()))
    use_profile = typer.prompt("Use a profile? (name or 'custom')", default="custom")

    if use_profile in PROFILES:
        api_key_var = {"openai": "OPENAI_API_KEY", "anthropic": "ANTHROPIC_API_KEY"}.get(use_profile, "")
        model_name = typer.prompt("Model name", default=PROFILES[use_profile]["request"]["body_template"]["model"])
        target_name = typer.prompt("Target name", default=use_profile)

        data = {**PROFILES[use_profile], "name": target_name}
        data["request"]["body_template"]["model"] = model_name
        if api_key_var:
            console.print(f"\n[yellow]Set {api_key_var} in your environment before scanning.[/]")
    else:
        endpoint = typer.prompt("API endpoint URL")
        auth_type = typer.prompt("Auth type (none/bearer/header)", default="bearer")
        target_name = typer.prompt("Target name", default="my-target")

        data = {
            "name": target_name,
            "endpoint": endpoint,
            "auth": {"type": auth_type},
            "request": {
                "method": "POST",
                "headers": {"Content-Type": "application/json"},
                "body_template": {"messages": "{{messages}}"},
            },
            "response_path": "choices[0].message.content",
        }
        if auth_type == "bearer":
            data["auth"]["token"] = "${API_KEY}"
            console.print("\n[yellow]Set API_KEY in your environment before scanning.[/]")

    out_path = Path(f"{target_name}.yaml")
    out_path.write_text(yaml.dump(data, default_flow_style=False, sort_keys=False))
    console.print(f"\n[green]Config saved to {out_path}[/]")
    console.print(f"Run: [bold]llm-pentest scan -c {out_path}[/]")


@app.command()
def profiles() -> None:
    """List available target profiles."""
    table = Table(title="Target Profiles")
    table.add_column("Name", style="cyan")
    table.add_column("Endpoint")
    table.add_column("Default Model")

    for name, cfg in PROFILES.items():
        table.add_row(name, cfg["endpoint"], cfg["request"]["body_template"].get("model", ""))

    console.print(table)


def _print_diff(d: dict) -> None:
    console.print()
    console.print("[bold]Baseline Comparison[/]")
    console.print(f"  Previous: {d['baseline_vulns']} vulns ({d['baseline_timestamp'][:19]})")
    console.print(f"  Current:  {d['current_vulns']} vulns ({d['current_timestamp'][:19]})")
    console.print(f"  [green]Fixed:      {d['fixed']}[/]")
    console.print(f"  [red]New:        {d['new']}[/]")
    console.print(f"  Persistent: {d['persistent']}")

    for item in d["fixed_details"]:
        console.print(f"    [green]✅ {item}[/]")
    for item in d["new_details"]:
        console.print(f"    [red]🆕 {item}[/]")


def _print_summary(result: ScanResult) -> None:
    console.print()
    console.print("[bold]Scan Complete[/]")
    console.print(f"  Total attacks:  {result.total_attacks}")
    console.print(f"  Successful:     [red]{result.successful_attacks}[/]")
    console.print(f"  Blocked:        [green]{result.total_attacks - result.successful_attacks}[/]")

    # OWASP summary
    owasp_hits: dict[str, int] = {}
    for r in result.results:
        if r.success:
            for o in r.owasp:
                owasp_hits[o] = owasp_hits.get(o, 0) + 1
    if owasp_hits:
        console.print("  OWASP LLM Top 10:")
        for code, count in sorted(owasp_hits.items()):
            label = OWASP_LABELS.get(code, code)
            console.print(f"    [red]{code}[/] {label}: {count} hits")

    for r in result.results:
        if r.success:
            console.print(f"  [red]⚠ {r.payload_id}[/] [{r.severity.value}] confidence {r.confidence:.0%} — {r.prompt[:70]}")


OWASP_LABELS = {
    "LLM01": "Prompt Injection",
    "LLM02": "Insecure Output Handling",
    "LLM03": "Training Data Poisoning",
    "LLM04": "Model Denial of Service",
    "LLM05": "Supply Chain Vulnerabilities",
    "LLM06": "Sensitive Information Disclosure",
    "LLM07": "Insecure Plugin Design",
    "LLM08": "Excessive Agency",
    "LLM09": "Overreliance",
    "LLM10": "Model Theft",
}


if __name__ == "__main__":
    app()
