
from __future__ import annotations

import sys
from importlib import metadata as importlib_metadata
from pathlib import Path

import click
from dotenv import load_dotenv
from rich.console import Console
from rich.panel import Panel
from rich.table import Table

from aegisrt import __version__

load_dotenv()

console = Console()

# ---------------------------------------------------------------------------
# Preset-backed config generation
# ---------------------------------------------------------------------------

_PRESETS = {
    "anthropic": {
        "target_type": "openai_compat",
        "url": "https://api.anthropic.com/v1/messages",
        "headers": '    x-api-key: "${ANTHROPIC_API_KEY}"\n    anthropic-version: "2023-06-01"',
        "env_var": "ANTHROPIC_API_KEY",
        "default_model": "claude-sonnet-4-20250514",
        "judge_type": "anthropic",
        "judge_base_url": "https://api.anthropic.com/v1",
        "judge_api_key": "${ANTHROPIC_API_KEY}",
        "judge_model": "claude-sonnet-4-20250514",
    },
    "openai": {
        "target_type": "openai_compat",
        "url": "https://api.openai.com/v1/chat/completions",
        "headers": '    Authorization: "Bearer ${OPENAI_API_KEY}"',
        "env_var": "OPENAI_API_KEY",
        "default_model": "gpt-4o",
        "judge_type": "openai",
        "judge_base_url": "https://api.openai.com/v1",
        "judge_api_key": "${OPENAI_API_KEY}",
        "judge_model": "gpt-4o-mini",
    },
    "ollama": {
        "target_type": "openai_compat",
        "url": "http://localhost:11434/v1/chat/completions",
        "headers": "",
        "env_var": "",
        "default_model": "llama3.1",
        "judge_type": "",
        "judge_base_url": "",
        "judge_api_key": "",
        "judge_model": "",
    },
}

_PROFILES = {
    "quick": [
        ("prompt_injection", "injection", "high", "[owasp-llm-01]"),
        ("data_leakage", "data_leakage", "critical", "[owasp-llm-06]"),
        ("refusal_bypass", "refusal_bypass", "high", "[owasp-llm-08]"),
        ("tool_misuse", "tool_misuse", "critical", "[owasp-llm-07]"),
    ],
    "standard": [
        ("prompt_injection", "injection", "high", "[owasp-llm-01]"),
        ("data_leakage", "data_leakage", "critical", "[owasp-llm-06]"),
        ("data_exfiltration", "exfiltration", "high", "[owasp-llm-06]"),
        ("refusal_bypass", "refusal_bypass", "high", "[owasp-llm-08]"),
        ("tool_misuse", "tool_misuse", "critical", "[owasp-llm-07]"),
        ("rag_manipulation", "rag_manipulation", "high", "[owasp-llm-05]"),
        ("system_integrity", "system_integrity", "critical", "[owasp-llm-01]"),
        ("output_policy", "output_policy", "medium", "[owasp-llm-02]"),
        ("hallucination", "factuality", "medium", "[owasp-llm-09]"),
        ("bias_stereotyping", "bias", "high", "[]"),
    ],
    "full": [
        ("prompt_injection", "injection", "high", "[owasp-llm-01]"),
        ("data_leakage", "data_leakage", "critical", "[owasp-llm-06]"),
        ("data_exfiltration", "exfiltration", "high", "[owasp-llm-06]"),
        ("refusal_bypass", "refusal_bypass", "high", "[owasp-llm-08]"),
        ("tool_misuse", "tool_misuse", "critical", "[owasp-llm-07]"),
        ("rag_manipulation", "rag_manipulation", "high", "[owasp-llm-05]"),
        ("system_integrity", "system_integrity", "critical", "[owasp-llm-01]"),
        ("output_policy", "output_policy", "medium", "[owasp-llm-02]"),
        ("hallucination", "factuality", "medium", "[owasp-llm-09]"),
        ("bias_stereotyping", "bias", "high", "[]"),
        ("cbrn", "cbrn", "critical", "[]"),
        ("cyber", "cyber", "critical", "[]"),
        ("persuasion", "persuasion", "high", "[]"),
    ],
}


def _build_config(preset_name: str, model: str, profile: str, judge: bool) -> str:
    """Build an aegisrt.yaml string from preset + model + profile + judge flag."""
    p = _PRESETS[preset_name]
    probes = _PROFILES[profile]
    detectors = "[llm_judge]" if judge else "[regex, policy]"

    lines = [
        f"# Generated by: aegisrt init --preset {preset_name} --model {model} --profile {profile}"
        + (" --judge" if judge else ""),
        "",
        "target:",
        f"  type: {p['target_type']}",
        f'  url: "{p["url"]}"',
        "  timeout_seconds: 30",
        "  retries: 2",
    ]

    if p["headers"]:
        lines.append("  headers:")
        lines.append(p["headers"])

    lines.extend([
        "  params:",
        f'    model: "{model}"',
        '    max_tokens: "1024"',
        "",
        "probes:",
    ])

    for probe_id, family, severity, tags in probes:
        lines.extend([
            f"  - id: {probe_id}",
            f"    family: {family}",
            "    generator: static",
            f"    detectors: {detectors}",
            f"    severity: {severity}",
            f"    tags: {tags}",
            "",
        ])

    if judge and p["judge_type"]:
        lines.extend([
            "providers:",
            "  judge:",
            f"    type: {p['judge_type']}",
            f'    model: "{p["judge_model"]}"',
            f'    base_url: "{p["judge_base_url"]}"',
            f'    api_key: "{p["judge_api_key"]}"',
            "",
        ])

    lines.extend([
        "runtime:",
        "  concurrency: 3",
        "  retries: 1",
        "  timeout_seconds: 60",
        "",
        "report:",
        "  formats: [terminal, json]",
        "  output_dir: .aegisrt",
        "  fail_on:",
        "    severity: high",
        "    min_confidence: 0.7",
        "",
    ])

    return "\n".join(lines)


