"""
AIR Blackbox CLI - AI governance control plane.

    air-blackbox setup       # One-command setup: install model + verify
    air-blackbox discover    # Shadow AI inventory + AI-BOM
    air-blackbox comply      # EU AI Act compliance from live traffic
    air-blackbox replay      # Incident reconstruction from audit chain
    air-blackbox export      # Signed evidence bundle for auditors
    air-blackbox validate    # Pre-execution runtime checks
    air-blackbox test        # End-to-end stack validation
    air-blackbox demo        # Zero-config demo with sample data
    air-blackbox init        # Initialize project templates
"""

import click
from datetime import datetime
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from air_blackbox import __version__ as _ab_version

console = Console()


AIR_BANNER = r"""[bold #00d4aa]
    _    ___ ____    ____  _            _    _
   / \  |_ _|  _ \  | __ )| | __ _  ___| | _| |__   _____  __
  / _ \  | || |_) | |  _ \| |/ _` |/ __| |/ / '_ \ / _ \ \/ /
 / ___ \ | ||  _ <  | |_) | | (_| | (__|   <| |_) | (_) >  <
/_/   \_\___|_| \_\ |____/|_|\__,_|\___|_|\_\_.__/ \___/_/\_\
[/bold #00d4aa]"""



def print_banner():
    from rich.rule import Rule
    console.print(AIR_BANNER)
    console.print("  [dim]EU AI Act Compliance · AI-BOM · Audit Chain · Incident Replay[/dim]")
    console.print()
    console.print("  " + "─" * 76, style="#1e2530")
    console.print("  [bold #f85149]⚠  High-risk obligations apply December 2, 2027 (2026 omnibus) - €35M or 7% global turnover[/bold #f85149]")
    console.print("  " + "─" * 76, style="#1e2530")
    console.print("  [dim]pip install air-blackbox  ·  github.com/airblackbox/airblackbox  ·  airblackbox.ai[/dim]")
    console.print()


@click.group()
@click.version_option(version=_ab_version, prog_name="air-blackbox")
@click.pass_context
def main(ctx):
    """AIR Blackbox - AI governance control plane.

    Route your AI traffic through the gateway and get compliance,
    security, inventory, and incident response out of the box.
    """
    if ctx.invoked_subcommand not in ("--version", None):
        print_banner()


@main.command()
def setup():
    """One-command setup: install the AI compliance model and verify everything works.

    This pulls the air-compliance model from Ollama registry and verifies
    the scanner is ready to use. Run this once after installing air-blackbox.

    Requirements: Ollama must be installed first (https://ollama.com)
    """
    import subprocess
    import shutil

    console.print(Panel.fit(
        "[bold cyan]AIR Blackbox Setup[/bold cyan]\n"
        "Setting up the AI compliance scanner...",
        border_style="cyan",
    ))

    # Step 1: Check Ollama
    console.print("\n[bold]Step 1/3:[/bold] Checking Ollama installation...")
    if shutil.which("ollama"):
        try:
            result = subprocess.run(["ollama", "--version"], capture_output=True, text=True, timeout=5)
            console.print(f"  [green]✓[/green] Ollama installed: {result.stdout.strip()}")
        except Exception:
            console.print("  [green]✓[/green] Ollama found")
    else:
        console.print("  [red]✗[/red] Ollama not installed")
        console.print("\n  Install Ollama first:")
        console.print("    Mac:   [cyan]brew install ollama[/cyan]")
        console.print("    Linux: [cyan]curl -fsSL https://ollama.com/install.sh | sh[/cyan]")
        console.print("    All:   [cyan]https://ollama.com/download[/cyan]")
        console.print("\n  Then run [cyan]air-blackbox setup[/cyan] again.")
        return

    # Step 2: Pull model
    console.print("\n[bold]Step 2/3:[/bold] Pulling air-compliance model from registry...")
    console.print("  This downloads ~8GB (one-time). Grab a coffee.\n")

    try:
        result = subprocess.run(
            ["ollama", "pull", "airblackbox/air-compliance"],
            timeout=600,
        )
        if result.returncode == 0:
            # Create local alias
            subprocess.run(
                ["ollama", "cp", "airblackbox/air-compliance", "air-compliance"],
                capture_output=True, timeout=30,
            )
            console.print("  [green]✓[/green] Model pulled and ready")
        else:
            console.print("  [red]✗[/red] Failed to pull model")
            console.print("  Try manually: [cyan]ollama pull airblackbox/air-compliance[/cyan]")
            return
    except subprocess.TimeoutExpired:
        console.print("  [red]✗[/red] Download timed out. Try: [cyan]ollama pull airblackbox/air-compliance[/cyan]")
        return

    # Step 3: Verify
    console.print("\n[bold]Step 3/3:[/bold] Verifying scanner...")
    try:
        result = subprocess.run(["ollama", "list"], capture_output=True, text=True, timeout=10)
        if "air-compliance" in result.stdout:
            console.print("  [green]✓[/green] Model verified in Ollama")
        else:
            console.print("  [yellow]⚠[/yellow] Model pulled but not showing in list. Try restarting Ollama.")
    except Exception:
        # Best-effort verification; do not fail setup if this check errors
        console.print("  [yellow]⚠[/yellow] Could not verify model in Ollama (verification step failed).")

    console.print(Panel.fit(
        "[bold green]Setup complete![/bold green]\n\n"
        "Run your first scan:\n"
        "  [cyan]air-blackbox comply --scan .[/cyan]\n\n"
        "Or try the demo:\n"
        "  [cyan]air-blackbox demo[/cyan]",
        border_style="green",
    ))


