"""
cli.py — SichGate command-line interface.

Three commands the customer ever runs:
  sichgate run      — execute attacks, output JSON
  sichgate activate — one-time license activation (Phase 2)
  sichgate start    — launch localhost dashboard (Phase 2)
"""
from __future__ import annotations

import json
import sys
import logging
from pathlib import Path
from typing import Optional

import typer
from rich.console import Console
from rich.table import Table
from rich import print as rprint

app = typer.Typer(
    name="sichgate",
    help="AI red teaming — independent adversarial validation.",
    add_completion=False,
)
console = Console()

# Register compliance sub-commands (Gaps #1-4)
from sichgate.compliance_cli import compliance_app  # noqa: E402
app.add_typer(compliance_app, name="compliance")

logging.basicConfig(level=logging.WARNING)


# ---------------------------------------------------------------------------
# sichgate estimate-cost (Show cost before running)
# ---------------------------------------------------------------------------

@app.command(name="estimate-cost")
def estimate_cost(
    mode: str = typer.Option(
        "api",
        "--mode",
        help="Estimation mode: api (27 queries) or local (200+ queries)"
    ),
    provider: str = typer.Option(
        "gpt-4o",
        "--provider",
        help="API provider (only for api mode): gpt-4o, gpt-3.5, claude-3-haiku, claude-3-sonnet"
    ),
) -> None:
    """Estimate API costs and query counts before running attacks."""
    
    from sichgate.api_config import get_default_api_config, APIRunConfig
    from sichgate.cost_estimator import CostEstimator
    
    if mode not in ["api", "local"]:
        console.print(f"[red]Invalid mode: {mode}. Use 'api' or 'local'.[/red]")
        raise typer.Exit(1)
    
    # Get appropriate config
    if mode == "api":
        config = get_default_api_config()
    else:
        from sichgate.api_config import get_default_local_config
        config = get_default_local_config()
    
    # Create estimator and show estimate
    estimator = CostEstimator(mode=mode)
    estimate = estimator.estimate_run(config, provider if mode == "api" else "gpt-4o")
    
    console.print(estimator.format_estimate_for_cli(estimate))
    
    # Show attack breakdown
    console.print("\n[bold]Attack Breakdown:[/bold]\n")
    
    if mode == "api":
        attacks_info = [
            ("Prompt Injection Direct & Indirect", 8, "Catch injection attacks"),
            ("Instruction Hijacking", 8, "Override model instructions"),
            ("Safe Messaging Violations", 5, "Bypass safety guidelines"),
            ("Hallucination Probes", 6, "Detect false information"),
            ("Token Smuggling", 4, "Encoding-based attacks"),
            ("PII Leakage", 3, "Information disclosure"),
        ]
    else:
        attacks_info = [
            ("Prompt Injection", 8, "High priority"),
            ("Instruction Hijacking", 8, "High priority"),
            ("Safe Messaging", 5, "Safety critical"),
            ("Hallucination", 6, "Safety critical"),
            ("Bias & Fairness", 48, "Comprehensive suite"),
            ("Consistency Checks", 20, "Robustness testing"),
            ("Crescendo Attack", 15, "Iterative exploitation"),
            ("Token Smuggling", 4, "Encoding attacks"),
            ("PII Leakage", 3, "Data exposure"),
            ("Membership Inference", 25, "Privacy attacks"),
            ("Data Extraction", 30, "Training data recovery"),
            ("Model Extraction", 40, "Model stealing"),
            ("Evasion Attacks", 50, "Robustness evasion"),
        ]
    
    from rich.table import Table
    table = Table(show_header=True, header_style="bold cyan")
    table.add_column("Attack Name")
    table.add_column("Queries", justify="right")
    table.add_column("Notes")
    
    for attack_name, query_count, notes in attacks_info:
        table.add_row(attack_name, str(query_count), notes)
    
    console.print(table)


# ---------------------------------------------------------------------------
# sichgate run
# ---------------------------------------------------------------------------