def _tty_prompt() -> tuple[str, str, str, bool]:
    """Ask 2-4 short questions when running in a terminal."""
    console.print()
    console.print("[bold]What are you testing?[/bold]")
    console.print("  1. Claude (Anthropic)")
    console.print("  2. GPT (OpenAI)")
    console.print("  3. Local model (Ollama, vLLM, LM Studio)")
    console.print()

    preset_choice = click.prompt("Choose", type=click.IntRange(1, 3), default=1)
    preset_map = {1: "anthropic", 2: "openai", 3: "ollama"}
    preset = preset_map[preset_choice]
    default_model = _PRESETS[preset]["default_model"]

    model = click.prompt("Model", default=default_model)

    console.print()
    console.print("[bold]How thorough?[/bold]")
    console.print("  1. Quick  — 4 probes, fast feedback")
    console.print("  2. Standard — 10 probes, covers OWASP Top 10")
    console.print("  3. Full — 13 probes, includes CBRN/cyber/persuasion")
    console.print()

    profile_choice = click.prompt("Choose", type=click.IntRange(1, 3), default=1)
    profile_map = {1: "quick", 2: "standard", 3: "full"}
    profile = profile_map[profile_choice]

    judge = False
    if preset != "ollama":
        judge = click.confirm(
            "Use LLM-as-judge for accurate grading? (uses extra API calls)", default=True
        )

    return preset, model, profile, judge

def _package_version(distribution: str) -> str:
    try:
        return importlib_metadata.version(distribution)
    except importlib_metadata.PackageNotFoundError:
        return "unknown"

@click.group()
@click.version_option(version=__version__, prog_name="aegisrt")
def main() -> None:
    pass

try:
    from aegisrt.web.cli_commands import serve
    main.add_command(serve)
except ImportError:

    @main.command()
    def serve() -> None:
        """Launch the AegisRT web dashboard."""
        import click as _click

        _click.echo(
            "Error: The AegisRT dashboard requires web extras.\n"
            "Install with: pip install 'aegisrt[web]'"
        )
        raise SystemExit(1)

@main.command()
@click.option("-o", "--output", default="aegisrt.yaml", help="Output path.")
@click.option(
    "--preset",
    type=click.Choice(list(_PRESETS.keys()), case_sensitive=False),
    default=None,
    help="Provider preset (anthropic, openai, ollama).",
)
@click.option("--model", default=None, help="Model name (e.g. claude-sonnet-4-20250514, gpt-4o, llama3.1).")
@click.option(
    "--profile",
    type=click.Choice(list(_PROFILES.keys()), case_sensitive=False),
    default=None,
    help="Probe profile: quick (4), standard (10), full (13).",
)
@click.option("--judge", is_flag=True, default=False, help="Enable LLM-as-judge evaluation.")
def init(output: str, preset: str | None, model: str | None, profile: str | None, judge: bool) -> None:
    """Generate an aegisrt.yaml config file."""
    out_path = Path(output)
    if out_path.exists():
        if not click.confirm(f"  {out_path} already exists. Overwrite?"):
            console.print("[yellow]Aborted.[/yellow]")
            return

    # If no flags given and we're in a TTY, run the interactive prompt
    if preset is None and sys.stdin.isatty():
        preset, model, profile, judge = _tty_prompt()
    elif preset is None:
        # Non-TTY with no flags: default to anthropic/quick
        preset = "anthropic"

    # Fill defaults for anything not specified
    if model is None:
        model = _PRESETS[preset]["default_model"]
    if profile is None:
        profile = "quick"

    config_text = _build_config(preset, model, profile, judge)
    out_path.write_text(config_text, encoding="utf-8")

    env_var = _PRESETS[preset]["env_var"]
    env_hint = f"\n\nSet your API key first:\n  [bold]export {env_var}=...[/bold] or create a .env file" if env_var else ""

    console.print(
        Panel(
            f"[green]Created:[/green] {out_path}\n"
            f"  Preset: {preset} | Model: {model} | Profile: {profile} | Judge: {'yes' if judge else 'no'}"
            f"{env_hint}\n\n"
            f"Run your first scan:\n"
            f"  [bold]aegisrt run -c {out_path}[/bold]",
            title="AegisRT Init",
            border_style="green",
        )
    )

@main.command(name="run")
@click.option(
    "-c",
    "--config",
    "config_path",
    default="aegisrt.yaml",
    help="Path to the AegisRT config file.",
)
@click.option(
    "--dry-run",
    is_flag=True,
    default=False,
    help="Show what would be run without executing against the target.",
)
@click.option(
    "--watch",
    is_flag=True,
    default=False,
    help="Watch config file for changes and re-run automatically.",
)
@click.option(
    "--no-cache",
    is_flag=True,
    default=False,
    help="Disable response caching for this run.",
)
def run_cmd(config_path: str, dry_run: bool, watch: bool, no_cache: bool) -> None:
    from aegisrt.config.loader import load_config
    from aegisrt.core.runner import SecurityRunner
    from aegisrt.reports.terminal import TerminalReporter
    from aegisrt.reports.json_report import JsonReportWriter
    from aegisrt.reports.html_report import HtmlReportWriter
    from aegisrt.reports.sarif_report import SarifReportWriter
    from aegisrt.reports.junit_report import JunitReportWriter
    from aegisrt.storage.sqlite import ResultStore

    console.print(f"[bold]AegisRT v{__version__}[/bold] -- starting run\n")

    try:
        cfg = load_config(config_path)
    except FileNotFoundError:
        console.print(f"[red]Config file not found:[/red] {config_path}")
        console.print("Run [bold]aegisrt init[/bold] to create one.")
        sys.exit(1)
    except Exception as exc:
        console.print(f"[red]Config error:[/red] {exc}")
        sys.exit(1)

    console.print(f"  Target type : [cyan]{cfg.target.type}[/cyan]")
    console.print(f"  Target URL  : [cyan]{cfg.target.url or 'N/A'}[/cyan]")
    console.print(f"  Probes      : [cyan]{len(cfg.probes)}[/cyan]")
    console.print()

    if dry_run:
        _dry_run_summary(cfg)
        return

    runner = SecurityRunner(cfg, no_cache=no_cache)
    try:
        report = runner.run()
    except NotImplementedError as exc:
        console.print(f"[red]Error:[/red] {exc}")
        sys.exit(1)
    except Exception as exc:
        console.print(f"[red]Run failed:[/red] {exc}")
        sys.exit(1)

    output_dir = ".aegisrt"
    formats: list[str] = ["terminal", "json"]
    if cfg.report:
        output_dir = cfg.report.output_dir
        formats = cfg.report.formats

    run_dir = Path(output_dir) / "runs" / report.run_id
    run_dir.mkdir(parents=True, exist_ok=True)

    duration = report.summary.get("duration_seconds", 0.0)

    for fmt in formats:
        if fmt == "terminal":
            TerminalReporter(console).report(report, duration_seconds=duration)
        elif fmt == "json":
            out = JsonReportWriter().write(report, run_dir / "report.json")
            console.print(f"  [dim]JSON report written to {out}[/dim]")
        elif fmt == "html":
            out = HtmlReportWriter().write(report, run_dir / "report.html")
            console.print(f"  [dim]HTML report written to {out}[/dim]")
        elif fmt == "sarif":
            out = SarifReportWriter().write(report, run_dir / "report.sarif.json")
            console.print(f"  [dim]SARIF report written to {out}[/dim]")
        elif fmt == "junit":
            out = JunitReportWriter().write(report, run_dir / "report.junit.xml")
            console.print(f"  [dim]JUnit report written to {out}[/dim]")
        elif fmt == "csv":
            from aegisrt.reports.csv_report import CsvReportWriter
            out = CsvReportWriter().write(report, run_dir / "report.csv")
            console.print(f"  [dim]CSV report written to {out}[/dim]")

    store = ResultStore(db_path=Path(output_dir) / "results.db")
    try:
        store.save_run(report)
        console.print(f"\n  [dim]Results saved to SQLite ({output_dir}/results.db)[/dim]")
    finally:
        store.close()

    exit_code = _apply_fail_policy(cfg, report)

    _print_summary(report)

    if watch:
        from aegisrt.core.watcher import ConfigWatcher

        def _rerun() -> None:
            try:
                from aegisrt.config.loader import load_config as _load
                from aegisrt.core.runner import SecurityRunner as _Runner

                new_cfg = _load(config_path)
                new_runner = _Runner(new_cfg)
                new_report = new_runner.run()

                new_output_dir = ".aegisrt"
                new_formats: list[str] = ["terminal", "json"]
                if new_cfg.report:
                    new_output_dir = new_cfg.report.output_dir
                    new_formats = new_cfg.report.formats

                new_run_dir = Path(new_output_dir) / "runs" / new_report.run_id
                new_run_dir.mkdir(parents=True, exist_ok=True)
                new_duration = new_report.summary.get("duration_seconds", 0.0)

                for f in new_formats:
                    if f == "terminal":
                        TerminalReporter(console).report(
                            new_report, duration_seconds=new_duration
                        )

                _print_summary(new_report)
            except Exception as exc:
                console.print(f"[red]Re-run failed:[/red] {exc}")

        watcher = ConfigWatcher(config_path, _rerun)
        watcher.start()
        return

    sys.exit(exit_code)