@main.command()
@click.option("--gateway", default="http://localhost:8080", help="Gateway URL")
@click.option("--scan", default=".", help="Path to scan for code-level checks")
@click.option("--runs-dir", default=None, help="Path to .air.json records directory")
@click.option("--format", "fmt", type=click.Choice(["table", "json"]), default="table")
@click.option("--verbose", "-v", is_flag=True, help="Show detection type and fix hints")
@click.option("--deep", is_flag=True, default=False, hidden=True, help="(deprecated, now default) Run LLM deep analysis")
@click.option("--no-llm", is_flag=True, help="Skip LLM analysis, regex-only scan")
@click.option("--model", default="air-compliance", help="Ollama model for deep scan")
@click.option("--no-save", is_flag=True, help="Don't save results to compliance history")
@click.option("--standard", default="all", type=click.Choice(["eu", "us", "all"]), help="Jurisdiction: eu, us, or all")
def comply(gateway, scan, runs_dir, fmt, verbose, deep, no_llm, model, no_save, standard):
    """Check AI compliance across EU and US jurisdictions."""
    from air_blackbox.gateway_client import GatewayClient
    from air_blackbox.compliance.engine import run_all_checks
    std_label = {"eu": "EU AI Act", "us": "US State Laws", "all": "EU + US Jurisdictions"}[standard]
    console.print(f"\n[bold blue]AIR Blackbox[/] - {std_label} Compliance Check\n")
    with console.status("[bold green]Connecting to gateway..."):
        client = GatewayClient(gateway_url=gateway, runs_dir=runs_dir, scan_path=scan)
        status = client.get_status()
    if status.reachable:
        console.print(f"  [green]●[/] Gateway connected at [bold]{gateway}[/]")
    else:
        console.print(f"  [dim]○ Gateway not running at {gateway} (optional - enables runtime checks)[/]")
    if status.total_runs > 0:
        src = "gateway" if status.reachable else "trust layer records"
        console.print(f"  [green]●[/] [bold]{status.total_runs:,}[/] logged events from {src} ({', '.join(status.models_observed[:3])})")
    else:
        console.print(f"  [yellow]●[/] No traffic data found")
    console.print(f"  [dim]Scanning: {scan}[/]\n")
    articles, detected_frameworks, rec_pkg, crosswalk_hint = run_all_checks(status, scan, standard=standard)

    # Hybrid mode: auto-run LLM analysis unless --no-llm
    deep_findings = []
    if not no_llm:
        from air_blackbox.compliance.deep_scan import deep_scan, _ollama_available, _model_available
        import os
        if _ollama_available() and _model_available(model):
            console.print(f"[bold]Running hybrid analysis (regex + AI model)...[/]\n")
            # Collect all Python files (supports single-file and directory scanning)
            py_files = []
            if os.path.isfile(scan) and scan.endswith(".py"):
                py_files = [os.path.abspath(scan)]
            else:
                skip_dirs = {"node_modules", ".git", "__pycache__", ".venv", "venv",
                             "dist", "build", ".eggs", "site-packages", ".tox",
                             ".mypy_cache", ".pytest_cache"}
                for root, dirs, files in os.walk(scan):
                    dirs[:] = [d for d in dirs if d not in skip_dirs and not d.endswith(".egg-info")]
                    for f in files:
                        if f.endswith(".py"):
                            py_files.append(os.path.join(root, f))
            total_files = len(py_files)

            # === Smart sampling: pick compliance-relevant files ===
            # Priority keywords - files most likely to contain compliance patterns
            priority_keywords = [
                "agent", "pipeline", "tool", "llm", "model", "chat",
                "safety", "guard", "policy", "policies", "hitl", "human",
                "trace", "tracing", "logging", "log", "audit", "monitor",
                "auth", "token", "scope", "permission", "identity",
                "validate", "validation", "schema", "pii", "redact",
                "retry", "fallback", "error", "exception", "handler",
                "inject", "sanitize", "filter", "boundary", "limit",
                "config", "settings", "core", "main", "app", "run",
            ]
            # Score and rank files by relevance
            def _score_file(fp):
                rel = os.path.relpath(fp, scan).lower()
                basename = os.path.basename(fp).lower()
                score = 0
                # Skip test files - they don't reflect compliance posture
                parts = rel.replace("\\", "/").split("/")
                if any(p in {"tests", "test", "testing"} for p in parts):
                    return -1
                if basename.startswith("test_") or basename == "conftest.py":
                    return -1
                # Boost files with compliance-relevant names
                for kw in priority_keywords:
                    if kw in basename:
                        score += 3
                    elif kw in rel:
                        score += 1
                # Boost core source files
                if any(p in {"src", "core", "lib", "components"} for p in parts):
                    score += 2
                # Boost larger files (more substance)
                try:
                    size = os.path.getsize(fp)
                    if size > 5000:
                        score += 2
                    elif size > 1000:
                        score += 1
                except OSError:
                    pass
                return score

            scored = [(fp, _score_file(fp)) for fp in py_files]
            scored = [(fp, s) for fp, s in scored if s >= 0]  # exclude tests
            scored.sort(key=lambda x: x[1], reverse=True)

            # Build code sample from top-ranked files
            # Cap individual files at 8KB to preserve full context, total at 48KB
            MAX_PER_FILE = 8000
            MAX_TOTAL = 48000
            code_parts = []
            total_chars = 0
            files_included = 0
            for fp, score in scored:
                try:
                    with open(fp, "r", encoding="utf-8", errors="ignore") as fh:
                        content = fh.read()
                    if len(content.strip()) < 50:
                        continue
                    # Truncate large files to get more variety in the sample
                    if len(content) > MAX_PER_FILE:
                        content = content[:MAX_PER_FILE] + "\n# ... (file truncated for sampling)"
                    code_parts.append(f"# File: {os.path.relpath(fp, scan)}\n{content}")
                    total_chars += len(content)
                    files_included += 1
                    if total_chars > MAX_TOTAL:
                        break
                except Exception:
                    continue

            merged_code = "\n\n".join(code_parts)
            sample_desc = f"targeted sample of {files_included} compliance-relevant source files"

            if verbose:
                console.print(f"  [dim]AI model sampling: {files_included} files from {total_files} total ({total_chars:,} chars)[/]")
                # Show top 5 sampled files
                shown = 0
                for fp, score in scored:
                    if shown >= 5:
                        break
                    rel = os.path.relpath(fp, scan)
                    try:
                        size = os.path.getsize(fp)
                        if size >= 50:
                            console.print(f"    [dim]score={score}: {rel} ({size:,} bytes)[/]")
                            shown += 1
                    except OSError:
                        pass
                if files_included > 5:
                    console.print(f"    [dim]... and {files_included - 5} more files[/]")

            # Build rule-based context summary for the model
            rule_context_lines = []
            article_map = {9: "Risk Management", 10: "Data Governance",
                          11: "Technical Documentation", 12: "Record-Keeping",
                          14: "Human Oversight", 15: "Accuracy & Security"}
            for article in articles:
                art_num = article.get("number", 0)
                if art_num not in article_map:
                    continue
                passes = []
                fails = []
                warns = []
                for check in article.get("checks", []):
                    name = check.get("name", "")
                    evidence = check.get("evidence", "")
                    status = check.get("status", "")
                    summary = f"{name}: {evidence[:80]}" if evidence else name
                    if status == "pass":
                        passes.append(summary)
                    elif status == "fail":
                        fails.append(summary)
                    elif status == "warn":
                        warns.append(summary)
                line = f"Article {art_num} ({article_map[art_num]}): "
                if passes:
                    line += f"{len(passes)} PASS ({'; '.join(passes[:2])})"
                if fails:
                    line += f", {len(fails)} FAIL" if passes else f"{len(fails)} FAIL"
                if warns:
                    line += f", {len(warns)} WARN" if (passes or fails) else f"{len(warns)} WARN"
                rule_context_lines.append(line)
            rule_context = "\n".join(rule_context_lines)

            # Only run AI model if we have actual code to analyze
            if files_included == 0 or not merged_code.strip():
                if verbose:
                    console.print(f"  [dim]No Python files found for AI analysis - skipping model[/]")
                result = {"available": False, "findings": [], "model": model, "error": None}
            else:
                if verbose:
                    os.environ["AIR_VERBOSE"] = "1"
                result = deep_scan(merged_code, model=model,
                                  sample_context=sample_desc,
                                  total_files=total_files,
                                  rule_context=rule_context)
                if verbose:
                    os.environ.pop("AIR_VERBOSE", None)
            if result.get("available") and not result.get("error"):
                deep_findings = result.get("findings", [])

                # ── Smart reconciliation: override model FAIL when rule-based has strong PASS ──
                # Build a map of rule-based pass counts per article
                rule_pass_counts = {}
                rule_evidence_map = {}
                for article in articles:
                    art_num = article.get("number", 0)
                    passes = [c for c in article.get("checks", []) if c.get("status") == "pass"]
                    rule_pass_counts[art_num] = len(passes)
                    if passes:
                        # Collect the best evidence summaries
                        rule_evidence_map[art_num] = "; ".join(
                            c.get("evidence", "")[:60] for c in passes[:3]
                        )

                overrides = 0
                for finding in deep_findings:
                    art = finding.get("article", 0)
                    model_status = finding.get("status", "")
                    rule_passes = rule_pass_counts.get(art, 0)

                    # If model says FAIL but rule-based has 2+ PASS checks → override to PASS
                    if model_status == "fail" and rule_passes >= 2:
                        finding["status"] = "pass"
                        rule_ev = rule_evidence_map.get(art, "")
                        finding["evidence"] = (
                            f"[Corrected by rule-based analysis] "
                            f"Rule-based scanner found {rule_passes} passing checks: {rule_ev}. "
                            f"Model's original assessment: {finding.get('evidence', '')}"
                        )
                        finding["fix_hint"] = ""
                        overrides += 1
                    # If model says FAIL but rule-based has 1 PASS → upgrade to WARN
                    elif model_status == "fail" and rule_passes == 1:
                        finding["status"] = "warn"
                        rule_ev = rule_evidence_map.get(art, "")
                        finding["evidence"] = (
                            f"[Partial - rule-based found evidence] {rule_ev}. "
                            f"Model noted: {finding.get('evidence', '')}"
                        )
                        overrides += 1

                console.print(f"  [green]●[/] AI model analyzed [bold]{files_included}[/] files ({total_chars:,} chars) from {total_files} total")
                console.print(f"  [green]●[/] AI model found [bold]{len(deep_findings)}[/] finding(s) using [bold]{model}[/]")
                if overrides > 0:
                    console.print(f"  [green]●[/] Smart reconciliation: [bold]{overrides}[/] model verdict(s) corrected by rule-based evidence")
                console.print(f"  [green]●[/] Hybrid mode: rule-based + AI analysis merged\n")
            elif result.get("error"):
                console.print(f"  [yellow]●[/] AI model: {result['error']}")
                console.print(f"  [dim]Falling back to regex-only scan[/]\n")
        else:
            if verbose:
                console.print(f"  [dim]AI model not available - using regex-only scan[/]")
                console.print(f"  [dim]Optional: install the local model - https://github.com/airblackbox/airblackbox/tree/main/model[/]\n")

    # Save to compliance history
    if not no_save:
        try:
            from air_blackbox.compliance.history import save_scan
            scan_id = save_scan(articles, scan_path=scan, version="1.6.3",
                                deep_findings=deep_findings if deep_findings else None)
            if verbose:
                console.print(f"  [dim]Saved to compliance history (scan #{scan_id})[/]\n")
        except Exception:
            # Don't break the scan if history save fails
            if verbose:
                console.print("  [dim]Could not save to compliance history[/]")

    if fmt == "json":
        import json
        output_data = articles
        if deep_findings:
            output_data = list(articles)  # shallow copy
            output_data.append({
                "number": 0,
                "title": "LLM Deep Analysis",
                "checks": [{
                    "name": f.get("name", ""),
                    "status": f.get("status", "warn"),
                    "evidence": f.get("evidence", ""),
                    "fix_hint": f.get("fix_hint", ""),
                    "tier": "static",
                    "detection": "auto",
                    "source": "llm",
                } for f in deep_findings],
            })
        click.echo(json.dumps(output_data, indent=2))
        return
    for article in articles:
        table = Table(title=f"Article {article['number']} - {article['title']}",
            show_header=True, header_style="bold white on dark_blue", title_style="bold")
        table.add_column("Check", style="bold", width=28)
        table.add_column("Tier", width=8, justify="center")
        table.add_column("Status", width=10, justify="center")
        if verbose:
            table.add_column("Type", width=8, justify="center")
        table.add_column("Evidence", width=42 if not verbose else 36)
        for check in article["checks"]:
            si = {"pass": "[bold green]✅ PASS[/]", "warn": "[bold yellow]⚠️  WARN[/]", "fail": "[bold red]❌ FAIL[/]"}.get(check["status"])
            tier = check.get("tier", "static")
            tier_label = "[green]STATIC[/]" if tier == "static" else "[blue]RUNTIME[/]"
            db = {"auto": "[green]AUTO[/]", "hybrid": "[yellow]HYBRID[/]", "manual": "[red]MANUAL[/]"}.get(check.get("detection", ""), "")
            ev = check["evidence"]
            if verbose and check.get("fix_hint"):
                ev += f"\n[dim italic]Fix: {check['fix_hint']}[/]"
            row = [check["name"], tier_label, si]
            if verbose:
                row.append(db)
            row.append(ev)
            table.add_row(*row)
        console.print(table)
        console.print()

    # Display deep findings if any (supplementary - not counted in main score)
    if deep_findings and verbose:
        # Only show LLM findings that ADD info beyond what rules found
        rule_articles = {a["number"] for a in articles}
        novel_findings = [f for f in deep_findings if f.get("article", 0) not in rule_articles]
        supplementary = [f for f in deep_findings if f.get("article", 0) in rule_articles]
        if novel_findings or supplementary:
            deep_table = Table(title="AI Model Insights (supplementary)",
                show_header=True, header_style="bold white on dark_blue", title_style="dim")
            deep_table.add_column("Article", width=10, justify="center")
            deep_table.add_column("AI Assessment", style="bold", width=30)
            deep_table.add_column("Status", width=10, justify="center")
            deep_table.add_column("Evidence", width=40)
            for f in deep_findings:
                si = {"pass": "[bold green]✅ PASS[/]", "warn": "[bold yellow]⚠️  WARN[/]", "fail": "[bold red]❌ FAIL[/]"}.get(f.get("status", "warn"))
                ev = f.get("evidence", "")
                if f.get("fix_hint"):
                    ev += f"\n[dim italic]Fix: {f['fix_hint']}[/]"
                deep_table.add_row(f"Art {f.get('article', '?')}", f.get("name", ""), si, ev)
            console.print(deep_table)
            console.print("[dim]  Note: AI model assessed a code sample (max 48KB). Rule-based checks above are more accurate.[/]")
            console.print()

    total = sum(len(a["checks"]) for a in articles)
    passing = sum(1 for a in articles for c in a["checks"] if c["status"] == "pass")
    warning = sum(1 for a in articles for c in a["checks"] if c["status"] == "warn")
    failing = sum(1 for a in articles for c in a["checks"] if c["status"] == "fail")
    # Two-tier breakdown
    static_checks = [c for a in articles for c in a["checks"] if c.get("tier", "static") == "static"]
    runtime_checks = [c for a in articles for c in a["checks"] if c.get("tier") == "runtime"]
    s_pass = sum(1 for c in static_checks if c["status"] == "pass")
    s_total = len(static_checks)
    r_pass = sum(1 for c in runtime_checks if c["status"] == "pass")
    r_total = len(runtime_checks)
    parts = f"[bold green]{passing}[/] passing  [bold yellow]{warning}[/] warnings  [bold red]{failing}[/] failing  out of [bold]{total}[/] checks"
    parts += f"\n\n  [green]Static analysis[/]:  [bold]{s_pass}/{s_total}[/] passing  (code patterns, docs, config)"
    parts += f"\n  [blue]Runtime checks[/]:   [bold]{r_pass}/{r_total}[/] passing  (requires gateway or trust layer)"
    if deep_findings:
        deep_pass = sum(1 for f in deep_findings if f.get("status") == "pass")
        deep_fail = sum(1 for f in deep_findings if f.get("status") == "fail")
        deep_warn = sum(1 for f in deep_findings if f.get("status") == "warn")
        parts += f"\n  [magenta]AI model[/]:          [bold]{deep_pass}[/] pass, [bold]{deep_warn}[/] warn, [bold]{deep_fail}[/] fail (supplementary, not counted above)"
    if r_total > 0 and r_pass < r_total:
        parts += f"\n\n  [dim]Unlock runtime checks: pip install {rec_pkg}[/]"
    if verbose:
        auto = sum(1 for a in articles for c in a["checks"] if c.get("detection") == "auto")
        hybrid = sum(1 for a in articles for c in a["checks"] if c.get("detection") == "hybrid")
        manual = sum(1 for a in articles for c in a["checks"] if c.get("detection") == "manual")
        parts += f"\n  [dim]Detection: {auto} auto, {hybrid} hybrid, {manual} manual ({(auto + hybrid) / total * 100:.0f}% automated)[/]"
    console.print(Panel(parts, title="[bold]Compliance Summary[/]", border_style="blue"))
    if failing > 0 and not verbose:
        console.print("\n[dim]Run with -v to see fix hints for each failing check.[/]")

    # --- Trust layer recommendation ---
    from air_blackbox.compliance.engine import TRUST_LAYER_MAP
    if detected_frameworks:
        rec_lines = []
        for fw in detected_frameworks:
            pkg = TRUST_LAYER_MAP.get(fw, None)
            if pkg:
                rec_lines.append(f"  pip install {pkg}")
        if rec_lines:
            console.print(f"\n[bold yellow]Detected frameworks:[/] {', '.join(detected_frameworks)}")
            console.print("[dim]Add trust layers for runtime compliance:[/]")
            for line in rec_lines:
                console.print(f"[bold green]{line}[/]")
    else:
        console.print(f"\n[dim]Add a trust layer for runtime compliance: pip install {rec_pkg}[/]")
    console.print(f"[dim]All 10 trust layer packages: https://airblackbox.ai[/]\n")

    # --- Crosswalk teaser ---
    if crosswalk_hint:
        console.print(f"[bold cyan]{crosswalk_hint}[/]\n")

    # --- Telemetry (anonymous, opt-out with AIR_BLACKBOX_TELEMETRY=off) ---
    try:
        from air_blackbox.telemetry import send_event
        import os as _os
        py_count = 0
        if _os.path.isfile(scan) and scan.endswith(".py"):
            py_count = 1
        else:
            for _root, _dirs, _files in _os.walk(scan):
                py_count += sum(1 for f in _files if f.endswith(".py"))
        all_checks = [c for a in articles for c in a.get("checks", [])]
        send_event(
            command="comply",
            python_files=py_count,
            checks_passing=sum(1 for c in all_checks if c.get("status") == "pass"),
            checks_warning=sum(1 for c in all_checks if c.get("status") == "warn"),
            checks_failing=sum(1 for c in all_checks if c.get("status") == "fail"),
            total_checks=len(all_checks),
            version=_ab_version,
        )
    except Exception:
        pass  # Telemetry should never break the tool