@app.command()
def run(
    model: str = typer.Option(..., "--model", "-m", help="Path to model or HuggingFace model ID"),
    attacks: str = typer.Option(
        "all-llm",
        "--attacks", "-a",
        help=(
            "Comma-separated attack names. Examples: prompt-injection,jailbreak,pii-leakage  "
            "Groups: all-llm, all-evasion, all"
        ),
    ),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write results to JSON file"),
    optimize: bool = typer.Option(False, "--optimize", help="[Pro] Auto-tune attacks with Optuna"),
    depth: str = typer.Option("quick", "--depth", help="Extraction depth: quick|standard|deep"),
    adapter: str = typer.Option("hf", "--adapter", help="Model adapter: hf|openai_compat|callable"),
    api_mode: bool = typer.Option(False, "--api-mode", help="[Pro] Lightweight attack set for API deployments (27 queries vs 200+)"),
    provider: str = typer.Option("gpt-4o", "--provider", help="API provider: gpt-4o, gpt-3.5, claude-3-haiku, claude-3-sonnet"),
    skip_estimate: bool = typer.Option(False, "--skip-estimate", help="Skip cost estimate confirmation (auto-proceed)"),
    verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose logging"),
    domain_data: Optional[str] = typer.Option(
        None, "--domain-data", help="Path to domain data for data-dependent attacks"
    ),
    aibom: bool = typer.Option(False, "--aibom", help="[2.1] Capture CycloneDX 1.6 AIBoM and embed in output"),
    attest: bool = typer.Option(False, "--attest", help="[2.1] Sign findings report with ed25519 attestation"),
    key_path: Optional[str] = typer.Option(None, "--key-path", help="[2.1] Path to ed25519 private key PEM (generated if absent)"),
) -> None:
    """Run the attack battery against a model."""

    if verbose:
        logging.getLogger().setLevel(logging.INFO)

    # Optuna warning
    if optimize:
        console.print(
            "[yellow]⚠ Note: auto-optimization is compute-intensive. "
            "Recommended for 16GB+ RAM. Expected runtime: 15-60 min depending on hardware.[/yellow]"
        )

    # If API mode, show cost estimate and ask for confirmation
    if api_mode and not skip_estimate:
        from sichgate.api_config import APIRunConfig, get_default_api_config
        from sichgate.cost_estimator import CostEstimator
        
        config = get_default_api_config()
        estimator = CostEstimator(mode="api")
        estimate = estimator.estimate_run(config, provider)
        
        # Display estimate
        console.print(estimator.format_estimate_for_cli(estimate))
        
        # Ask for confirmation
        proceed = typer.confirm("[bold]Proceed with API testing?[/bold]")
        if not proceed:
            console.print("[yellow]Cancelled.[/yellow]")
            raise typer.Exit(0)

    # Resolve attack classes
    from sichgate.runner import build_registry, AttackRunner

    registry = build_registry()
    attack_names = [a.strip() for a in attacks.split(",")]
    attack_classes = []
    unknown = []

    for name in attack_names:
        if name in registry:
            for cls in registry[name]:
                if cls not in attack_classes:
                    attack_classes.append(cls)
        else:
            unknown.append(name)

    if unknown:
        console.print(f"[red]Unknown attacks: {', '.join(unknown)}[/red]")
        console.print(f"Available: {', '.join(registry.keys())}")
        raise typer.Exit(1)

    mode_str = "[cyan]API mode[/cyan] — lightweight battery" if api_mode else "full battery"
    console.print(f"[bold]SichGate[/bold] — {mode_str} with {len(attack_classes)} attacks against [cyan]{model}[/cyan]")

    runner = AttackRunner(optimize=optimize)

    with console.status("Running attacks..."):
        summary = runner.run(
            model_path=model,
            attack_classes=attack_classes,
            adapter_type=adapter,
            depth=depth,
        )

    # Print summary table
    _print_summary(summary)

    result_dict = summary.to_dict()

    # Phase 2.1 — AIBoM capture
    if aibom:
        from sichgate.aibom_capture import capture_model_aibom
        with console.status("Capturing AIBoM..."):
            bom = capture_model_aibom(model_path=model)
        result_dict["aibom"] = bom.to_cyclonedx()
        console.print("[green]✓ AIBoM (CycloneDX 1.6 ML-BOM) captured[/green]")

    # Phase 2.1 — Cryptographic attestation
    if attest:
        from sichgate.crypto_attestation import CryptoAttestation
        att = CryptoAttestation(key_path=key_path)
        if key_path and not Path(key_path).exists():
            att.save_key(key_path)
            console.print(f"[green]✓ New signing key saved to {key_path}[/green]")
        with console.status("Signing report..."):
            signed = att.sign_findings_report(result_dict)
        result_dict = signed
        console.print("[green]✓ Report signed with ed25519 attestation[/green]")

    # Write output
    if output:
        output.write_text(json.dumps(result_dict, indent=2))
        console.print(f"\n[green]✓ Results written to {output}[/green]")
    else:
        console.print("\n[dim]Tip: use --output results.json to save[/dim]")

    # Exit non-zero if critical findings
    if summary.critical_count > 0:
        raise typer.Exit(2)