def _dry_run_summary(cfg) -> None:
    from aegisrt.plugins.entrypoints import register_builtin_probes
    from aegisrt.core.runner import SecurityRunner

    registry = register_builtin_probes()
    runner = SecurityRunner(cfg)

    enabled_probes = [p for p in cfg.probes if p.enabled]
    probe_details: list[tuple[str, str, str, int]] = []
    total_cases = 0

    for probe_cfg in enabled_probes:
        probe_cls = runner._resolve_probe_class(registry, probe_cfg)

        if probe_cls is None:
            probe_details.append((
                probe_cfg.id,
                probe_cfg.family or "-",
                probe_cfg.severity or "unknown",
                0,
            ))
            continue

        instance = runner._configure_probe_instance(probe_cls(), probe_cfg)
        case_count = 0
        try:
            generator_name = probe_cfg.generator or ""
            if generator_name == "dataset":
                generator = runner._resolve_generator(instance, probe_cfg)
                case_count = len(generator._load_records())
            elif generator_name in {"llm", "adaptive", "multilingual"}:
                case_count = len(runner._resolve_probe_seeds(instance, probe_cfg))
            else:
                generator = runner._resolve_generator(instance, probe_cfg)
                cases = generator.generate(
                    runner._resolve_probe_seeds(instance, probe_cfg),
                    instance.id,
                )
                case_count = len(cases)
        except Exception:
            case_count = 0

        total_cases += case_count
        probe_details.append((
            instance.id,
            instance.family,
            probe_cfg.severity or instance.severity,
            case_count,
        ))

    concurrency = 4
    if cfg.runtime:
        concurrency = cfg.runtime.concurrency

    table = Table(
        title="Probes to Execute",
        show_header=True,
        header_style="bold cyan",
    )
    table.add_column("Probe ID", style="bold")
    table.add_column("Family")
    table.add_column("Severity")
    table.add_column("Test Cases", justify="right")

    for probe_id, family, severity, case_count in probe_details:
        sev_color = {
            "critical": "bold red",
            "high": "red",
            "medium": "yellow",
            "low": "green",
        }.get(severity, "white")
        table.add_row(
            probe_id,
            family,
            f"[{sev_color}]{severity.upper()}[/{sev_color}]",
            str(case_count),
        )

    console.print(table)

    lines = [
        f"Target type   : {cfg.target.type}",
        f"Target URL    : {cfg.target.url or 'N/A'}",
        f"Probes        : {len(probe_details)}",
        f"Total cases   : {total_cases}",
        f"Concurrency   : {concurrency}",
    ]
    if cfg.report:
        lines.append(f"Report formats: {', '.join(cfg.report.formats)}")
        lines.append(f"Output dir    : {cfg.report.output_dir}")

    console.print()
    console.print(
        Panel(
            "\n".join(lines),
            title="Dry Run Summary",
            subtitle="No tests were executed",
            border_style="yellow",
        )
    )

def _apply_fail_policy(cfg, report) -> int:
    if cfg.report and cfg.report.fail_on:
        policy = cfg.report.fail_on
        severity_order = {"low": 0, "medium": 1, "high": 2, "critical": 3}
        threshold_level = severity_order.get(policy.severity, 2)
        min_confidence = policy.min_confidence
        configured_severity_by_probe = {
            probe.id: severity_order.get(probe.severity, 0)
            for probe in cfg.probes
        }

        for r in report.results:
            if not r.passed:
                result_level = severity_order.get(r.severity, 0)
                configured_level = configured_severity_by_probe.get(r.probe_id, 0)
                effective_level = max(result_level, configured_level)
                if effective_level >= threshold_level and r.confidence >= min_confidence:
                    return 1
    else:
        if any(not r.passed for r in report.results):
            return 1

    return 0