@main.command()
@click.option("--gateway", default="http://localhost:8080", help="Gateway URL")
@click.option("--runs-dir", default=None, help="Path to .air.json records")
@click.option("--approved", default=None, help="Path to approved models YAML")
@click.option("--format", "fmt", type=click.Choice(["table", "cyclonedx", "json"]), default="table")
@click.option("--output", "-o", default=None, help="Output file path")
@click.option("--init-registry", is_flag=True, help="Generate approved-models.yaml from current traffic")
def discover(gateway, runs_dir, approved, fmt, output, init_registry):
    """Discover AI models, tools, and services in your environment."""
    from air_blackbox.gateway_client import GatewayClient
    from air_blackbox.aibom.generator import generate_aibom
    from air_blackbox.aibom.shadow import detect_shadow_ai, generate_approved_registry
    import json as jsonlib

    console.print("\n[bold blue]AIR Blackbox[/] - AI Discovery & Inventory\n")
    with console.status("[bold green]Scanning environment..."):
        client = GatewayClient(gateway_url=gateway, runs_dir=runs_dir)
        status = client.get_status()
    if status.total_runs == 0 and not status.reachable:
        console.print("[yellow]No traffic data found.[/] Start gateway and route AI traffic through it.\n")
        return

    # Generate approved registry if requested
    if init_registry:
        registry = generate_approved_registry(status)
        reg_path = "approved-models.json"
        with open(reg_path, "w") as f:
            jsonlib.dump(registry, f, indent=2)
        console.print(f"  [green]✓[/] Generated [bold]{reg_path}[/] with {len(registry['models'])} models, {len(registry['providers'])} providers")
        console.print(f"  [dim]Future runs of discover will flag anything not in this list.[/]\n")
        return

    # CycloneDX output
    if fmt == "cyclonedx" or fmt == "json":
        bom = generate_aibom(status)
        bom_json = jsonlib.dumps(bom, indent=2)
        if output:
            with open(output, "w") as f:
                f.write(bom_json)
            console.print(f"  [green]✓[/] AI-BOM written to [bold]{output}[/]")
            console.print(f"  [dim]{len(bom['components'])} components, CycloneDX 1.6[/]\n")
        else:
            click.echo(bom_json)
        return

    # Table output (default)
    console.print(f"  Total logged events: [bold]{status.total_runs:,}[/]")
    console.print(f"  Period: {status.date_range_start or 'N/A'} → {status.date_range_end or 'N/A'}")
    console.print(f"  Total tokens: [bold]{status.total_tokens:,}[/]\n")

    # Models table
    if status.models_observed:
        t = Table(title="Models Detected", show_header=True, header_style="bold white on dark_blue")
        t.add_column("Model", style="bold", width=25)
        t.add_column("Provider", width=12)
        t.add_column("Status", justify="center", width=14)
        for m in status.models_observed:
            from air_blackbox.aibom.generator import _guess_provider
            t.add_row(m, _guess_provider(m), "[green]✅ Observed[/]")
        console.print(t)
        console.print()

    # Providers table
    if status.providers_observed:
        t = Table(title="API Providers", show_header=True, header_style="bold white on dark_blue")
        t.add_column("Provider", style="bold")
        t.add_column("Status", justify="center")
        for p in status.providers_observed:
            t.add_row(p, "[green]✅ Active[/]")
        console.print(t)
        console.print()

    # Tools table
    tools = set()
    for r in status.recent_runs:
        for tc in r.get("tool_calls", []):
            if tc: tools.add(tc)
    if tools:
        t = Table(title="Agent Tools Detected", show_header=True, header_style="bold white on dark_blue")
        t.add_column("Tool", style="bold")
        t.add_column("Status", justify="center")
        for tool in sorted(tools):
            t.add_row(tool, "[green]✅ Observed[/]")
        console.print(t)
        console.print()

    # Shadow AI alerts
    alerts = detect_shadow_ai(status, approved)
    if alerts:
        t = Table(title="Shadow AI Alerts", show_header=True, header_style="bold white on red")
        t.add_column("Model", style="bold", width=20)
        t.add_column("Severity", justify="center", width=10)
        t.add_column("Reason", width=50)
        for a in alerts:
            sev_color = {"high": "red", "medium": "yellow", "low": "dim"}.get(a.severity, "white")
            t.add_row(a.model, f"[{sev_color}]{a.severity.upper()}[/{sev_color}]", a.reason)
        console.print(t)
        console.print()

    # Summary
    bom = generate_aibom(status)
    console.print(Panel(
        f"[bold]{len(bom['components'])}[/] components inventoried: "
        f"{len(status.models_observed)} models, {len(status.providers_observed)} providers, {len(tools)} tools\n\n"
        f"[green]air-blackbox discover --format=cyclonedx -o aibom.json[/]  Export full AI-BOM\n"
        f"[green]air-blackbox discover --init-registry[/]                    Create approved models list\n"
        f"[green]air-blackbox discover --approved=approved-models.json[/]    Check against approved list",
        title="[bold blue]AI-BOM Summary[/]",
        border_style="blue",
    ))

    # --- Telemetry ---
    try:
        from air_blackbox.telemetry import send_event
        send_event(command="discover", version=_ab_version)
    except Exception:
        pass


@main.command()
@click.option("--gateway", default="http://localhost:8080", help="Gateway URL")
@click.option("--runs-dir", default=None, help="Path to .air.json records")
@click.option("--episode", default=None, help="Episode ID to replay")
@click.option("--last", default=10, help="Show last N runs")
@click.option("--verify", is_flag=True, help="Verify HMAC audit chain")
def replay(gateway, runs_dir, episode, last, verify):
    """Reconstruct AI incidents from the audit chain."""
    from air_blackbox.replay.engine import ReplayEngine

    console.print("\n[bold blue]AIR Blackbox[/] - Incident Replay\n")

    with console.status("[bold green]Loading audit records..."):
        engine = ReplayEngine(runs_dir=runs_dir or "./runs")
        count = engine.load()

    if count == 0:
        console.print("[yellow]No audit records found.[/] Run 'air-blackbox demo' or route traffic through gateway.\n")
        return

    # Verify chain if requested
    if verify:
        console.print("[bold]Verifying HMAC audit chain...[/]\n")
        result = engine.verify_chain()
        if result.intact:
            console.print(f"  [green]✅ CHAIN INTACT[/] - {result.verified_records:,} records verified. No tampering detected.\n")
        else:
            console.print(f"  [red]❌ CHAIN BROKEN[/] at record {result.first_break_at} (run: {result.first_break_run_id})")
            console.print(f"  [red]  {result.verified_records} of {result.total_records} records verified before break.[/]\n")
        return

    # Detail view for single episode
    if episode:
        rec = engine.get_run(episode)
        if not rec:
            console.print(f"[red]Run '{episode}' not found.[/] Use 'air-blackbox replay' to see all runs.\n")
            return
        console.print(f"  [bold]Run Detail: {rec.run_id}[/]\n")
        console.print(f"  Model:     {rec.model}")
        console.print(f"  Provider:  {rec.provider}")
        console.print(f"  Timestamp: {rec.timestamp}")
        console.print(f"  Duration:  {rec.duration_ms}ms")
        console.print(f"  Tokens:    {rec.tokens}")
        console.print(f"  Status:    {'[green]success[/]' if rec.status == 'success' else '[red]' + rec.status + '[/]'}")
        if rec.tool_calls:
            console.print(f"  Tools:     {', '.join(rec.tool_calls)}")
        if rec.pii_alerts:
            console.print(f"  [yellow]PII Alerts:  {len(rec.pii_alerts)} detected[/]")
        if rec.injection_alerts:
            console.print(f"  [red]Injection:   {len(rec.injection_alerts)} detected[/]")
        if rec.error:
            console.print(f"  [red]Error:       {rec.error}[/]")
        console.print()
        return

    # Stats summary
    stats = engine.get_stats()
    console.print(f"  [bold]{stats['total_records']:,}[/] total records")
    if stats.get("date_range"):
        console.print(f"  Period: {stats['date_range'][0]} → {stats['date_range'][1]}")
    console.print(f"  Total tokens: {stats['total_tokens']:,} | Avg latency: {stats['avg_duration_ms']}ms")
    if stats["pii_alerts"] > 0:
        console.print(f"  [yellow]PII alerts: {stats['pii_alerts']}[/]")
    if stats["injection_alerts"] > 0:
        console.print(f"  [red]Injection attempts: {stats['injection_alerts']}[/]")
    console.print()

    # Runs table
    records = engine.records[-last:]
    records.reverse()
    t = Table(title=f"Last {len(records)} Runs", show_header=True, header_style="bold white on dark_blue")
    t.add_column("Run ID", width=20)
    t.add_column("Model", width=15)
    t.add_column("Tokens", justify="right", width=8)
    t.add_column("Latency", justify="right", width=8)
    t.add_column("Status", justify="center", width=10)
    t.add_column("Timestamp", width=22)
    for rec in records:
        st = "[green]✅[/]" if rec.status == "success" else "[red]❌[/]"
        t.add_row(rec.run_id[:20], rec.model, str(rec.tokens.get("total", 0)),
                  f"{rec.duration_ms}ms", st, rec.timestamp[:22])
    console.print(t)
    console.print()
    console.print("[dim]Detail view: air-blackbox replay --episode=<run_id>[/]")
    console.print("[dim]Verify chain: air-blackbox replay --verify[/]\n")

    # --- Telemetry ---
    try:
        from air_blackbox.telemetry import send_event
        send_event(command="replay", version=_ab_version)
    except Exception:
        pass