def _print_summary(summary) -> None:
    table = Table(title=f"Run {summary.run_id[:8]}...", show_header=True)
    table.add_column("Attack", style="cyan")
    table.add_column("Status", justify="center")
    table.add_column("Severity", justify="center")
    table.add_column("Notes")

    status_colors = {"fail": "red", "warning": "yellow", "pass": "green"}
    severity_colors = {"critical": "bold red", "high": "red", "medium": "yellow", "low": "green"}

    for r in summary.results:
        status = r["status"]
        severity = r["severity"]
        table.add_row(
            r["attack_name"],
            f"[{status_colors.get(status, 'white')}]{status}[/]",
            f"[{severity_colors.get(severity, 'white')}]{severity}[/]",
            r["notes"][:80] + "..." if len(r["notes"]) > 80 else r["notes"],
        )

    console.print(table)
    console.print(
        f"\n[bold]Summary:[/bold] "
        f"[red]{summary.failed} failed[/red]  "
        f"[yellow]{summary.warnings} warnings[/yellow]  "
        f"[green]{summary.passed} passed[/green]  "
        f"| Critical: [bold red]{summary.critical_count}[/bold red]"
    )


# ---------------------------------------------------------------------------
# sichgate activate (Phase 2 — license key)
# ---------------------------------------------------------------------------

@app.command()
def activate(
    key: Optional[str] = typer.Argument(None, help="License key (prompted if not provided)"),
) -> None:
    """Activate SichGate Pro license. Phones home once, then fully offline."""
    try:
        from sichgate.license import LicenseManager
    except ImportError:
        console.print("[red]License module not available in this build.[/red]")
        raise typer.Exit(1)

    if not key:
        key = typer.prompt("Enter your SichGate Pro license key")

    lm = LicenseManager()
    with console.status("Activating license..."):
        ok, message = lm.activate(key)

    if ok:
        console.print(f"[green]✓ {message}[/green]")
        console.print("[dim]License stored. No further internet access required.[/dim]")
    else:
        console.print(f"[red]✗ Activation failed: {message}[/red]")
        raise typer.Exit(1)


@app.command()
def deactivate() -> None:
    """Free this machine's license seat (e.g. before switching to a new machine)."""
    try:
        from sichgate.license import LicenseManager
    except ImportError:
        console.print("[red]License module not available.[/red]")
        raise typer.Exit(1)

    lm = LicenseManager()
    ok, message = lm.deactivate()
    if ok:
        console.print(f"[green]✓ {message}[/green]")
    else:
        console.print(f"[red]✗ {message}[/red]")
        raise typer.Exit(1)


# ---------------------------------------------------------------------------
# sichgate start (Phase 2 — launches dashboard)
# ---------------------------------------------------------------------------

@app.command()
def start(
    port: int = typer.Option(8000, "--port", help="API port"),
    open_browser: bool = typer.Option(True, "--browser/--no-browser", help="Open dashboard in browser"),
) -> None:
    """Launch the SichGate Pro local dashboard (localhost)."""
    try:
        from sichgate.api.server import create_app
    except ImportError:
        console.print("[red]Server module not available. Install sichgate-pro.[/red]")
        raise typer.Exit(1)

    console.print(f"[bold green]SichGate Pro[/bold green] starting on http://localhost:{port}")
    console.print("[dim]Everything runs locally. No data leaves this machine.[/dim]")

    if open_browser:
        import threading, webbrowser, time
        def _open():
            time.sleep(1.5)
            webbrowser.open(f"http://localhost:{port}")
        threading.Thread(target=_open, daemon=True).start()

    import uvicorn
    uvicorn_app = create_app()
    uvicorn.run(uvicorn_app, host="127.0.0.1", port=port, log_level="warning")