def _print_summary(report) -> None:
    summary = report.summary
    total = summary.get("total", 0)
    passed = summary.get("passed", 0)
    failed = summary.get("failed", 0)
    duration = summary.get("duration_seconds", 0.0)
    pass_rate = summary.get("pass_rate", 1.0)

    by_severity: dict[str, int] = {}
    for r in report.results:
        if not r.passed:
            sev = r.severity.lower()
            by_severity[sev] = by_severity.get(sev, 0) + 1

    sev_parts = []
    for sev in ("critical", "high", "medium", "low"):
        count = by_severity.get(sev, 0)
        if count > 0:
            sev_parts.append(f"{sev.upper()}: {count}")

    sev_line = ", ".join(sev_parts) if sev_parts else "none"

    verdict_style = "bold white on green" if failed == 0 else "bold white on red"
    verdict_text = "PASS" if failed == 0 else "FAIL"

    lines = [
        f"Run ID       : {report.run_id}",
        f"Total tests  : {total}",
        f"Passed       : {passed}",
        f"Failed       : {failed}",
        f"Pass rate    : {pass_rate * 100:.1f}%",
        f"Duration     : {duration:.2f}s",
        f"Failed by sev: {sev_line}",
        "",
        f"Verdict      : [{verdict_style}] {verdict_text} [/{verdict_style}]",
    ]

    console.print()
    console.print(
        Panel(
            "\n".join(lines),
            title="Run Summary",
            border_style="blue",
        )
    )

@main.group(name="report")
def report_group() -> None:
    pass

@report_group.command(name="latest")
@click.option("--format", "fmt", default="terminal", help="Output format (terminal, json, html, sarif, junit).")
def report_latest(fmt: str) -> None:
    from aegisrt.storage.sqlite import ResultStore

    store = ResultStore()
    try:
        runs = store.list_runs()
        if not runs:
            console.print("[yellow]No runs found.[/yellow] Run [bold]aegisrt run[/bold] first.")
            return

        latest_run_id = runs[0]["run_id"]
        report = store.load_run(latest_run_id)
        if report is None:
            console.print(f"[red]Could not load run:[/red] {latest_run_id}")
            return

        _output_report(report, fmt)
    finally:
        store.close()

@report_group.command(name="show")
@click.argument("run_id")
@click.option("--format", "fmt", default="terminal", help="Output format (terminal, json, html, sarif, junit).")
def report_show(run_id: str, fmt: str) -> None:
    from aegisrt.storage.sqlite import ResultStore

    store = ResultStore()
    try:
        report = store.load_run(run_id)
        if report is None:
            console.print(f"[red]Run not found:[/red] {run_id}")
            console.print("[dim]Use 'aegisrt report latest' to see the most recent run.[/dim]")
            return

        _output_report(report, fmt)
    finally:
        store.close()

def _output_report(report, fmt: str) -> None:
    from aegisrt.reports.terminal import TerminalReporter
    from aegisrt.reports.json_report import JsonReportWriter
    from aegisrt.reports.html_report import HtmlReportWriter
    from aegisrt.reports.sarif_report import SarifReportWriter
    from aegisrt.reports.junit_report import JunitReportWriter

    duration = report.summary.get("duration_seconds", 0.0)

    if fmt == "terminal":
        TerminalReporter(console).report(report, duration_seconds=duration)
    elif fmt == "json":
        out = JsonReportWriter().write(report, f".aegisrt/runs/{report.run_id}/report.json")
        console.print(f"[green]JSON report written to {out}[/green]")
    elif fmt == "html":
        out = HtmlReportWriter().write(report, f".aegisrt/runs/{report.run_id}/report.html")
        console.print(f"[green]HTML report written to {out}[/green]")
    elif fmt == "sarif":
        out = SarifReportWriter().write(report, f".aegisrt/runs/{report.run_id}/report.sarif.json")
        console.print(f"[green]SARIF report written to {out}[/green]")
    elif fmt == "junit":
        out = JunitReportWriter().write(report, f".aegisrt/runs/{report.run_id}/report.junit.xml")
        console.print(f"[green]JUnit report written to {out}[/green]")
    elif fmt == "csv":
        from aegisrt.reports.csv_report import CsvReportWriter
        out = CsvReportWriter().write(report, f".aegisrt/runs/{report.run_id}/report.csv")
        console.print(f"[green]CSV report written to {out}[/green]")
    else:
        console.print(f"[red]Unknown format:[/red] {fmt}")
        console.print("[dim]Supported: terminal, json, html, sarif, junit, csv[/dim]")

@main.group(name="list")
def list_group() -> None:
    pass

@list_group.command(name="probes")
def list_probes() -> None:
    from aegisrt.plugins.entrypoints import register_builtin_probes

    registry = register_builtin_probes()

    if not registry:
        console.print("[yellow]No probes found.[/yellow]")
        return

    table = Table(title="Available Probes", show_header=True, header_style="bold cyan")
    table.add_column("ID", style="bold")
    table.add_column("Family")
    table.add_column("Severity")
    table.add_column("Description", max_width=50)
    table.add_column("Seeds", justify="right")

    for probe_id, probe_cls in sorted(registry.items()):
        instance = probe_cls()
        seeds = instance.get_seeds()
        sev = getattr(instance, "severity", "medium")
        sev_color = {
            "critical": "bold red",
            "high": "red",
            "medium": "yellow",
            "low": "green",
        }.get(sev, "white")

        table.add_row(
            probe_id,
            getattr(instance, "family", "unknown"),
            f"[{sev_color}]{sev.upper()}[/{sev_color}]",
            getattr(instance, "description", ""),
            str(len(seeds)),
        )

    console.print(table)

@list_group.command(name="suites")
def list_suites() -> None:
    from aegisrt.suites.registry import SuiteRegistry

    registry = SuiteRegistry()
    registry.load_builtins()

    suite_names = registry.list_suites()
    if not suite_names:
        console.print("[yellow]No suites found.[/yellow]")
        return

    table = Table(title="Available Suites", show_header=True, header_style="bold cyan")
    table.add_column("Name", style="bold")
    table.add_column("Description")
    table.add_column("Probes", justify="right")

    for name in suite_names:
        suite = registry.get(name)
        if suite is None:
            continue
        table.add_row(
            suite.name,
            suite.description,
            str(len(suite.probes)),
        )

    console.print(table)

@main.command()
@click.argument("path", default=".")
def discover(path: str) -> None:
    from aegisrt.audit.discover import DiscoveryScanner

    console.print(f"[bold]Scanning[/bold] {path} for LLM integrations...\n")

    scanner = DiscoveryScanner()
    result = scanner.scan(path)

    if not result.findings:
        console.print("[green]No LLM integrations found.[/green]")
        return

    table = Table(title="LLM Integration Discovery", show_header=True, header_style="bold cyan")
    table.add_column("File", style="dim")
    table.add_column("Line", justify="right")
    table.add_column("Module", style="bold")
    table.add_column("Integration Type")
    table.add_column("Recommended Probes", max_width=40)

    for finding in result.findings:
        table.add_row(
            finding.file_path,
            str(finding.line),
            finding.module_name,
            finding.integration_type,
            ", ".join(finding.recommended_probes),
        )

    console.print(table)

    summary = result.summary
    console.print(
        f"\n  Scanned files : {summary.get('scanned_files', 0)}"
        f"\n  Findings      : {summary.get('total_findings', 0)}"
        f"\n  Modules found : {', '.join(summary.get('modules_found', []))}"
    )