@main.command()
@click.option("--gateway",  default="http://localhost:8080", help="Gateway URL")
@click.option("--runs-dir", default=None, help="Path to .air.json records")
@click.option("--scan",     default=".", help="Path to scan for code-level checks")
@click.option("--range",    "time_range", default="30d", help="Time range")
@click.option("--format",   "fmt", type=click.Choice(["json", "pdf", "evidence"]), default="json")
@click.option("--signing-key", default=None, help="HMAC signing key for the audit chain (or set TRUST_SIGNING_KEY)")
@click.option("--output",   "-o", default=None, help="Output file path")
def export(gateway, runs_dir, scan, time_range, fmt, signing_key, output):
    """Generate signed evidence bundles for auditors and insurers.

    \b
    Formats:
        json      - machine-readable signed evidence bundle (default)
        pdf       - formatted PDF compliance report for humans / auditors
        evidence  - self-verifying .air-evidence.zip with a standalone verify.py
                    an auditor runs offline (stdlib only) to get PASS/FAIL

    \b
    Examples:
        air-blackbox export
        air-blackbox export --format pdf
        air-blackbox export --format evidence --signing-key "$TRUST_SIGNING_KEY"
        air-blackbox export --scan ~/myproject --format pdf
    """
    from air_blackbox.export.bundle import generate_evidence_bundle
    import json as jsonlib

    console.print("\n[bold cyan]AIR Blackbox[/] - Evidence Export\n")

    if fmt == "evidence":
        import os as _os
        from air_blackbox.export.evidence_bundle import generate_evidence_zip
        from air_blackbox.replay.engine import ReplayEngine
        from air_blackbox.compliance.engine import run_all_checks
        from air_blackbox.gateway_client import GatewayClient

        key = signing_key or _os.environ.get("TRUST_SIGNING_KEY", "air-blackbox-default")
        rdir = runs_dir or "./runs"

        engine = ReplayEngine(runs_dir=rdir)
        engine.load()
        records = list(getattr(engine, "_raw_records", []))

        client = GatewayClient(gateway_url=gateway, runs_dir=rdir, scan_path=scan)
        _ra = run_all_checks(client.get_status(), scan)
        scan_results = _ra[0] if isinstance(_ra, tuple) else _ra

        zip_path = generate_evidence_zip(
            chain_entries=records,
            scan_results=scan_results,
            signing_key=key,
            output_dir=output or ".",
        )

        console.print(f"  [bold]Audit records:[/]  {len(records)}")
        console.print(f"  [bold]Signing key:[/]    {'(default)' if key == 'air-blackbox-default' else '(custom)'}")
        console.print()
        console.print(Panel(
            f"Self-verifying bundle written to [bold]{zip_path}[/]\n\n"
            f"Contains: audit_chain.json + scan_results.json + bundle_meta.json + verify.py\n"
            f"An auditor extracts it and runs (no install needed):\n"
            f"  [bold]python verify.py --key <your-signing-key>[/]\n"
            f"and gets PASS or FAIL. Verification uses only the Python standard library.",
            title="[bold green]Evidence Bundle Complete[/]",
            border_style="green",
        ))
        try:
            from air_blackbox.telemetry import send_event
            send_event(command="export", version=_ab_version)
        except Exception:
            pass
        return

    with console.status("[bold green]Generating evidence bundle..."):
        bundle = generate_evidence_bundle(gateway_url=gateway, runs_dir=runs_dir, scan_path=scan)

    summary = bundle.get("compliance", {}).get("summary", {})
    trail   = bundle.get("audit_trail", {})
    chain   = trail.get("chain_verification", {})

    passing  = summary.get("passing",  0)
    warnings = summary.get("warnings", 0)
    failing  = summary.get("failing",  0)

    console.print(f"  [bold]Compliance:[/]   {passing} passing · {warnings} warnings · {failing} failing")
    console.print(f"  [bold]AI-BOM:[/]        {len(bundle.get('aibom', {}).get('components', []))} components")
    console.print(f"  [bold]Audit trail:[/]   {trail.get('total_records', 0)} records")
    console.print(f"  [bold]Chain:[/]         {'[green]INTACT[/]' if chain.get('intact') else '[yellow]No signing key set[/]'}")
    console.print()

    if fmt == "pdf":
        from air_blackbox.export.pdf_report import generate_pdf, REPORTLAB_OK
        if not REPORTLAB_OK:
            console.print("[red]reportlab not installed.[/] Run: [bold]pip install reportlab[/]")
            raise SystemExit(1)

        out_path = output or f"AIR_Blackbox_Compliance_Report_{datetime.utcnow().strftime('%Y%m%d')}.pdf"
        with console.status("[bold green]Rendering PDF report..."):
            generate_pdf(bundle, out_path)

        console.print(Panel(
            f"PDF report written to [bold]{out_path}[/]\n\n"
            f"Contains: compliance scorecard · per-article findings · audit trail · priority fix list\n"
            f"Ready to hand to your auditor, compliance team, or share with stakeholders.",
            title="[bold green]PDF Export Complete[/]",
            border_style="green",
        ))
    else:
        # Default: JSON evidence bundle
        out_path = output or f"air-blackbox-evidence-{datetime.utcnow().strftime('%Y%m%d')}.json"
        with open(out_path, "w") as f:
            jsonlib.dump(bundle, f, indent=2)

        console.print(Panel(
            f"Evidence bundle written to [bold]{out_path}[/]\n\n"
            f"Contains: compliance scan + AI-BOM (CycloneDX) + audit trail + HMAC attestation\n"
            f"Hand this file to your auditor or insurer as a single verifiable document.\n\n"
            f"[dim]Tip: use [bold]--format pdf[/bold] to generate a human-readable PDF report[/dim]",
            title="[bold green]Export Complete[/]",
            border_style="green",
        ))

    # --- Telemetry ---
    try:
        from air_blackbox.telemetry import send_event
        send_event(command="export", version=_ab_version)
    except Exception:
        pass


@main.command()
@click.option("--output", "-o", default=".", help="Directory to create demo data in")
def demo(output):
    """Run a zero-config demo - generates sample data and shows compliance.

    Creates sample .air.json records and compliance doc templates so you
    can experience the full tool without Docker or a running gateway.

    \b
    Try it:
        air-blackbox demo
        air-blackbox comply -v
        air-blackbox discover
        air-blackbox replay
    """
    from air_blackbox.demo_generator import generate_demo_data
    from air_blackbox.gateway_client import GatewayClient
    from air_blackbox.compliance.engine import run_all_checks
    import time

    console.print("\n[bold blue]AIR Blackbox[/] - Zero-Config Demo\n")
    console.print("[dim]Generating sample AI agent traffic...[/]\n")
    time.sleep(0.5)

    # Generate sample data
    result = generate_demo_data(output)

    console.print(f"  [green]✓[/] Created [bold]{result['runs_created']}[/] sample .air.json records in [bold]{result['runs_dir']}[/]")
    console.print(f"  [green]✓[/] Models: {', '.join(result['models'])}")
    console.print(f"  [green]✓[/] Providers: {', '.join(result['providers'])}")
    console.print(f"  [green]✓[/] Total tokens: {result['total_tokens']:,}")
    console.print(f"  [green]✓[/] Generated RISK_ASSESSMENT.md template")
    console.print(f"  [green]✓[/] Generated DATA_GOVERNANCE.md template")
    console.print(f"  [green]✓[/] Generated sample_agent.py (with good + bad patterns)")
    console.print()

    time.sleep(0.3)
    console.print("[dim]Running compliance check against demo data...[/]\n")

    # Now run compliance against the generated data
    client = GatewayClient(runs_dir=result["runs_dir"])
    status = client.get_status()

    if status.reachable:
        console.print(f"  [green]●[/] Gateway detected at {status.url}")
    else:
        console.print(f"  [yellow]●[/] No gateway running (offline mode - using .air.json records)")

    console.print(f"  [green]●[/] [bold]{status.total_runs}[/] events loaded")
    console.print()

    articles, _, _, _ = run_all_checks(status, output)

    for article in articles:
        for check in article["checks"]:
            icon = {"pass": "✅", "warn": "⚠️ ", "fail": "❌"}.get(check["status"], "?")
            tier = check.get("tier", "static")
            tier_tag = "[green]S[/]" if tier == "static" else "[blue]R[/]"
            det = {"auto": "[green]AUTO[/]", "hybrid": "[yellow]HYBR[/]", "manual": "[red]MANU[/]"}.get(check.get("detection", ""), "")
            console.print(f"  {icon} Art. {article['number']:>2} {tier_tag} {det} {check['name']}")

    total = sum(len(a["checks"]) for a in articles)
    passing = sum(1 for a in articles for c in a["checks"] if c["status"] == "pass")

    console.print(f"\n  [bold]{passing}/{total}[/] checks passing")
    console.print()
    console.print(Panel(
        "[bold]What just happened:[/]\n\n"
        "1. Generated 10 sample AI agent records (like a real agent would create)\n"
        "2. Created EU AI Act compliance doc templates (Articles 9 + 10)\n"
        "3. Ran compliance check against the sample data\n\n"
        "[bold]Try these next:[/]\n\n"
        "  [green]air-blackbox comply -v[/]     Full compliance with fix hints\n"
        "  [green]air-blackbox discover[/]      See models and providers detected\n"
        "  [green]air-blackbox replay[/]        See the audit trail timeline\n"
        "  [green]docker compose up[/]          Start full gateway for live traffic",
        title="[bold blue]Demo Complete[/]",
        border_style="blue",
    ))

    # --- Telemetry ---
    try:
        from air_blackbox.telemetry import send_event
        send_event(command="demo", version=_ab_version)
    except Exception:
        pass


if __name__ == "__main__":
    main()


@main.command()
@click.option("--output", "-o", default=".", help="Directory to initialize")
def init(output):
    """Initialize a project for AIR Blackbox compliance.

    Creates compliance doc templates and a .air-blackbox.yaml config file.
    """
    from air_blackbox.demo_generator import _RISK_TEMPLATE, _DATA_GOV_TEMPLATE
    import os

    console.print("\n[bold blue]AIR Blackbox[/] - Project Init\n")

    files_created = []
    for fname, content in [
        ("RISK_ASSESSMENT.md", _RISK_TEMPLATE),
        ("DATA_GOVERNANCE.md", _DATA_GOV_TEMPLATE),
    ]:
        fpath = os.path.join(output, fname)
        if not os.path.exists(fpath):
            with open(fpath, "w") as f:
                f.write(content)
            files_created.append(fname)
            console.print(f"  [green]✓[/] Created {fname}")
        else:
            console.print(f"  [dim]⏭  {fname} already exists[/]")

    if files_created:
        console.print(f"\n  [bold]{len(files_created)}[/] files created. Run [green]air-blackbox comply -v[/] to check status.\n")
    else:
        console.print("\n  All files already exist. Run [green]air-blackbox comply -v[/] to check status.\n")


@main.command()
@click.option("--tool", default=None, help="Tool name to validate")
@click.option("--args", "arguments", default=None, help="Tool arguments as JSON")
@click.option("--content", default=None, help="LLM output content to validate")
@click.option("--allowlist", default=None, help="Comma-separated list of approved tools")
def validate(tool, arguments, content, allowlist):
    """Validate an agent action BEFORE execution.

    Pre-execution runtime certification - proves the output was
    checked against rules before it was acted on.

    \b
    Examples:
        air-blackbox validate --tool=db_query --args='{"query":"SELECT * FROM users"}'
        air-blackbox validate --content="Here is the result..."
        air-blackbox validate --tool=web_search --allowlist=web_search,calculator
    """
    from air_blackbox.validate import RuntimeValidator, ToolAllowlistRule
    import json as jsonlib

    console.print("\n[bold blue]AIR Blackbox[/] - Runtime Validation\n")

    validator = RuntimeValidator()

    if allowlist:
        validator.add_rule(ToolAllowlistRule(allowlist.split(",")))

    action = {}
    action_type = "tool_call"
    if tool:
        action["tool_name"] = tool
    if arguments:
        try:
            action["arguments"] = jsonlib.loads(arguments)
        except jsonlib.JSONDecodeError:
            action["arguments"] = {"raw": arguments}
    if content:
        action["content"] = content
        action_type = "llm_response"

    report = validator.validate(action, action_type=action_type)

    t = Table(title="Validation Results", show_header=True, header_style="bold white on dark_blue")
    t.add_column("Rule", style="bold", width=22)
    t.add_column("Result", justify="center", width=10)
    t.add_column("Severity", justify="center", width=10)
    t.add_column("Message", width=45)

    for r in report.results:
        icon = "[green]✅ PASS[/]" if r.passed else "[red]❌ FAIL[/]"
        sev = {"block": "[red]BLOCK[/]", "warn": "[yellow]WARN[/]", "info": "[dim]INFO[/]"}.get(r.severity, r.severity)
        t.add_row(r.rule_name, icon, sev, r.message)
    console.print(t)
    console.print()

    if report.passed:
        console.print(f"  [green]✅ VALIDATED[/] - action approved for execution ({report.validated_in_ms}ms)")
    else:
        console.print(f"  [red]❌ BLOCKED[/] - action failed validation ({report.validated_in_ms}ms)")
    console.print(f"  [dim]Validation record: {report.action_id}.air.json[/]\n")