# ---------------------------------------------------------------------------
# sichgate diff  (Phase 2.1 — differential redteaming)
# ---------------------------------------------------------------------------

@app.command()
def diff(
    base_model: str = typer.Option(..., "--base-model", help="Base model path or HuggingFace model ID"),
    custom_model: str = typer.Option(..., "--custom-model", help="Fine-tuned model path or HuggingFace model ID"),
    attacks: str = typer.Option(
        "all-llm",
        "--attacks", "-a",
        help="Comma-separated attack names or groups (e.g. all-llm)",
    ),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write differential report to JSON file"),
    adapter: str = typer.Option("hf", "--adapter", help="Model adapter: hf|openai_compat|callable"),
    aibom: bool = typer.Option(True, "--aibom/--no-aibom", help="Embed CycloneDX 1.6 AIBoM for both models"),
    attest: bool = typer.Option(False, "--attest", help="Sign differential report with ed25519"),
    key_path: Optional[str] = typer.Option(None, "--key-path", help="Path to ed25519 private key PEM"),
    verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose logging"),
) -> None:
    """
    [2.1] Differential redteaming: compare base vs fine-tuned model attack results.

    Runs the full attack battery against both models and produces a DifferentialReport
    showing which vulnerabilities worsened, improved, or remained unchanged after
    fine-tuning.  Outputs comply with EU AI Act Article 10 evidence requirements.

    Example:

        sichgate diff --base-model qwen2-1.5b \\
                      --custom-model ./medical-finetuned.gguf \\
                      --output differential_report.json
    """
    if verbose:
        logging.getLogger().setLevel(logging.INFO)

    from sichgate.runner import build_registry, AttackRunner
    from sichgate.differential_redteaming import DifferentialRedteaming

    registry = build_registry()
    attack_names = [a.strip() for a in attacks.split(",")]
    attack_classes = []
    unknown = []
    for name in attack_names:
        if name in registry:
            for cls in registry[name]:
                if cls not in attack_classes:
                    attack_classes.append(cls)
        else:
            unknown.append(name)

    if unknown:
        console.print(f"[red]Unknown attacks: {', '.join(unknown)}[/red]")
        raise typer.Exit(1)

    runner = AttackRunner()

    console.print(f"[bold]SichGate Diff[/bold] — running {len(attack_classes)} attacks on [cyan]{base_model}[/cyan] (base)")
    with console.status(f"Scanning base model: {base_model}..."):
        base_summary = runner.run(
            model_path=base_model,
            attack_classes=attack_classes,
            adapter_type=adapter,
        )

    console.print(f"[bold]SichGate Diff[/bold] — running {len(attack_classes)} attacks on [cyan]{custom_model}[/cyan] (custom)")
    with console.status(f"Scanning custom model: {custom_model}..."):
        custom_summary = runner.run(
            model_path=custom_model,
            attack_classes=attack_classes,
            adapter_type=adapter,
        )

    # Optional AIBoM capture
    base_bom_dict: Optional[dict] = None
    custom_bom_dict: Optional[dict] = None
    if aibom:
        from sichgate.aibom_capture import capture_model_aibom
        base_bom_dict = capture_model_aibom(model_path=base_model).to_cyclonedx()
        custom_bom_dict = capture_model_aibom(model_path=custom_model).to_cyclonedx()
        console.print("[green]✓ AIBoM captured for both models[/green]")

    # Optional attestation
    attestation_obj = None
    if attest:
        from sichgate.crypto_attestation import CryptoAttestation
        attestation_obj = CryptoAttestation(key_path=key_path)
        if key_path and not Path(key_path).exists():
            attestation_obj.save_key(key_path)
            console.print(f"[green]✓ New signing key saved to {key_path}[/green]")

    engine = DifferentialRedteaming()
    report = engine.compare(
        base_results=base_summary.results,
        custom_results=custom_summary.results,
        base_model_id=base_model,
        custom_model_id=custom_model,
        base_aibom=base_bom_dict,
        custom_aibom=custom_bom_dict,
        attestation=attestation_obj,
    )

    # Print differential summary
    _print_differential_summary(report)

    if output:
        output.write_text(report.to_json())
        console.print(f"\n[green]✓ Differential report written to {output}[/green]")
    else:
        console.print("\n[dim]Tip: use --output diff_report.json to save[/dim]")

    if report.critical_alerts:
        raise typer.Exit(2)