@main.command()
@click.argument("path", default=".")
def audit(path: str) -> None:
    from aegisrt.audit.rules import ALL_RULES
    from aegisrt.audit.python_ast import parse_file

    console.print(f"[bold]Auditing[/bold] {path} for LLM security anti-patterns...\n")

    root = Path(path)
    all_findings = []
    scanned = 0

    if root.is_file():
        py_files = [root]
    else:
        py_files = sorted(root.rglob("*.py"))

    for py_file in py_files:
        parts = py_file.parts
        if any(
            p.startswith(".") or p in {"__pycache__", "node_modules", ".venv", "venv"}
            for p in parts
        ):
            continue

        tree = parse_file(py_file)
        if tree is None:
            continue
        scanned += 1

        for rule in ALL_RULES:
            findings = rule.match(tree, str(py_file))
            all_findings.extend(findings)

    if not all_findings:
        console.print(f"[green]No findings in {scanned} files scanned.[/green]")
        return

    table = Table(title="Audit Findings", show_header=True, header_style="bold cyan")
    table.add_column("Rule", style="bold")
    table.add_column("File", style="dim")
    table.add_column("Line", justify="right")
    table.add_column("Severity")
    table.add_column("Message", max_width=60)

    has_high_or_critical = False
    for finding in all_findings:
        sev = finding.severity.lower()
        sev_color = {
            "critical": "bold red",
            "high": "red",
            "medium": "yellow",
            "low": "green",
        }.get(sev, "white")

        if sev in ("high", "critical"):
            has_high_or_critical = True

        table.add_row(
            finding.rule_id,
            finding.file_path,
            str(finding.line),
            f"[{sev_color}]{sev.upper()}[/{sev_color}]",
            finding.message,
        )

    console.print(table)
    console.print(f"\n  Scanned files : {scanned}")
    console.print(f"  Total findings: {len(all_findings)}")

    if has_high_or_critical:
        console.print("\n[bold red]High/critical findings detected.[/bold red]")
        sys.exit(1)

@main.command()
def doctor() -> None:
    console.print(f"[bold]AegisRT Doctor v{__version__}[/bold]\n")

    py_ver = sys.version.split()[0]
    py_ok = sys.version_info >= (3, 10)
    py_status = "[green]OK[/green]" if py_ok else "[red]FAIL[/red]"
    console.print(f"  Python version : {py_ver}  {py_status}")

    if not py_ok:
        console.print("  [red]Python 3.10+ is required.[/red]")

    console.print("\n  [bold]Required dependencies:[/bold]")
    deps = [
        ("pydantic", "pydantic", "pydantic"),
        ("PyYAML", "yaml", "pyyaml"),
        ("click", "click", "click"),
        ("rich", "rich", "rich"),
        ("httpx", "httpx", "httpx"),
        ("jinja2", "jinja2", "jinja2"),
    ]
    all_ok = True
    for distribution, import_name, label in deps:
        try:
            __import__(import_name)
            ver = _package_version(distribution)
            console.print(f"    {label:14s} : {ver}  [green]OK[/green]")
        except ImportError:
            console.print(f"    {label:14s} : [red]NOT INSTALLED[/red]")
            all_ok = False

    config_path = Path("aegisrt.yaml")
    if config_path.exists():
        try:
            from aegisrt.config.loader import load_config

            load_config(config_path)
            console.print(f"\n  Config ({config_path}) : [green]VALID[/green]")
        except Exception as exc:
            console.print(f"\n  Config ({config_path}) : [red]INVALID[/red] - {exc}")
    else:
        console.print(f"\n  Config ({config_path}) : [yellow]NOT FOUND[/yellow]")
        console.print("  Run [bold]aegisrt init[/bold] to create one.")

    console.print("\n  [bold]Built-in probes:[/bold]")
    try:
        from aegisrt.plugins.entrypoints import register_builtin_probes

        registry = register_builtin_probes()
        if registry:
            for probe_id in sorted(registry.keys()):
                console.print(f"    {probe_id:30s} [green]OK[/green]")
        else:
            console.print("    [yellow]No probes loaded[/yellow]")
    except Exception as exc:
        console.print(f"    [red]Failed to load probes:[/red] {exc}")

    console.print("\n  [bold]Optional dependencies:[/bold]")
    for dep in ["openai", "anthropic"]:
        try:
            __import__(dep)
            ver = _package_version(dep)
            console.print(f"    {dep:12s} : {ver}  [green]OK[/green]")
        except ImportError:
            console.print(f"    {dep:12s} : [dim]not installed[/dim]")

    if py_ok and all_ok:
        console.print("\n  [bold green]All checks passed.[/bold green]")
    else:
        console.print("\n  [bold yellow]Some checks failed. See above.[/bold yellow]")

@main.command()
@click.argument("run_id")
def replay(run_id: str) -> None:
    from aegisrt.storage.sqlite import ResultStore
    from aegisrt.reports.terminal import TerminalReporter

    store = ResultStore()
    try:
        report = store.load_run(run_id)
        if report is None:
            console.print(f"[red]Run not found:[/red] {run_id}")
            console.print("[dim]Use 'aegisrt report latest' to see available runs.[/dim]")
            return

        duration = report.summary.get("duration_seconds", 0.0)
        TerminalReporter(console).report(report, duration_seconds=duration)
    finally:
        store.close()

@main.command()
@click.argument("run_id_1")
@click.argument("run_id_2")
@click.option("--json-output", "json_out", default=None, help="Write JSON diff to file.")
def diff(run_id_1: str, run_id_2: str, json_out: str | None) -> None:
    from aegisrt.storage.sqlite import ResultStore
    from aegisrt.core.diff import compare_runs
    from aegisrt.reports.diff_report import DiffReporter

    store = ResultStore()
    try:
        base = store.load_run(run_id_1)
        if base is None:
            console.print(f"[red]Base run not found:[/red] {run_id_1}")
            sys.exit(1)

        compare = store.load_run(run_id_2)
        if compare is None:
            console.print(f"[red]Compare run not found:[/red] {run_id_2}")
            sys.exit(1)

        run_diff = compare_runs(base, compare)

        reporter = DiffReporter(console)
        reporter.report(run_diff)

        if json_out:
            out_path = reporter.write_json(run_diff, json_out)
            console.print(f"\n  [dim]JSON diff written to {out_path}[/dim]")

        if run_diff.regressions > 0:
            sys.exit(1)
    finally:
        store.close()