@main.command()
@click.option("--path", default=None, help="Filter by scan path")
@click.option("--compare", is_flag=True, help="Compare last two scans and show diff")
@click.option("--export", "export_path", default=None, help="Export history to JSON file")
@click.option("--limit", default=20, help="Max number of scans to show")
def history(path, compare, export_path, limit):
    """View compliance scan history and trends.

    Tracks every scan in a local SQLite database (~/.air-blackbox/compliance.db).
    Shows score trends over time, diffs between scans, and exports for reporting.

    \b
    Examples:
        air-blackbox history                    # Show recent scans
        air-blackbox history --compare          # Diff last two scans
        air-blackbox history --export report.json
        air-blackbox history --path ./my-project
    """
    from air_blackbox.compliance.history import (
        get_history, get_last_scan, get_scan_findings,
        compare_scans, export_history,
    )
    import json as jsonlib

    console.print("\n[bold blue]AIR Blackbox[/] - Compliance History\n")

    # Export mode
    if export_path:
        data = export_history(scan_path=path, limit=limit)
        with open(export_path, "w") as f:
            jsonlib.dump(data, f, indent=2)
        console.print(f"  [green]✓[/] Exported {data['scan_count']} scan(s) to [bold]{export_path}[/]\n")
        return

    # Compare mode
    if compare:
        scans = get_history(scan_path=path, limit=2)
        if len(scans) < 2:
            console.print("  [yellow]Need at least 2 scans to compare.[/] Run `air-blackbox comply` first.\n")
            return
        newer = scans[0]
        older = scans[1]
        diff = compare_scans(older["id"], newer["id"])

        console.print(f"  Comparing scan #{older['id']} → #{newer['id']}\n")
        console.print(f"  Score: [bold]{older['score_percent']}%[/] → [bold]{newer['score_percent']}%[/]", end="")
        delta = newer["score_percent"] - older["score_percent"]
        if delta > 0:
            console.print(f"  [bold green](+{delta}%)[/]")
        elif delta < 0:
            console.print(f"  [bold red]({delta}%)[/]")
        else:
            console.print(f"  [dim](no change)[/]")
        console.print()

        if diff["improved"]:
            console.print(f"  [bold green]Improved ({len(diff['improved'])}):[/]")
            for item in diff["improved"]:
                console.print(f"    [green]↑[/] Art {item['article']}: {item['name']} ({item['was']} → {item['now']})")
        if diff["regressed"]:
            console.print(f"\n  [bold red]Regressed ({len(diff['regressed'])}):[/]")
            for item in diff["regressed"]:
                console.print(f"    [red]↓[/] Art {item['article']}: {item['name']} ({item['was']} → {item['now']})")
        if diff["new_checks"]:
            console.print(f"\n  [bold blue]New checks ({len(diff['new_checks'])}):[/]")
            for item in diff["new_checks"]:
                si = {"pass": "[green]pass[/]", "warn": "[yellow]warn[/]", "fail": "[red]fail[/]"}.get(item["status"], item["status"])
                console.print(f"    [blue]●[/] Art {item['article']}: {item['name']} ({si})")
        if not diff["improved"] and not diff["regressed"] and not diff["new_checks"]:
            console.print("  [dim]No changes between scans.[/]")
        console.print()
        return

    # Default: show history trend
    scans = get_history(scan_path=path, limit=limit)
    if not scans:
        console.print("  [yellow]No scan history found.[/] Run `air-blackbox comply` to start tracking.\n")
        return

    table = Table(title="Compliance Scan History", show_header=True,
                  header_style="bold white on dark_blue", title_style="bold")
    table.add_column("#", width=4, justify="right")
    table.add_column("Date", width=20)
    table.add_column("Path", width=20)
    table.add_column("Score", width=10, justify="center")
    table.add_column("Pass", width=6, justify="center")
    table.add_column("Warn", width=6, justify="center")
    table.add_column("Fail", width=6, justify="center")
    table.add_column("Static", width=10, justify="center")
    table.add_column("Runtime", width=10, justify="center")
    table.add_column("Deep", width=5, justify="center")

    for s in scans:
        score_pct = s["score_percent"]
        score_color = "green" if score_pct >= 70 else "yellow" if score_pct >= 40 else "red"
        ts = s["timestamp"][:16].replace("T", " ")
        scan_path_short = s["scan_path"]
        if len(scan_path_short) > 20:
            scan_path_short = "..." + scan_path_short[-17:]
        deep_icon = "✓" if s["deep_scan"] else ""
        table.add_row(
            str(s["id"]),
            ts,
            scan_path_short,
            f"[{score_color}]{score_pct}%[/{score_color}]",
            f"[green]{s['passing']}[/]",
            f"[yellow]{s['warnings']}[/]",
            f"[red]{s['failing']}[/]",
            f"{s['static_passing']}/{s['static_total']}",
            f"{s['runtime_passing']}/{s['runtime_total']}",
            deep_icon,
        )

    console.print(table)

    # Trend line (last 5 scans)
    if len(scans) >= 2:
        trend = scans[:5][::-1]  # oldest first
        trend_str = " → ".join(f"{s['score_percent']}%" for s in trend)
        latest = scans[0]["score_percent"]
        oldest = trend[0]["score_percent"]
        delta = latest - oldest
        if delta > 0:
            console.print(f"\n  [bold]Trend:[/] {trend_str}  [bold green](+{delta}% overall)[/]")
        elif delta < 0:
            console.print(f"\n  [bold]Trend:[/] {trend_str}  [bold red]({delta}% overall)[/]")
        else:
            console.print(f"\n  [bold]Trend:[/] {trend_str}  [dim](flat)[/]")

    console.print(f"\n  [dim]Run with --compare to diff last two scans, or --export to save as JSON[/]\n")