def _print_differential_summary(report) -> None:
    table = Table(title=f"Differential Report {report.run_id[:8]}...", show_header=True)
    table.add_column("Attack Type", style="cyan")
    table.add_column("Base Rate", justify="right")
    table.add_column("Custom Rate", justify="right")
    table.add_column("Delta", justify="right")
    table.add_column("Trend", justify="center")
    table.add_column("Alert")

    trend_colors = {"WORSENED": "red", "IMPROVED": "green", "UNCHANGED": "white"}
    for attack_type, delta in report.differential_analysis.items():
        color = trend_colors.get(delta.trend, "white")
        alert = delta.compliance_alert or ""
        table.add_row(
            attack_type,
            f"{delta.base_success_rate:.1%}",
            f"{delta.custom_success_rate:.1%}",
            f"[{color}]{delta.delta:+.1%}[/]",
            f"[{color}]{delta.trend}[/]",
            f"[bold red]{alert}[/bold red]" if "CRITICAL" in alert else alert,
        )

    console.print(table)
    delta_color = "red" if report.overall_risk_delta > 0 else "green"
    console.print(
        f"\n[bold]Overall risk delta:[/bold] [{delta_color}]{report.overall_risk_delta:+.1%}[/]  "
        f"| Critical alerts: [bold red]{len(report.critical_alerts)}[/bold red]"
    )
    for alert in report.critical_alerts:
        console.print(f"  [bold red]⚠ {alert}[/bold red]")


# ---------------------------------------------------------------------------
# sichgate technical-file  (Phase 2.1 — EU AI Act Article 11)
# ---------------------------------------------------------------------------