@main.command()
def history() -> None:
    from aegisrt.storage.sqlite import ResultStore

    store = ResultStore()
    try:
        runs = store.list_runs()
        if not runs:
            console.print("[yellow]No runs found.[/yellow] Run [bold]aegisrt run[/bold] first.")
            return

        table = Table(
            title="Run History",
            show_header=True,
            header_style="bold cyan",
        )
        table.add_column("Run ID", style="bold")
        table.add_column("Timestamp")
        table.add_column("Total", justify="right")
        table.add_column("Passed", justify="right", style="green")
        table.add_column("Failed", justify="right", style="red")
        table.add_column("Pass Rate", justify="right")
        table.add_column("Trend")

        runs_chrono = list(reversed(runs))
        prev_pass_rate: float | None = None

        trend_data: list[tuple[dict, str]] = []
        for run in runs_chrono:
            summary = run.get("summary", {})
            total = summary.get("total", 0)
            passed = summary.get("passed", 0)
            failed = summary.get("failed", 0)
            pass_rate = summary.get("pass_rate", 1.0 if total == 0 else 0.0)

            if prev_pass_rate is not None:
                if pass_rate > prev_pass_rate:
                    trend = "[green]^ UP[/green]"
                elif pass_rate < prev_pass_rate:
                    trend = "[red]v DOWN[/red]"
                else:
                    trend = "[dim]= SAME[/dim]"
            else:
                trend = "[dim]--[/dim]"

            prev_pass_rate = pass_rate
            trend_data.append((run, trend))

        for run, trend in reversed(trend_data):
            summary = run.get("summary", {})
            total = summary.get("total", 0)
            passed = summary.get("passed", 0)
            failed = summary.get("failed", 0)
            pass_rate = summary.get("pass_rate", 1.0 if total == 0 else 0.0)

            table.add_row(
                run["run_id"],
                run["timestamp"],
                str(total),
                str(passed),
                str(failed),
                f"{pass_rate * 100:.1f}%",
                trend,
            )

        console.print(table)
    finally:
        store.close()

@main.group(name="benchmark")
def benchmark_group() -> None:
    pass

@benchmark_group.command(name="run")
@click.option(
    "-c",
    "--config",
    "config_path",
    default="aegisrt-benchmark.yaml",
    help="Path to the benchmark config file.",
)
def benchmark_run(config_path: str) -> None:
    from aegisrt.config.loader import load_benchmark_config
    from aegisrt.core.benchmark import BenchmarkRunner
    from aegisrt.reports.benchmark_report import (
        BenchmarkCsvReportWriter,
        BenchmarkHtmlReportWriter,
        BenchmarkTerminalReporter,
    )

    console.print(f"[bold]AegisRT Benchmark v{__version__}[/bold]\n")

    try:
        cfg = load_benchmark_config(config_path)
    except FileNotFoundError:
        console.print(f"[red]Config file not found:[/red] {config_path}")
        sys.exit(1)
    except Exception as exc:
        console.print(f"[red]Config error:[/red] {exc}")
        sys.exit(1)

    console.print(f"  Targets    : [cyan]{len(cfg.targets)}[/cyan]")
    for t in cfg.targets:
        console.print(f"    - {t.name} ({t.type}: {t.url or 'N/A'})")
    console.print(f"  Probes     : [cyan]{len(cfg.probes)}[/cyan]")
    console.print()

    runner = BenchmarkRunner(cfg)
    try:
        report = runner.run()
    except Exception as exc:
        console.print(f"[red]Benchmark failed:[/red] {exc}")
        sys.exit(1)

    BenchmarkTerminalReporter(console).report(report)

    output_dir = ".aegisrt"
    formats: list[str] = ["terminal", "json"]
    if cfg.report:
        output_dir = cfg.report.output_dir
        formats = cfg.report.formats

    run_dir = Path(output_dir) / "benchmarks" / report.benchmark_id
    run_dir.mkdir(parents=True, exist_ok=True)

    json_path = run_dir / "benchmark.json"
    json_path.write_text(report.model_dump_json(indent=2), encoding="utf-8")
    console.print(f"\n  [dim]JSON report: {json_path}[/dim]")

    if "html" in formats:
        html_path = BenchmarkHtmlReportWriter().write(
            report, run_dir / "benchmark.html"
        )
        console.print(f"  [dim]HTML report: {html_path}[/dim]")

    if "csv" in formats:
        csv_path = BenchmarkCsvReportWriter().write(
            report, run_dir / "benchmark.csv"
        )
        console.print(f"  [dim]CSV report: {csv_path}[/dim]")

@benchmark_group.command(name="compare")
@click.argument("benchmark_id_1")
@click.argument("benchmark_id_2")
@click.option(
    "-d",
    "--dir",
    "output_dir",
    default=".aegisrt",
    help="Output directory where benchmarks are stored.",
)
def benchmark_compare(benchmark_id_1: str, benchmark_id_2: str, output_dir: str) -> None:
    import json as json_mod

    from aegisrt.core.benchmark import BenchmarkReport
    from aegisrt.reports.benchmark_report import BenchmarkTerminalReporter

    def _load(bid: str) -> BenchmarkReport | None:
        p = Path(output_dir) / "benchmarks" / bid / "benchmark.json"
        if not p.exists():
            return None
        data = json_mod.loads(p.read_text(encoding="utf-8"))
        return BenchmarkReport.model_validate(data)

    r1 = _load(benchmark_id_1)
    if r1 is None:
        console.print(f"[red]Benchmark not found:[/red] {benchmark_id_1}")
        sys.exit(1)

    r2 = _load(benchmark_id_2)
    if r2 is None:
        console.print(f"[red]Benchmark not found:[/red] {benchmark_id_2}")
        sys.exit(1)

    console.print(f"[bold]Comparing benchmarks[/bold]\n")
    console.print(f"  [cyan]A[/cyan]: {benchmark_id_1} ({r1.timestamp})")
    console.print(f"  [cyan]B[/cyan]: {benchmark_id_2} ({r2.timestamp})")
    console.print()

    console.print("[bold]--- Benchmark A ---[/bold]")
    BenchmarkTerminalReporter(console).report(r1)
    console.print()
    console.print("[bold]--- Benchmark B ---[/bold]")
    BenchmarkTerminalReporter(console).report(r2)

    console.print()
    table = Table(title="Delta (B - A)", show_header=True, header_style="bold cyan")
    table.add_column("Target", style="bold")
    table.add_column("A Rate", justify="right")
    table.add_column("B Rate", justify="right")
    table.add_column("Delta", justify="right")

    a_scores = {s.target_name: s for s in r1.scores}
    b_scores = {s.target_name: s for s in r2.scores}
    all_targets = sorted(set(a_scores.keys()) | set(b_scores.keys()))

    for name in all_targets:
        a_rate = a_scores[name].pass_rate if name in a_scores else 0.0
        b_rate = b_scores[name].pass_rate if name in b_scores else 0.0
        delta = (b_rate - a_rate) * 100
        color = "green" if delta > 0 else ("red" if delta < 0 else "dim")
        sign = "+" if delta > 0 else ""
        table.add_row(
            name,
            f"{a_rate * 100:.1f}%",
            f"{b_rate * 100:.1f}%",
            f"[{color}]{sign}{delta:.1f}pp[/{color}]",
        )

    console.print(table)