@main.command()
@click.option("--gateway", default="http://localhost:8080", help="Gateway URL")
@click.option("--verbose", "-v", is_flag=True, help="Show detailed output for each test")
def test(gateway, verbose):
    """Run end-to-end validation of the AIR Blackbox stack.

    Tests every subsystem - validation engine, compliance engine,
    audit records, HMAC chain, and optionally the live gateway.

    \b
    Examples:
        air-blackbox test              # Test SDK (no gateway needed)
        air-blackbox test -v           # Verbose output
        air-blackbox test --gateway http://localhost:8080  # Include gateway tests
    """
    import time
    import json as jsonlib
    import tempfile
    import os

    console.print("\n[bold blue]AIR Blackbox[/] - Stack Validation Test\n")

    results = []
    start_time = time.time()

    def _run_test(name, fn):
        """Run a single test, catch exceptions, record result."""
        try:
            passed, detail = fn()
            results.append({"name": name, "passed": passed, "detail": detail})
            icon = "[green]✅[/]" if passed else "[red]❌[/]"
            console.print(f"  {icon} {name}")
            if verbose and detail:
                console.print(f"     [dim]{detail}[/]")
        except Exception as e:
            results.append({"name": name, "passed": False, "detail": str(e)})
            console.print(f"  [red]❌[/] {name}")
            if verbose:
                console.print(f"     [red]{str(e)[:120]}[/]")

    # ── Test 1: SDK imports ──────────────────────────────────────────
    def test_sdk_imports():
        from air_blackbox.validate import RuntimeValidator, ToolAllowlistRule
        from air_blackbox.validate import ContentPolicyRule, PiiOutputRule
        from air_blackbox.compliance.engine import run_all_checks
        from air_blackbox.gateway_client import GatewayClient, GatewayStatus
        from air_blackbox.aibom.generator import generate_aibom
        from air_blackbox.replay.engine import ReplayEngine
        return True, "All core modules imported successfully"

    console.print("[bold]SDK Tests[/]\n")
    _run_test("SDK module imports", test_sdk_imports)

    # ── Test 2: Validation engine ────────────────────────────────────
    def test_validation_engine():
        from air_blackbox.validate import RuntimeValidator, ToolAllowlistRule
        with tempfile.TemporaryDirectory() as tmpdir:
            v = RuntimeValidator(runs_dir=tmpdir)
            v.add_rule(ToolAllowlistRule(["web_search", "calculator"]))
            # Should pass - tool is on allowlist
            r1 = v.validate({"tool_name": "web_search", "arguments": {"q": "hello"}})
            assert r1.passed, "Approved tool should pass"
            # Should fail - tool is NOT on allowlist
            r2 = v.validate({"tool_name": "exec_shell", "arguments": {"cmd": "rm -rf /"}})
            assert not r2.passed, "Blocked tool should fail"
            return True, f"2/2 validation scenarios correct ({r1.validated_in_ms + r2.validated_in_ms}ms)"

    _run_test("Validation engine (approve/block)", test_validation_engine)

    # ── Test 3: Content policy detection ─────────────────────────────
    def test_content_policy():
        from air_blackbox.validate import RuntimeValidator
        with tempfile.TemporaryDirectory() as tmpdir:
            v = RuntimeValidator(runs_dir=tmpdir)
            # Safe content should pass
            r1 = v.validate({"content": "The weather today is sunny."}, action_type="llm_response")
            assert r1.passed, "Safe content should pass"
            # Dangerous content should be blocked
            r2 = v.validate({"tool_name": "db", "arguments": {"query": "DROP TABLE users"}})
            assert not r2.passed, "SQL injection should be blocked"
            return True, "Safe content passed, dangerous content blocked"

    _run_test("Content policy (safe vs dangerous)", test_content_policy)

    # ── Test 4: PII detection ────────────────────────────────────────
    def test_pii_detection():
        from air_blackbox.validate import RuntimeValidator
        with tempfile.TemporaryDirectory() as tmpdir:
            v = RuntimeValidator(runs_dir=tmpdir)
            # Content with PII should warn
            r = v.validate({"content": "Contact john@example.com or SSN 123-45-6789"}, action_type="llm_response")
            pii_results = [x for x in r.results if x.rule_name == "pii_output_check"]
            assert len(pii_results) > 0, "PII rule should run"
            assert not pii_results[0].passed, "PII should be detected"
            types = pii_results[0].details.get("pii_types", [])
            assert "email" in types, "Email should be detected"
            assert "ssn" in types, "SSN should be detected"
            return True, f"Detected PII types: {', '.join(types)}"

    _run_test("PII detection (email, SSN)", test_pii_detection)

    # ── Test 5: Hallucination guard ──────────────────────────────────
    def test_hallucination_guard():
        from air_blackbox.validate import RuntimeValidator
        with tempfile.TemporaryDirectory() as tmpdir:
            v = RuntimeValidator(runs_dir=tmpdir)
            r = v.validate(
                {"content": "Visit https://www.fake.com/api for more info"},
                action_type="llm_response"
            )
            hal_results = [x for x in r.results if x.rule_name == "hallucination_guard"]
            assert len(hal_results) > 0, "Hallucination rule should run"
            assert not hal_results[0].passed, "Fake URL should be flagged"
            return True, "Suspicious URL detected and flagged"

    _run_test("Hallucination guard (fake URLs)", test_hallucination_guard)

    # ── Test 6: Audit record write/read ──────────────────────────────
    def test_audit_records():
        from air_blackbox.validate import RuntimeValidator
        with tempfile.TemporaryDirectory() as tmpdir:
            v = RuntimeValidator(runs_dir=tmpdir)
            report = v.validate({"tool_name": "test_tool", "arguments": {}})
            # Check that a .air.json file was written
            air_files = [f for f in os.listdir(tmpdir) if f.endswith(".air.json")]
            assert len(air_files) >= 1, "Should write at least 1 audit record"
            # Read it back and verify structure
            with open(os.path.join(tmpdir, air_files[0])) as f:
                record = jsonlib.load(f)
            assert record.get("type") == "validation", "Record type should be 'validation'"
            assert "run_id" in record, "Record should have run_id"
            assert "timestamp" in record, "Record should have timestamp"
            assert "checks" in record, "Record should have checks"
            return True, f"Wrote and verified {len(air_files)} audit record(s)"

    _run_test("Audit record write/read", test_audit_records)

    # ── Test 7: Compliance engine ────────────────────────────────────
    def test_compliance_engine():
        from air_blackbox.compliance.engine import run_all_checks
        from air_blackbox.gateway_client import GatewayStatus
        status = GatewayStatus(
            reachable=False, total_runs=5,
            models_observed=["gpt-4o"], providers_observed=["openai"],
            total_tokens=1000, date_range_start="2026-01-01", date_range_end="2026-03-13",
            recent_runs=[{"run_id": "test-1", "model": "gpt-4o", "timestamp": "2026-03-13", "status": "success"}]
        )
        with tempfile.TemporaryDirectory() as tmpdir:
            articles, _, _, _ = run_all_checks(status, tmpdir)
            assert len(articles) == 6, f"Should have 6 articles, got {len(articles)}"
            total_checks = sum(len(a["checks"]) for a in articles)
            assert total_checks > 0, "Should have checks"
            article_nums = [a["number"] for a in articles]
            assert article_nums == [9, 10, 11, 12, 14, 15], f"Wrong articles: {article_nums}"
            return True, f"6 articles, {total_checks} checks executed"

    _run_test("Compliance engine (Articles 9-15)", test_compliance_engine)

    # ── Test 8: AI-BOM generation ────────────────────────────────────
    def test_aibom_generation():
        from air_blackbox.aibom.generator import generate_aibom
        from air_blackbox.gateway_client import GatewayStatus
        status = GatewayStatus(
            total_runs=3, models_observed=["gpt-4o", "claude-3-opus"],
            providers_observed=["openai", "anthropic"], total_tokens=5000,
        )
        bom = generate_aibom(status)
        assert bom.get("bomFormat") == "CycloneDX", "Should be CycloneDX format"
        assert "components" in bom, "Should have components"
        assert len(bom["components"]) >= 2, "Should have at least 2 components"
        return True, f"CycloneDX BOM with {len(bom['components'])} components"

    _run_test("AI-BOM generation (CycloneDX)", test_aibom_generation)

    # ── Test 9: Replay engine ────────────────────────────────────────
    def test_replay_engine():
        from air_blackbox.replay.engine import ReplayEngine
        with tempfile.TemporaryDirectory() as tmpdir:
            # Write a sample .air.json record
            sample = {
                "version": "1.0.0", "run_id": "test-replay-1",
                "timestamp": "2026-03-13T10:00:00Z",
                "model": "gpt-4o", "provider": "openai",
                "tokens": {"prompt": 100, "completion": 50, "total": 150},
                "duration_ms": 234, "status": "success",
                "tool_calls": ["web_search"], "pii_alerts": [], "injection_alerts": [],
            }
            with open(os.path.join(tmpdir, "test-replay-1.air.json"), "w") as f:
                jsonlib.dump(sample, f)
            engine = ReplayEngine(runs_dir=tmpdir)
            count = engine.load()
            assert count >= 1, "Should load at least 1 record"
            stats = engine.get_stats()
            assert stats["total_records"] >= 1, "Should have records in stats"
            return True, f"Loaded {count} record(s), stats computed"

    _run_test("Replay engine (load + stats)", test_replay_engine)

    # ── Code Scanner Tests ───────────────────────────────────────────
    console.print(f"\n[bold]Code Scanner Tests[/]\n")

    def test_scanner_error_handling():
        """Verify scanner detects try/except around LLM calls."""
        from air_blackbox.compliance.code_scanner import scan_codebase
        with tempfile.TemporaryDirectory() as tmpdir:
            # File WITH error handling
            with open(os.path.join(tmpdir, "good.py"), "w") as f:
                f.write("from openai import OpenAI\nclient = OpenAI()\ntry:\n    client.chat.completions.create(model='gpt-4o')\nexcept Exception:\n    pass\n")
            findings = scan_codebase(tmpdir)
            eh = [f for f in findings if f.name == "LLM call error handling"]
            assert len(eh) == 1, "Should find error handling check"
            assert eh[0].status == "pass", f"Should pass, got {eh[0].status}"
            return True, "Error handling detection working"

    _run_test("Scanner: error handling detection", test_scanner_error_handling)

    def test_scanner_no_error_handling():
        """Verify scanner catches missing error handling."""
        from air_blackbox.compliance.code_scanner import scan_codebase
        with tempfile.TemporaryDirectory() as tmpdir:
            with open(os.path.join(tmpdir, "bad.py"), "w") as f:
                f.write("from openai import OpenAI\nclient = OpenAI()\nclient.chat.completions.create(model='gpt-4o')\n")
            findings = scan_codebase(tmpdir)
            eh = [f for f in findings if f.name == "LLM call error handling"]
            assert len(eh) == 1, "Should find error handling check"
            assert eh[0].status == "fail", f"Should fail without try/except, got {eh[0].status}"
            return True, "Missing error handling correctly flagged"

    _run_test("Scanner: missing error handling", test_scanner_no_error_handling)

    def test_scanner_tracing_patterns():
        """Verify scanner detects modern tracing (instrumentation, event bus, OTel)."""
        from air_blackbox.compliance.code_scanner import scan_codebase
        with tempfile.TemporaryDirectory() as tmpdir:
            with open(os.path.join(tmpdir, "traced.py"), "w") as f:
                f.write("from opentelemetry import trace\ntracer = trace.get_tracer(__name__)\nwith tracer.start_span('agent_call'):\n    pass\n")
            findings = scan_codebase(tmpdir)
            tr = [f for f in findings if f.name == "Tracing / observability"]
            assert len(tr) == 1 and tr[0].status == "pass", "OTel tracing should pass"
        # Test instrumentation (LlamaIndex pattern)
        with tempfile.TemporaryDirectory() as tmpdir:
            with open(os.path.join(tmpdir, "instrumented.py"), "w") as f:
                f.write("from llama_index.core.instrumentation import dispatcher\ndispatcher.add_event_handler(my_handler)\n")
            findings = scan_codebase(tmpdir)
            tr = [f for f in findings if f.name == "Tracing / observability"]
            assert len(tr) == 1 and tr[0].status == "pass", "Instrumentation pattern should pass"
        return True, "OTel + instrumentation patterns detected"

    _run_test("Scanner: tracing patterns (OTel, instrumentation)", test_scanner_tracing_patterns)

    def test_scanner_hitl_patterns():
        """Verify scanner detects HITL patterns from multiple frameworks."""
        from air_blackbox.compliance.code_scanner import scan_codebase
        patterns_to_test = [
            ("haystack_hitl.py", "confirmation_strategy = 'always_ask'\n"),
            ("crewai_hitl.py", "agent = Agent(allow_delegation=True)\n"),
            ("langgraph_hitl.py", "workflow.add_node('human', interrupt_before=['action'])\n"),
        ]
        for fname, content in patterns_to_test:
            with tempfile.TemporaryDirectory() as tmpdir:
                with open(os.path.join(tmpdir, fname), "w") as f:
                    f.write(content)
                findings = scan_codebase(tmpdir)
                hitl = [f for f in findings if f.name == "Human-in-the-loop patterns"]
                assert len(hitl) == 1 and hitl[0].status == "pass", f"HITL should pass for {fname}, got {hitl[0].status if hitl else 'none'}"
        return True, "Haystack, CrewAI, LangGraph HITL patterns all detected"

    _run_test("Scanner: HITL patterns (Haystack, CrewAI, LangGraph)", test_scanner_hitl_patterns)

    def test_scanner_injection_defense():
        """Verify scanner detects guardrail patterns from CrewAI and LlamaIndex."""
        from air_blackbox.compliance.code_scanner import scan_codebase
        with tempfile.TemporaryDirectory() as tmpdir:
            with open(os.path.join(tmpdir, "guarded.py"), "w") as f:
                f.write("from crewai import Agent\nagent = Agent(hallucination_guardrail=True)\n")
            findings = scan_codebase(tmpdir)
            inj = [f for f in findings if f.name == "Prompt injection defense"]
            assert len(inj) == 1 and inj[0].status == "pass", "CrewAI guardrail should pass"
        return True, "CrewAI hallucination_guardrail detected"

    _run_test("Scanner: injection defense (CrewAI guardrail)", test_scanner_injection_defense)

    def test_scanner_output_validation():
        """Verify scanner detects CrewAI output_pydantic and expected_output."""
        from air_blackbox.compliance.code_scanner import scan_codebase
        with tempfile.TemporaryDirectory() as tmpdir:
            with open(os.path.join(tmpdir, "tasks.py"), "w") as f:
                f.write("from crewai import Task\ntask = Task(description='Analyze', output_pydantic=AnalysisResult, expected_output='JSON report')\n")
            findings = scan_codebase(tmpdir)
            ov = [f for f in findings if f.name == "LLM output validation"]
            assert len(ov) == 1 and ov[0].status == "pass", "CrewAI output_pydantic should pass"
        return True, "CrewAI output_pydantic + expected_output detected"

    _run_test("Scanner: output validation (CrewAI patterns)", test_scanner_output_validation)

    def test_scanner_identity_binding():
        """Verify scanner detects CrewAI Fingerprint and Haystack memory patterns."""
        from air_blackbox.compliance.code_scanner import scan_codebase
        with tempfile.TemporaryDirectory() as tmpdir:
            with open(os.path.join(tmpdir, "identity.py"), "w") as f:
                f.write("from crewai.utilities import Fingerprint\nagent_id = Fingerprint()\ncard = AgentCard(name='worker')\n")
            findings = scan_codebase(tmpdir)
            ib = [f for f in findings if f.name == "Agent-to-user identity binding"]
            assert len(ib) == 1 and ib[0].status == "pass", "CrewAI Fingerprint should pass"
        return True, "CrewAI Fingerprint + AgentCard detected"

    _run_test("Scanner: identity binding (CrewAI Fingerprint)", test_scanner_identity_binding)

    def test_scanner_audit_trail_events():
        """Verify scanner detects CrewAI event bus patterns."""
        from air_blackbox.compliance.code_scanner import scan_codebase
        with tempfile.TemporaryDirectory() as tmpdir:
            with open(os.path.join(tmpdir, "events.py"), "w") as f:
                f.write("from crewai.utilities import agent_events, crew_events\ndef on_event(event): pass\nemit_event('task_completed', data)\n")
            findings = scan_codebase(tmpdir)
            at = [f for f in findings if f.name == "Agent action audit trail"]
            assert len(at) == 1 and at[0].status == "pass", "CrewAI event bus should pass"
        return True, "CrewAI event bus audit trail detected"

    _run_test("Scanner: audit trail (CrewAI event bus)", test_scanner_audit_trail_events)

    def test_scanner_false_positive_pii():
        """Regression: 'private' alone should NOT trigger PII detection."""
        from air_blackbox.compliance.code_scanner import scan_codebase
        with tempfile.TemporaryDirectory() as tmpdir:
            with open(os.path.join(tmpdir, "normal.py"), "w") as f:
                f.write("class MyClass:\n    def __init__(self):\n        self.private = True\n        self._private_method = lambda: None\n")
            findings = scan_codebase(tmpdir)
            pii = [f for f in findings if f.name == "PII handling in code"]
            # Should NOT pass - 'private' alone is not PII handling
            assert len(pii) == 1 and pii[0].status != "pass", f"Bare 'private' should not trigger PII pass, got {pii[0].status}"
        return True, "False positive: bare 'private' correctly ignored"

    _run_test("Scanner: PII false positive regression", test_scanner_false_positive_pii)

    def test_scanner_skip_deprecated():
        """Verify scanner skips deprecated and archived directories."""
        from air_blackbox.compliance.code_scanner import scan_codebase
        with tempfile.TemporaryDirectory() as tmpdir:
            # Create deprecated dir with Python files - should be skipped
            dep_dir = os.path.join(tmpdir, "deprecated")
            os.makedirs(dep_dir)
            with open(os.path.join(dep_dir, "old.py"), "w") as f:
                f.write("import logging\nlogging.getLogger(__name__)\n")
            findings = scan_codebase(tmpdir)
            # Should find no Python files (only deprecated dir has them)
            no_files = [f for f in findings if "No Python files" in f.evidence]
            assert len(no_files) >= 1, "Should skip deprecated directory"
        return True, "Deprecated directories correctly skipped"

    _run_test("Scanner: skip deprecated directories", test_scanner_skip_deprecated)

    # ── Two-Tier Scoring Tests ───────────────────────────────────────
    console.print(f"\n[bold]Two-Tier Scoring Tests[/]\n")

    def test_tier_labels():
        """Verify every check has a tier and gateway checks are labeled 'runtime'."""
        from air_blackbox.compliance.engine import run_all_checks
        from air_blackbox.gateway_client import GatewayStatus
        status = GatewayStatus(reachable=False, total_runs=0)
        with tempfile.TemporaryDirectory() as tmpdir:
            articles, _, _, _ = run_all_checks(status, tmpdir)
            all_checks = [c for a in articles for c in a["checks"]]
            # Every check must have a tier
            for c in all_checks:
                assert "tier" in c, f"Check '{c['name']}' missing tier"
                assert c["tier"] in ("static", "runtime"), f"Check '{c['name']}' has invalid tier: {c['tier']}"
            static = [c for c in all_checks if c["tier"] == "static"]
            runtime = [c for c in all_checks if c["tier"] == "runtime"]
            assert len(static) > 0, "Should have static checks"
            assert len(runtime) > 0, "Should have runtime checks"
            return True, f"{len(static)} static + {len(runtime)} runtime checks, all labeled"

    _run_test("Tier labels on all checks", test_tier_labels)

    def test_runtime_checks_identified():
        """Verify known gateway-dependent checks are tier='runtime'."""
        from air_blackbox.compliance.engine import run_all_checks
        from air_blackbox.gateway_client import GatewayStatus
        status = GatewayStatus(reachable=False, total_runs=0)
        with tempfile.TemporaryDirectory() as tmpdir:
            articles, _, _, _ = run_all_checks(status, tmpdir)
            all_checks = {c["name"]: c for a in articles for c in a["checks"]}
            runtime_expected = [
                "Risk mitigations active", "PII detection in prompts",
                "Data vault (controlled storage)", "Runtime system inventory (AI-BOM data)",
                "Automatic event logging", "Tamper-evident audit chain",
                "Log detail and traceability", "Kill switch / stop mechanism",
                "Prompt injection protection",
            ]
            for name in runtime_expected:
                assert name in all_checks, f"Check '{name}' not found"
                assert all_checks[name]["tier"] == "runtime", f"'{name}' should be runtime, got {all_checks[name]['tier']}"
            return True, f"All {len(runtime_expected)} gateway checks correctly labeled runtime"

    _run_test("Runtime checks correctly identified", test_runtime_checks_identified)

    def test_version_consistency():
        """Verify version is consistent across pyproject.toml, __init__.py, and cli."""
        import air_blackbox
        cli_version = _ab_version
        init_version = air_blackbox.__version__
        assert init_version == cli_version, f"__init__ ({init_version}) != cli ({cli_version})"
        return True, f"Version {init_version} consistent across modules"

    _run_test("Version consistency", test_version_consistency)

    # ── Gateway tests (optional) ─────────────────────────────────────
    console.print(f"\n[bold]Gateway Tests[/]\n")

    def test_gateway_health():
        from air_blackbox.gateway_client import GatewayClient
        client = GatewayClient(gateway_url=gateway)
        status = client.get_status()
        if status.reachable:
            return True, f"Gateway reachable at {gateway}"
        else:
            return False, f"Gateway not reachable at {gateway} (start with: docker compose up)"

    _run_test("Gateway connectivity", test_gateway_health)

    def test_gateway_audit_endpoint():
        import httpx
        try:
            r = httpx.get(f"{gateway}/v1/audit", timeout=5.0)
            if r.status_code == 200:
                data = r.json()
                chain = data.get("audit_chain", {})
                return True, f"Audit endpoint OK - chain length: {chain.get('length', 0)}, intact: {chain.get('intact', False)}"
            return False, f"Audit endpoint returned {r.status_code}"
        except Exception:
            return False, f"Audit endpoint not reachable (gateway may not be running)"

    _run_test("Gateway audit endpoint", test_gateway_audit_endpoint)

    def test_gateway_proxy():
        import httpx
        try:
            r = httpx.get(f"{gateway}/v1/models", timeout=5.0)
            if r.status_code == 200:
                data = r.json()
                models = [m.get("id", "?") for m in data.get("data", [])[:3]]
                return True, f"Proxy forwarding OK - models: {', '.join(models)}"
            elif r.status_code == 401:
                return True, "Proxy reached upstream (401 = API key needed, but proxy works)"
            return False, f"Proxy returned {r.status_code}"
        except Exception:
            return False, "Proxy not reachable (gateway may not be running)"

    _run_test("Gateway proxy forwarding", test_gateway_proxy)

    # ── Summary ──────────────────────────────────────────────────────
    elapsed = int((time.time() - start_time) * 1000)
    total = len(results)
    passed = sum(1 for r in results if r["passed"])
    failed = total - passed

    console.print()
    if failed == 0:
        console.print(Panel(
            f"[bold green]{passed}/{total}[/] tests passing in {elapsed}ms\n\n"
            f"Your AIR Blackbox stack is healthy.",
            title="[bold green]All Tests Passed[/]",
            border_style="green",
        ))
    else:
        # Separate SDK failures from gateway failures (gateway not running is expected)
        sdk_failures = [r for r in results if not r["passed"] and "Gateway" not in r["name"] and "audit endpoint" not in r["name"] and "Proxy" not in r["name"]]
        gw_failures = [r for r in results if not r["passed"] and r not in sdk_failures]

        if sdk_failures:
            lines = f"[bold red]{failed}[/] test(s) failed out of {total} ({elapsed}ms)\n"
            for r in sdk_failures:
                lines += f"\n  [red]●[/] {r['name']}: {r['detail']}"
            console.print(Panel(lines, title="[bold red]Tests Failed[/]", border_style="red"))
        else:
            console.print(Panel(
                f"[bold green]{passed}/{total}[/] tests passing in {elapsed}ms\n\n"
                f"SDK tests: [bold green]all passing[/]\n"
                f"Gateway tests: [bold yellow]{len(gw_failures)} skipped[/] (gateway not running)\n\n"
                f"[dim]Start gateway with: docker compose up[/]",
                title="[bold green]SDK Tests Passed[/]",
                border_style="green",
            ))
    console.print()