@app.command(name="technical-file")
def technical_file(
    findings: Path = typer.Option(
        ..., "--findings", help="Path to findings JSON (from sichgate run --output)"
    ),
    aibom: Optional[Path] = typer.Option(
        None, "--aibom", help="Path to AIBoM JSON (from sichgate run --aibom --output)"
    ),
    differential: Optional[Path] = typer.Option(
        None, "--differential", help="Path to differential report JSON (from sichgate diff)"
    ),
    attest: bool = typer.Option(
        False, "--attest", help="Sign the Technical File with ed25519 attestation"
    ),
    key_path: Optional[str] = typer.Option(
        None, "--key-path", help="Path to ed25519 private key PEM (generated if absent)"
    ),
    output: Path = typer.Option(
        Path("technical_file.pdf"),
        "--output", "-o",
        help="Output PDF path (HTML and JSON saved alongside)",
    ),
) -> None:
    """
    [2.1] Generate an EU AI Act Article 11 Technical File.

    Assembles findings, AIBoM, and optional differential report into a
    structured Technical File document, exports it as PDF, HTML, and JSON.

    Example:

        sichgate technical-file \\
          --findings findings.json \\
          --aibom aibom.json \\
          --differential differential_report.json \\
          --attest \\
          --output technical_file.pdf
    """
    import json as _json

    from sichgate.technical_file_builder import TechnicalFileBuilder
    from sichgate.pdf_exporter import PDFExporter

    # ── Load findings ──────────────────────────────────────────────────
    if not findings.exists():
        console.print(f"[red]Findings file not found: {findings}[/red]")
        raise typer.Exit(1)

    with findings.open() as fh:
        findings_data = _json.load(fh)

    # Support both a raw list and a RunSummary dict with "results" key
    if isinstance(findings_data, dict):
        findings_list = findings_data.get("results", [])
        aibom_from_findings = findings_data.get("aibom")
    else:
        findings_list = findings_data
        aibom_from_findings = None

    # ── Load AIBoM ─────────────────────────────────────────────────────
    aibom_dict: Optional[dict] = None
    if aibom and aibom.exists():
        with aibom.open() as fh:
            aibom_dict = _json.load(fh)
        console.print("[dim]AIBoM loaded from file[/dim]")
    elif aibom_from_findings:
        aibom_dict = aibom_from_findings
        console.print("[dim]AIBoM loaded from findings file[/dim]")

    # ── Load differential report ───────────────────────────────────────
    diff_dict: Optional[dict] = None
    if differential and differential.exists():
        with differential.open() as fh:
            diff_dict = _json.load(fh)
        console.print("[dim]Differential report loaded[/dim]")

    # ── Build Technical File ───────────────────────────────────────────
    with console.status("Building Technical File..."):
        builder = TechnicalFileBuilder(
            findings=findings_list,
            aibom=aibom_dict,
            differential_report=diff_dict,
        )
        tf = builder.build()

    # ── Optional attestation ───────────────────────────────────────────
    if attest:
        from sichgate.crypto_attestation import CryptoAttestation
        att = CryptoAttestation(key_path=key_path)
        if key_path and not Path(key_path).exists():
            att.save_key(key_path)
            console.print(f"[green]✓ New signing key saved to {key_path}[/green]")
        with console.status("Signing Technical File..."):
            tf_dict = tf.to_dict()
            signed = att.sign_findings_report(tf_dict)
            tf.attestation = signed.get("attestation")
        console.print("[green]✓ Technical File signed with ed25519 attestation[/green]")

    # ── Persist to database ────────────────────────────────────────────
    try:
        from sichgate.database import DatabaseManager
        db = DatabaseManager()
        record_id = db.save_technical_file(tf)
        console.print(f"[dim]Saved to database (ID: {record_id[:8]}...)[/dim]")
    except Exception as exc:  # noqa: BLE001
        console.print(f"[dim yellow]DB save skipped: {exc}[/dim yellow]")

    # ── Save JSON artefact ─────────────────────────────────────────────
    json_path = output.with_suffix(".json")
    json_path.write_text(tf.to_json(), encoding="utf-8")
    console.print(f"[green]✓ Technical File JSON saved: {json_path}[/green]")

    # ── Export PDF ─────────────────────────────────────────────────────
    try:
        exporter = PDFExporter()
        with console.status("Exporting PDF..."):
            pdf_path = exporter.export_pdf(tf, str(output))
        html_path = output.with_suffix(".html")
        console.print(f"[green]✓ Technical File generated: {pdf_path}[/green]")
        console.print(f"[green]✓ HTML artefact saved: {html_path}[/green]")
    except RuntimeError as exc:
        console.print(f"[yellow]⚠ PDF export skipped: {exc}[/yellow]")
        console.print("[dim]Install weasyprint for PDF output: pip install weasyprint[/dim]")


# ---------------------------------------------------------------------------
# sichgate list-attacks
# ---------------------------------------------------------------------------

@app.command(name="check")
def system_check(
    verbose: bool = typer.Option(False, "--verbose", "-v", help="Show detailed diagnostics"),
    fail_on_warn: bool = typer.Option(False, "--strict", help="Exit non-zero on WARN or BLOCK"),
) -> None:
    """Run system health checks before starting a red team run."""
    from sichgate.system_health import run_checks, print_report, CheckLevel
    report = run_checks()
    print_report(report, verbose=verbose)
    if not report.can_proceed:
        raise typer.Exit(2)
    if fail_on_warn and report.worst_level in (CheckLevel.WARN, CheckLevel.BLOCK):
        raise typer.Exit(1)


@app.command(name="list-attacks")
def list_attacks() -> None:
    """List all available attack names and their groups."""
    from sichgate.runner import build_registry
    registry = build_registry()

    table = Table(title="Available Attacks", show_header=True)
    table.add_column("Name", style="cyan")
    table.add_column("Type")
    table.add_column("Attack Classes")

    for name, classes in registry.items():
        attack_types = list({getattr(c, "attack_type", "?") for c in classes})
        class_names = ", ".join(c.__name__ for c in classes)
        table.add_row(name, ", ".join(attack_types), class_names)

    console.print(table)


if __name__ == "__main__":
    app()