@benchmark_group.command(name="leaderboard")
@click.option(
    "-d",
    "--dir",
    "output_dir",
    default=".aegisrt",
    help="Output directory where benchmarks are stored.",
)
def benchmark_leaderboard(output_dir: str) -> None:
    import json as json_mod

    from aegisrt.core.benchmark import BenchmarkReport

    bench_dir = Path(output_dir) / "benchmarks"
    if not bench_dir.exists():
        console.print("[yellow]No benchmarks found.[/yellow]")
        return

    all_scores: dict[str, list[float]] = {}

    for bdir in sorted(bench_dir.iterdir()):
        jpath = bdir / "benchmark.json"
        if not jpath.exists():
            continue
        try:
            data = json_mod.loads(jpath.read_text(encoding="utf-8"))
            report = BenchmarkReport.model_validate(data)
        except Exception:
            continue
        for score in report.scores:
            if score.target_name not in all_scores:
                all_scores[score.target_name] = []
            all_scores[score.target_name].append(score.pass_rate)

    if not all_scores:
        console.print("[yellow]No benchmark results found.[/yellow]")
        return

    table = Table(title="All-Time Leaderboard", show_header=True, header_style="bold cyan")
    table.add_column("Rank", justify="center")
    table.add_column("Target", style="bold")
    table.add_column("Best Rate", justify="right")
    table.add_column("Avg Rate", justify="right")
    table.add_column("Runs", justify="right")

    ranked = sorted(
        all_scores.items(),
        key=lambda kv: max(kv[1]),
        reverse=True,
    )

    for i, (name, rates) in enumerate(ranked, 1):
        best = max(rates) * 100
        avg = (sum(rates) / len(rates)) * 100
        table.add_row(
            str(i),
            name,
            f"{best:.1f}%",
            f"{avg:.1f}%",
            str(len(rates)),
        )

    console.print(table)

@main.group(name="datasets")
def datasets_group() -> None:
    pass

@datasets_group.command(name="list")
def datasets_list() -> None:
    from aegisrt.datasets.registry import list_datasets

    datasets = list_datasets()
    if not datasets:
        console.print("[yellow]No datasets found.[/yellow]")
        return

    table = Table(
        title="Built-in Attack Datasets",
        show_header=True,
        header_style="bold cyan",
    )
    table.add_column("ID", style="bold")
    table.add_column("Name")
    table.add_column("Size", justify="right")
    table.add_column("Categories")
    table.add_column("Available")

    for ds in datasets:
        cats = ", ".join(ds.get("categories", [])[:4])
        if len(ds.get("categories", [])) > 4:
            cats += ", ..."
        available = "[green]Yes[/green]" if ds["available"] else "[red]No[/red]"
        table.add_row(
            ds["id"],
            ds["name"],
            str(ds.get("size", "?")),
            cats,
            available,
        )

    console.print(table)
    console.print(
        "\n  [dim]Use 'aegisrt datasets info <id>' for details.[/dim]"
    )
    console.print(
        "  [dim]Use 'builtin://<id>' as a dataset path in configs.[/dim]"
    )

@datasets_group.command(name="info")
@click.argument("dataset_id")
def datasets_info(dataset_id: str) -> None:
    from aegisrt.datasets.registry import get_dataset_info

    info = get_dataset_info(dataset_id)
    if info is None:
        console.print(f"[red]Dataset not found:[/red] {dataset_id}")
        console.print("[dim]Use 'aegisrt datasets list' to see available datasets.[/dim]")
        return

    lines = [
        f"ID          : {info['id']}",
        f"Name        : {info['name']}",
        f"Description : {info['description']}",
        f"Source      : {info.get('source', 'N/A')}",
        f"License     : {info.get('license', 'N/A')}",
        f"Size        : {info.get('actual_size', info.get('size', '?'))} prompts",
        f"Categories  : {', '.join(info.get('categories', []))}",
        f"File        : {info.get('path', 'N/A')}",
        f"Available   : {'Yes' if info['available'] else 'No'}",
    ]

    if "languages" in info:
        lines.append(f"Languages   : {', '.join(info['languages'])}")

    console.print(
        Panel(
            "\n".join(lines),
            title=f"Dataset: {info['name']}",
            border_style="cyan",
        )
    )

    console.print(
        f"\n  [dim]Config usage: generator: {{type: dataset, path: builtin://{dataset_id}}}[/dim]"
    )

@datasets_group.command(name="preview")
@click.argument("dataset_id")
@click.option("-n", "--count", default=5, help="Number of prompts to preview.")
def datasets_preview(dataset_id: str, count: int) -> None:
    from aegisrt.datasets.registry import load_dataset

    try:
        records = load_dataset(f"builtin://{dataset_id}")
    except (ValueError, FileNotFoundError) as exc:
        console.print(f"[red]Error:[/red] {exc}")
        return

    table = Table(
        title=f"Preview: {dataset_id} (first {count})",
        show_header=True,
        header_style="bold cyan",
    )
    table.add_column("#", justify="right", style="dim")
    table.add_column("Prompt", max_width=80)
    table.add_column("Category")

    for i, record in enumerate(records[:count]):
        table.add_row(
            str(i + 1),
            record.get("prompt", "")[:80],
            record.get("category", ""),
        )

    console.print(table)
    console.print(f"\n  [dim]Total: {len(records)} prompts[/dim]")