# ── Feedback command ─────────────────────────────────────────────────
@main.command()
@click.option("--scan-id", default=None, help="Scan ID from a previous comply run")
@click.option("--article", type=int, default=0, help="Article number (9, 10, 11, 12, 14, 15)")
@click.option("--false-positive", is_flag=True, help="Flag: the finding was a false positive")
@click.option("--correction", default=None, help="What the scanner should have said")
@click.option("--code-file", default=None, type=click.Path(exists=True), help="The scanned file")
@click.option("--interactive", is_flag=True, help="Interactive mode with prompts")
def feedback(scan_id, article, false_positive, correction, code_file, interactive):
    """Report a false positive or correction to improve the scanner.

    Your feedback gets written to training data so the next model retrain
    permanently fixes the issue. Every correction makes the scanner smarter.

    Examples:

        air-blackbox feedback --article 12 --false-positive \\
            --correction "structlog counts as audit logging" \\
            --code-file agent.py

        air-blackbox feedback --interactive
    """
    import os as _os
    from air_blackbox.feedback import accept_and_write

    console.print("\n[bold blue]AIR Blackbox[/] - Scanner Feedback\n")

    if interactive:
        scan_id = scan_id or click.prompt("Scan ID (or press Enter to skip)", default="manual")
        article = article or click.prompt("Article number", type=int)
        false_positive = false_positive or click.confirm("Was this a false positive?")
        correction = correction or click.prompt("What should the scanner have said?")
        if not code_file:
            code_file_input = click.prompt("Path to the scanned file (or Enter to skip)", default="")
            code_file = code_file_input if code_file_input else None

    if article not in (9, 10, 11, 12, 14, 15):
        console.print(f"[red]Invalid article: {article}. Must be 9, 10, 11, 12, 14, or 15.[/]")
        return

    code_snippet = ""
    if code_file:
        try:
            with open(code_file, "r", encoding="utf-8", errors="ignore") as f:
                code_snippet = f.read()
            if len(code_snippet) > 10000:
                code_snippet = code_snippet[:10000] + "\n# ... (truncated)"
        except Exception as e:
            console.print(f"[yellow]Could not read {code_file}: {e}[/]")

    feedback_text = correction or "False positive reported"
    severity = "high" if false_positive else "medium"

    try:
        result = accept_and_write(
            scan_id=scan_id or "manual",
            article=article,
            severity=severity,
            feedback_text=feedback_text,
            original_finding="",
            corrected_finding=correction or "",
            code_snippet=code_snippet,
            is_false_positive=false_positive,
        )

        console.print(f"  [green]Feedback recorded.[/] ID: [bold]{result['feedback_id']}[/]")
        console.print(f"  Training data: {result['output_path']}")
        console.print(f"  Total corrections: {result['total_in_file']}")
        console.print()
        console.print("  [dim]Run 'air-blackbox retrain status' to see pending corrections.[/]")
        console.print("  [dim]Run 'air-blackbox retrain all' to retrain the model.[/]")
        console.print()

    except ValueError as e:
        console.print(f"[red]Validation error: {e}[/]")


# ── Retrain command group ────────────────────────────────────────────
@main.group()
def retrain():
    """Retrain the compliance model from accumulated feedback.

    The retrain pipeline merges user corrections into the training corpus,
    fine-tunes Llama 3.2 1B via Unsloth + QLoRA, evaluates accuracy, and
    registers the updated model with Ollama.

    Subcommands:

        air-blackbox retrain status    Show pending feedback and model info

        air-blackbox retrain merge     Merge feedback into training data

        air-blackbox retrain run       Fine-tune and evaluate

        air-blackbox retrain publish   Register model with Ollama

        air-blackbox retrain all       Run full pipeline
    """
    pass


@retrain.command()
def status():
    """Show current model version, training data stats, and pending feedback."""
    import os as _os
    import shutil
    import subprocess
    import json as jsonlib

    console.print("\n[bold blue]AIR Blackbox[/] - Retrain Status\n")

    home_dir = _os.path.expanduser("~/.air-blackbox")
    feedback_path = _os.path.join(home_dir, "training_feedback.jsonl")

    # Count pending feedback.
    pending = 0
    fp_count = 0
    if _os.path.exists(feedback_path):
        with open(feedback_path) as f:
            for line in f:
                if line.strip():
                    pending += 1
                    try:
                        ex = jsonlib.loads(line)
                        output = ex.get("output", "")
                        if "NO ISSUE" in output or "false positive" in output.lower():
                            fp_count += 1
                    except jsonlib.JSONDecodeError:
                        pass

    console.print(f"  Pending feedback corrections: [bold]{pending}[/]")
    if fp_count:
        console.print(f"    False positive corrections: {fp_count}")
    console.print(f"  Feedback file: {feedback_path}")

    # Check Ollama model.
    if shutil.which("ollama"):
        try:
            proc = subprocess.run(
                ["ollama", "list"],
                capture_output=True,
                text=True,
                timeout=10,
            )
            models = [line for line in proc.stdout.split("\n") if "air-compliance" in line]
            if models:
                console.print(f"\n  [green]Installed models:[/]")
                for m in models:
                    console.print(f"    {m.strip()}")
            else:
                console.print(f"\n  [yellow]No air-compliance model installed.[/]")
                console.print(f"  [dim]Run 'air-blackbox setup' to install.[/]")
        except Exception:
            console.print(f"\n  [yellow]Could not query Ollama.[/]")
    else:
        console.print(f"\n  [yellow]Ollama not installed.[/]")

    console.print()


@retrain.command("merge")
@click.option("--canonical", default=None, help="Path to canonical training_data_v{N}.jsonl")
@click.option("--feedback", "feedback_path", default=None, help="Path to feedback JSONL")
@click.option("--output-dir", default=None, help="Directory for merged output")
@click.option("--min-count", default=5, help="Minimum feedback count before merging")
def retrain_merge(canonical, feedback_path, output_dir, min_count):
    """Merge accumulated feedback into the training corpus."""
    import os as _os
    from air_blackbox.retrain.merge import merge_training_data

    console.print("\n[bold blue]AIR Blackbox[/] - Merge Training Data\n")

    home_dir = _os.path.expanduser("~/.air-blackbox")
    feedback_path = feedback_path or _os.path.join(home_dir, "training_feedback.jsonl")
    output_dir = output_dir or _os.path.join(home_dir, "training")

    if not canonical:
        import glob as globmod
        candidates = sorted(globmod.glob(_os.path.join(output_dir, "training_data_v*.jsonl")))
        if not candidates:
            repo_dir = _os.path.join(_os.getcwd(), "training")
            candidates = sorted(globmod.glob(_os.path.join(repo_dir, "training_data_v*.jsonl")))
        if candidates:
            canonical = candidates[-1]
            console.print(f"  [dim]Using canonical: {canonical}[/]")
        else:
            console.print("[red]No canonical training data found. Specify with --canonical[/]")
            return

    report = merge_training_data(
        canonical_path=canonical,
        feedback_path=feedback_path,
        output_dir=output_dir,
        min_feedback_count=min_count,
    )

    if report.blocked:
        console.print(f"  [red]Merge blocked:[/] {report.block_reason}")
        return

    console.print(f"  Version: [bold]{report.version}[/]")
    console.print(f"  Canonical examples: {report.canonical_examples:,}")
    console.print(f"  New from feedback: [bold green]{report.new_from_feedback}[/]")
    console.print(f"  Duplicates skipped: {report.duplicates_skipped}")
    console.print(f"  Invalid skipped: {report.invalid_skipped}")
    console.print(f"  Total examples: [bold]{report.total_examples:,}[/]")
    console.print(f"  False positive fixes: {report.false_positive_corrections}")
    console.print(f"  Output: {report.output_path}")

    if report.warnings:
        for w in report.warnings:
            console.print(f"  [yellow]Warning: {w}[/]")

    console.print()


@retrain.command("run")
@click.option("--training-data", default=None, help="Path to merged training JSONL")
@click.option("--output-dir", default=None, help="Directory for model output")
@click.option("--epochs", default=3, help="Training epochs")
@click.option("--batch-size", default=4, help="Batch size")
def retrain_run(training_data, output_dir, epochs, batch_size):
    """Fine-tune the compliance model on merged training data."""
    import os as _os
    from air_blackbox.retrain.run import retrain as do_retrain, RetrainConfig

    console.print("\n[bold blue]AIR Blackbox[/] - Model Retraining\n")

    home_dir = _os.path.expanduser("~/.air-blackbox")
    output_dir = output_dir or _os.path.join(home_dir, "models")

    if not training_data:
        import glob as globmod
        training_dir = _os.path.join(home_dir, "training")
        candidates = sorted(globmod.glob(_os.path.join(training_dir, "training_data_v*.jsonl")))
        if candidates:
            training_data = candidates[-1]
            console.print(f"  [dim]Using: {training_data}[/]")
        else:
            console.print("[red]No training data found. Run 'air-blackbox retrain merge' first.[/]")
            return

    config = RetrainConfig(
        training_data=training_data,
        output_dir=output_dir,
        epochs=epochs,
        batch_size=batch_size,
    )

    console.print(f"  Base model: {config.base_model}")
    console.print(f"  Epochs: {config.epochs}")
    console.print(f"  Eval threshold: {config.eval_threshold:.0%}")
    console.print()

    with console.status("[bold green]Training model (this may take 1-2 hours)..."):
        result = do_retrain(config)

    if result.success:
        console.print(f"  [green]Training complete![/]")
        console.print(f"  Version: [bold]{result.version}[/]")
        console.print(f"  Loss: {result.training_loss:.4f}")
        console.print(f"  Accuracy: {result.eval_accuracy:.1%}")
        console.print(f"  False positive rate: {result.eval_false_positive_rate:.1%}")
        console.print(f"  GGUF: {result.gguf_path}")
        console.print(f"  Duration: {result.duration_seconds // 60}m {result.duration_seconds % 60}s")
    else:
        console.print(f"  [red]Training failed:[/] {result.error}")
        if result.model_path:
            console.print(f"  [dim]Failed model saved at: {result.model_path}[/]")

    console.print()


@retrain.command("publish")
@click.option("--gguf", default=None, help="Path to .gguf model file")
@click.option("--version", "version_str", default=None, help="Version string (e.g., v12)")
def retrain_publish(gguf, version_str):
    """Register the retrained model with Ollama."""
    import os as _os
    import re as re_mod
    from air_blackbox.retrain.publish import publish_model

    console.print("\n[bold blue]AIR Blackbox[/] - Publish Model\n")

    if not gguf:
        import glob as globmod
        home_dir = _os.path.expanduser("~/.air-blackbox/models")
        candidates = sorted(globmod.glob(_os.path.join(home_dir, "v*/gguf/*.gguf")))
        if candidates:
            gguf = candidates[-1]
            console.print(f"  [dim]Using: {gguf}[/]")
        else:
            console.print("[red]No GGUF model found. Run 'air-blackbox retrain run' first.[/]")
            return

    if not version_str:
        m = re_mod.search(r"v(\d+)", gguf)
        version_str = f"v{m.group(1)}" if m else "v0"

    result = publish_model(
        gguf_path=gguf,
        version=version_str,
        changelog_path=_os.path.join(_os.getcwd(), "CHANGELOG.md"),
        pyproject_path=_os.path.join(_os.getcwd(), "pyproject.toml"),
    )

    if result.success:
        console.print(f"  [green]Published![/] Model: [bold]{result.ollama_model}[/]")
        console.print(f"  Model card: {result.model_card_path}")
        console.print()
        console.print("  [dim]Users can now run 'air-blackbox setup' to get the updated model.[/]")
    else:
        console.print(f"  [red]Publish failed:[/] {result.error}")

    console.print()


@retrain.command("all")
@click.option("--canonical", default=None, help="Path to canonical training data")
@click.option("--min-count", default=5, help="Minimum feedback corrections before retraining")
def retrain_all(canonical, min_count):
    """Run the full retrain pipeline: merge, train, evaluate, publish.

    This is the one-command version. It stops and reports if any step fails.
    """
    import os as _os
    from air_blackbox.retrain.merge import merge_training_data
    from air_blackbox.retrain.run import retrain as do_retrain, RetrainConfig
    from air_blackbox.retrain.publish import publish_model

    console.print("\n[bold blue]AIR Blackbox[/] - Full Retrain Pipeline\n")

    home_dir = _os.path.expanduser("~/.air-blackbox")
    feedback_path = _os.path.join(home_dir, "training_feedback.jsonl")
    training_dir = _os.path.join(home_dir, "training")
    models_dir = _os.path.join(home_dir, "models")

    # Step 1: Merge.
    console.print("[bold]Step 1/3:[/] Merging feedback into training data...")

    if not canonical:
        import glob as globmod
        candidates = sorted(globmod.glob(_os.path.join(training_dir, "training_data_v*.jsonl")))
        if not candidates:
            repo_dir = _os.path.join(_os.getcwd(), "training")
            candidates = sorted(globmod.glob(_os.path.join(repo_dir, "training_data_v*.jsonl")))
        if candidates:
            canonical = candidates[-1]
        else:
            console.print("  [red]No canonical training data found.[/]")
            return

    merge_report = merge_training_data(
        canonical_path=canonical,
        feedback_path=feedback_path,
        output_dir=training_dir,
        min_feedback_count=min_count,
    )

    if merge_report.blocked:
        console.print(f"  [red]Blocked:[/] {merge_report.block_reason}")
        return

    console.print(f"  [green]Merged {merge_report.new_from_feedback} new corrections[/] ({merge_report.total_examples:,} total)")

    # Step 2: Retrain.
    console.print(f"\n[bold]Step 2/3:[/] Fine-tuning model ({merge_report.version})...")

    config = RetrainConfig(training_data=merge_report.output_path, output_dir=models_dir)

    with console.status("[bold green]Training (this may take 1-2 hours)..."):
        retrain_result = do_retrain(config)

    if not retrain_result.success:
        console.print(f"  [red]Training failed:[/] {retrain_result.error}")
        return

    console.print(f"  [green]Accuracy: {retrain_result.eval_accuracy:.1%}[/], FP rate: {retrain_result.eval_false_positive_rate:.1%}")

    # Step 3: Publish.
    console.print(f"\n[bold]Step 3/3:[/] Publishing to Ollama...")

    pub_result = publish_model(
        gguf_path=retrain_result.gguf_path,
        version=merge_report.version,
        merge_report=merge_report,
        retrain_result=retrain_result,
        changelog_path=_os.path.join(_os.getcwd(), "CHANGELOG.md"),
        pyproject_path=_os.path.join(_os.getcwd(), "pyproject.toml"),
    )

    if pub_result.success:
        console.print(f"  [green]Published![/] Model: [bold]{pub_result.ollama_model}[/]")
        console.print()
        console.print(Panel(
            f"[bold green]Retrain Complete[/]\n\n"
            f"  Version: {merge_report.version}\n"
            f"  New corrections: {merge_report.new_from_feedback}\n"
            f"  False positives fixed: {merge_report.false_positive_corrections}\n"
            f"  Accuracy: {retrain_result.eval_accuracy:.1%}\n"
            f"  Model: {pub_result.ollama_model}\n"
            f"  Duration: {retrain_result.duration_seconds // 60}m\n",
            border_style="green",
        ))
    else:
        console.print(f"  [red]Publish failed:[/] {pub_result.error}")

    console.print()