@main.command(name="compare-models")
@click.option(
    "-c",
    "--config",
    "config_paths",
    multiple=True,
    required=True,
    help="Path to a AegisRT config file (specify multiple times).",
)
def compare_models(config_paths: tuple[str, ...]) -> None:
    from aegisrt.config.loader import load_config
    from aegisrt.core.comparison import run_comparison

    if len(config_paths) < 2:
        console.print("[red]At least two config files are required for comparison.[/red]")
        sys.exit(1)

    configs = []
    for cp in config_paths:
        try:
            configs.append(load_config(cp))
        except FileNotFoundError:
            console.print(f"[red]Config file not found:[/red] {cp}")
            sys.exit(1)
        except Exception as exc:
            console.print(f"[red]Config error in {cp}:[/red] {exc}")
            sys.exit(1)

    console.print(f"[bold]Comparing {len(configs)} targets...[/bold]\n")

    comparison = run_comparison(configs)

    table = Table(
        title="Model Comparison",
        show_header=True,
        header_style="bold cyan",
    )
    table.add_column("Target", style="bold")
    table.add_column("Type")
    table.add_column("Total", justify="right")
    table.add_column("Passed", justify="right", style="green")
    table.add_column("Failed", justify="right", style="red")
    table.add_column("Pass Rate", justify="right")

    for target in comparison.targets:
        label = target["name"]
        stats = comparison.summary.get("per_target", {}).get(label, {})
        if "error" in stats:
            table.add_row(
                label,
                target["type"],
                "--",
                "--",
                "--",
                f"[red]ERROR: {stats['error'][:40]}[/red]",
            )
        else:
            table.add_row(
                label,
                target["type"],
                str(stats.get("total", 0)),
                str(stats.get("passed", 0)),
                str(stats.get("failed", 0)),
                f"{stats.get('pass_rate', 0) * 100:.1f}%",
            )

    console.print(table)

    probe_table = Table(
        title="Per-Probe Comparison",
        show_header=True,
        header_style="bold cyan",
    )
    probe_table.add_column("Probe Family", style="bold")
    for target in comparison.targets:
        probe_table.add_column(target["name"][:30], justify="center")

    all_probes: set[str] = set()
    for label, results in comparison.results_by_target.items():
        for r in results:
            all_probes.add(r.probe_id)

    for probe_id in sorted(all_probes):
        row = [probe_id]
        for target in comparison.targets:
            label = target["name"]
            target_results = [
                r for r in comparison.results_by_target.get(label, [])
                if r.probe_id == probe_id
            ]
            if not target_results:
                row.append("[dim]--[/dim]")
            else:
                passed = sum(1 for r in target_results if r.passed)
                total = len(target_results)
                rate = passed / total if total else 1.0
                color = "green" if rate >= 0.8 else ("yellow" if rate >= 0.5 else "red")
                row.append(f"[{color}]{passed}/{total} ({rate * 100:.0f}%)[/{color}]")
        probe_table.add_row(*row)

    console.print(probe_table)

    best_label = None
    best_rate = -1.0
    per_target = comparison.summary.get("per_target", {})
    for label, stats in per_target.items():
        rate = stats.get("pass_rate", 0.0)
        if rate > best_rate:
            best_rate = rate
            best_label = label

    if best_label:
        console.print(
            f"\n  [bold green]Most resistant target:[/bold green] {best_label} "
            f"({best_rate * 100:.1f}% pass rate)"
        )

@main.command()
@click.argument("run_id", default="latest")
@click.option("--html", "html_path", default=None, help="Write OWASP compliance HTML report to this path.")
def compliance(run_id: str, html_path: str | None) -> None:
    from aegisrt.storage.sqlite import ResultStore
    from aegisrt.compliance.owasp import generate_compliance_report

    store = ResultStore()
    try:
        if run_id == "latest":
            runs = store.list_runs()
            if not runs:
                console.print("[yellow]No runs found.[/yellow] Run [bold]aegisrt run[/bold] first.")
                return
            run_id = runs[0]["run_id"]

        report = store.load_run(run_id)
        if report is None:
            console.print(f"[red]Run not found:[/red] {run_id}")
            return

        compliance_data = generate_compliance_report(report.results)
        categories = compliance_data["categories"]
        summary = compliance_data["summary"]
        gaps = compliance_data["gaps"]

        console.print(
            Panel(
                f"Run ID     : {run_id}\n"
                f"Coverage   : {summary['coverage_pct']}% ({summary['tested']}/{summary['total_categories']} categories)\n"
                f"Passed     : {summary['passed']}\n"
                f"Failed     : {summary['failed']}\n"
                f"Not tested : {summary['not_tested']}",
                title="OWASP LLM Top 10 Compliance Summary",
                border_style="blue",
            )
        )

        table = Table(
            title="Category Breakdown",
            show_header=True,
            header_style="bold cyan",
        )
        table.add_column("ID", style="bold")
        table.add_column("Category")
        table.add_column("Severity")
        table.add_column("Status")
        table.add_column("Tests", justify="right")
        table.add_column("Passed", justify="right")
        table.add_column("Failed", justify="right")

        for owasp_id, cat in categories.items():
            sev = cat["severity"]
            sev_color = {
                "critical": "bold red",
                "high": "red",
                "medium": "yellow",
                "low": "green",
            }.get(sev, "white")

            status = cat["status"]
            status_color = {
                "passed": "green",
                "failed": "red",
                "not_tested": "dim",
                "audit_findings_only": "yellow",
            }.get(status, "white")

            table.add_row(
                owasp_id,
                cat["name"],
                f"[{sev_color}]{sev.upper()}[/{sev_color}]",
                f"[{status_color}]{status.upper().replace('_', ' ')}[/{status_color}]",
                str(cat["total"]),
                str(cat["passed"]),
                str(cat["failed"]),
            )

        console.print(table)

        if gaps:
            console.print("\n[bold yellow]Coverage Gaps[/bold yellow]\n")
            for gap in gaps:
                sev = gap["severity"]
                sev_color = {
                    "critical": "bold red",
                    "high": "red",
                    "medium": "yellow",
                }.get(sev, "white")
                console.print(
                    f"  [{sev_color}]{gap['id']}[/{sev_color}] "
                    f"[bold]{gap['name']}[/bold]"
                )
                console.print(f"    {gap['recommendation']}")
                console.print()

        if html_path:
            from aegisrt.reports.compliance_report import ComplianceReportWriter

            out = ComplianceReportWriter().write(report, html_path)
            console.print(f"\n[green]OWASP compliance HTML report written to {out}[/green]")

    finally:
        store.close()

if __name__ == "__main__":
    main()
