"""
compliance_cli.py — Gap #1-4: `sichgate compliance` command group.

Commands:
    sichgate compliance delta          — Gap #1: comparative vulnerability profiling
    sichgate compliance quantization   — Gap #2: quantization robustness variance report
    sichgate compliance audit-trail    — Gap #3: generate audit trail from run
    sichgate compliance mitigate       — Gap #3: add mitigation to a finding
    sichgate compliance export-report  — Gap #3: export PDF/JSON compliance report
    sichgate compliance composition-test — Gap #4: multi-model composition testing

Mount in cli.py with:
    from sichgate.compliance_cli import compliance_app
    app.add_typer(compliance_app, name="compliance")
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import List, Optional

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

compliance_app = typer.Typer(
    name="compliance",
    help="Compliance-grade testing: delta profiling, quantization variance, audit trails, composition.",
    add_completion=False,
)
console = Console()


# ---------------------------------------------------------------------------
# Gap #1: sichgate compliance delta
# ---------------------------------------------------------------------------

@compliance_app.command(name="delta")
def delta(
    base_model: str = typer.Option(..., "--base-model", help="Base model path or HuggingFace ID"),
    ft_model: str = typer.Option(..., "--ft-model", help="Fine-tuned model path or HuggingFace ID"),
    attacks: str = typer.Option("all-llm", "--attacks", help="Comma-separated attacks or group (e.g. all-llm, jailbreak)"),
    ft_data: Optional[str] = typer.Option(None, "--ft-data", help="Path to fine-tuning JSONL data (optional, for attribution context)"),
    adapter: str = typer.Option("hf", "--adapter", help="Model adapter: hf | openai_compat"),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write delta report JSON to file"),
    verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose logging"),
) -> None:
    """
    Gap #1: Comparative Vulnerability Profiling.

    Runs an identical attack battery against the base model and fine-tuned model,
    then computes a delta report classifying each vulnerability as FIXED, REGRESSED, or NEUTRAL.

    Example:
        sichgate compliance delta \\
          --base-model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \\
          --ft-model ./models/my-finetuned \\
          --attacks all-llm \\
          --output delta_report.json
    """
    import logging
    if verbose:
        logging.getLogger().setLevel(logging.INFO)

    from sichgate.runner import build_registry, AttackRunner
    from sichgate.profilers.baseline_profiler import BaselineProfiler
    from sichgate.models.adapter import ModelAdapter

    # Resolve attack classes
    registry = build_registry()
    attack_names = [a.strip() for a in attacks.split(",")]
    attack_classes = []
    for name in attack_names:
        if name in registry:
            for cls in registry[name]:
                if cls not in attack_classes:
                    attack_classes.append(cls)
        else:
            console.print(f"[red]Unknown attack: {name}[/red]")
            raise typer.Exit(1)

    console.print(
        f"[bold]SichGate[/bold] Gap #1: Comparative Vulnerability Profiling\n"
        f"  Base model:   [cyan]{base_model}[/cyan]\n"
        f"  FT model:     [cyan]{ft_model}[/cyan]\n"
        f"  Attacks:      [cyan]{len(attack_classes)} attack(s)[/cyan]"
    )

    profiler = BaselineProfiler()

    with console.status("Profiling base model..."):
        base_adapter = ModelAdapter(model_path=base_model, adapter_type=adapter)
        baseline = profiler.profile_base_model(
            adapter=base_adapter,
            attack_classes=attack_classes,
        )

    console.print(f"[green]✓ Baseline profile stored:[/green] {baseline.id}")

    with console.status("Profiling fine-tuned model and computing deltas..."):
        ft_adapter = ModelAdapter(model_path=ft_model, adapter_type=adapter)
        delta_report = profiler.profile_fine_tuned_model(
            adapter=ft_adapter,
            baseline_id=baseline.id,
            attack_classes=attack_classes,
            ft_data_path=ft_data,
        )

    # Print summary table
    table = Table(title="Vulnerability Delta Report", show_header=True)
    table.add_column("Attack", style="cyan")
    table.add_column("Base Rate", justify="right")
    table.add_column("FT Rate", justify="right")
    table.add_column("Delta", justify="right")
    table.add_column("Classification", justify="center")
    table.add_column("Attribution")

    CLASS_COLORS = {"FIXED": "green", "REGRESSED": "red", "NEUTRAL": "yellow"}
    for d in delta_report.deltas:
        cls_color = CLASS_COLORS.get(d.classification, "white")
        delta_str = f"{d.delta_success_rate:+.3f}"
        table.add_row(
            d.attack_name,
            f"{d.base_result.get('success_rate', 0.0):.3f}",
            f"{d.ft_result.get('success_rate', 0.0):.3f}",
            delta_str,
            f"[{cls_color}]{d.classification}[/{cls_color}]",
            d.attribution.source if d.attribution else "unknown",
        )

    console.print(table)

    s = delta_report.summary
    console.print(
        f"\n[bold]Summary:[/bold] "
        f"[green]{s['fixed']} FIXED[/green]  "
        f"[red]{s['regressed']} REGRESSED[/red]  "
        f"[yellow]{s['neutral']} NEUTRAL[/yellow]"
    )

    report_dict = delta_report.to_dict()
    if output:
        output.write_text(json.dumps(report_dict, indent=2))
        console.print(f"\n[green]✓ Delta report written to {output}[/green]")
    else:
        console.print("\n[dim]Tip: use --output delta_report.json to save[/dim]")


# ---------------------------------------------------------------------------
# Gap #2: sichgate compliance quantization
# ---------------------------------------------------------------------------

@compliance_app.command(name="quantization")
def quantization(
    model: str = typer.Option(..., "--model", "-m", help="Model path or HuggingFace ID"),
    quantization_levels: str = typer.Option("Q4,Q8,FP16", "--quantization-levels", help="Comma-separated levels: Q4,Q8,INT8,FP16"),
    attacks: str = typer.Option("all-llm", "--attacks", help="Comma-separated attack names or group"),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write variance report JSON to file"),
    verbose: bool = typer.Option(False, "--verbose", "-v"),
) -> None:
    """
    Gap #2: Quantization Robustness Testing.

    Tests a model at multiple quantization levels (Q4, Q8, FP16) and generates
    a variance report showing how attack success rates change with precision.

    Example:
        sichgate compliance quantization \\
          --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 \\
          --quantization-levels Q4,Q8,FP16 \\
          --attacks jailbreak,prompt-injection \\
          --output variance_report.json
    """
    import logging
    if verbose:
        logging.getLogger().setLevel(logging.INFO)

    from sichgate.runner import build_registry
    from sichgate.testing.quantization_matrix import QuantizationTestMatrix

    registry = build_registry()
    attack_names = [a.strip() for a in attacks.split(",")]
    attack_classes = []
    for name in attack_names:
        if name in registry:
            for cls in registry[name]:
                if cls not in attack_classes:
                    attack_classes.append(cls)
        else:
            console.print(f"[red]Unknown attack: {name}[/red]")
            raise typer.Exit(1)

    levels = [l.strip() for l in quantization_levels.split(",")]

    console.print(
        f"[bold]SichGate[/bold] Gap #2: Quantization Robustness Testing\n"
        f"  Model:  [cyan]{model}[/cyan]\n"
        f"  Levels: [cyan]{levels}[/cyan]\n"
        f"  Attacks: [cyan]{len(attack_classes)} attack(s)[/cyan]"
    )

    matrix = QuantizationTestMatrix(model)

    with console.status(f"Running attack matrix across {len(levels)} quantization levels..."):
        try:
            report = matrix.execute_attack_matrix(attack_classes, levels=levels)
        except Exception as exc:
            console.print(f"[red]Quantization testing failed: {exc}[/red]")
            raise typer.Exit(1)

    # Print variance table
    table = Table(title="Quantization Variance Report", show_header=True)
    table.add_column("Attack", style="cyan")
    for level in report.quantization_variants:
        table.add_column(level, justify="right")
    table.add_column("Max Delta", justify="right")
    table.add_column("Significant?", justify="center")

    for row in report.results:
        cells = [row.attack_name]
        for level in report.quantization_variants:
            qr = row.by_quantization.get(level)
            cells.append(f"{qr.success_rate:.3f}" if qr else "N/A")
        delta_str = f"{row.variance.max_delta:.3f}"
        sig = "[red]YES[/red]" if row.variance.significant else "[green]no[/green]"
        cells += [delta_str, sig]
        table.add_row(*cells)

    console.print(table)
    console.print(f"\n[bold]Recommendation:[/bold] {report.recommendation}")

    if report.high_variance_attacks:
        console.print(
            f"[yellow]High-variance attacks:[/yellow] {', '.join(report.high_variance_attacks)}"
        )

    report_dict = report.to_dict()
    if output:
        output.write_text(json.dumps(report_dict, indent=2))
        console.print(f"\n[green]✓ Variance report written to {output}[/green]")
    else:
        console.print("\n[dim]Tip: use --output variance_report.json to save[/dim]")


# ---------------------------------------------------------------------------
# Gap #3: sichgate compliance audit-trail
# ---------------------------------------------------------------------------

@compliance_app.command(name="audit-trail")
def audit_trail(
    run_id: Optional[str] = typer.Option(None, "--test-run-id", help="Run ID to create audit trail from"),
    list_findings: bool = typer.Option(False, "--list", help="List existing audit trail findings"),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write audit trail JSON to file"),
) -> None:
    """
    Gap #3: Generate or list compliance audit trail entries.

    Creates standards-mapped audit trail entries from a completed test run,
    or lists existing findings.

    Example:
        sichgate compliance audit-trail --test-run-id run_20260226_001 --output audit_trail.json
        sichgate compliance audit-trail --list
    """
    from sichgate.database import DatabaseManager
    from sichgate.compliance.mapper import ComplianceMapper

    db = DatabaseManager()
    mapper = ComplianceMapper()

    if list_findings:
        findings = db.list_compliance_findings()
        if not findings:
            console.print("[yellow]No compliance findings found. Run a test first.[/yellow]")
            return

        table = Table(title="Compliance Audit Trail", show_header=True)
        table.add_column("Finding ID", style="cyan", no_wrap=True)
        table.add_column("Attack")
        table.add_column("Vuln Type")
        table.add_column("Status", justify="center")
        table.add_column("Created At")

        STATUS_COLORS = {
            "NEW": "red", "MITIGATED": "yellow",
            "RESOLVED": "green", "ACCEPTED_RISK": "dim"
        }
        for f in findings:
            status = f.get("status", "NEW")
            color = STATUS_COLORS.get(status, "white")
            table.add_row(
                f.get("finding_id", "")[:20],
                f.get("attack_name", ""),
                f.get("vulnerability_type", ""),
                f"[{color}]{status}[/{color}]",
                (f.get("created_at") or "")[:19],
            )
        console.print(table)
        return

    if not run_id:
        console.print("[red]Error: provide --test-run-id or --list[/red]")
        raise typer.Exit(1)

    run = db.get_run(run_id)
    if not run:
        console.print(f"[red]Run '{run_id}' not found.[/red]")
        raise typer.Exit(1)

    findings = mapper.map_run(run)
    _persist_findings(findings)

    console.print(
        f"[green]✓ Audit trail created:[/green] "
        f"{len(findings)} finding(s) mapped to compliance standards"
    )

    if output:
        output.write_text(json.dumps([f["trail"] for f in findings], indent=2))
        console.print(f"[green]✓ Audit trail written to {output}[/green]")


# ---------------------------------------------------------------------------
# Gap #3: sichgate compliance mitigate
# ---------------------------------------------------------------------------

@compliance_app.command(name="mitigate")
def mitigate(
    finding_id: str = typer.Option(..., "--finding-id", help="Finding ID from audit trail"),
    mitigation_desc: str = typer.Option(..., "--mitigation", help="Description of the mitigation applied"),
    mitigation_type: str = typer.Option("other", "--mitigation-type", help="Type: content_filter | guardrails | fine-tuning | prompt_engineering | other"),
    added_by: str = typer.Option("compliance_officer", "--added-by", help="Person adding the mitigation"),
    retest_schedule: Optional[str] = typer.Option(None, "--retest-schedule", help="ISO 8601 date for retest (e.g. 2026-03-01)"),
) -> None:
    """
    Gap #3: Record a mitigation against a compliance finding.

    Example:
        sichgate compliance mitigate \\
          --finding-id find_abc123 \\
          --mitigation "Added output validator rejecting jailbreak patterns" \\
          --mitigation-type content_filter \\
          --retest-schedule 2026-03-01
    """
    import uuid
    from datetime import datetime
    from sichgate.database import get_session_factory, ComplianceAuditTrailRecord, MitigationRecord

    SessionFactory = get_session_factory()
    with SessionFactory() as session:
        trail = session.query(ComplianceAuditTrailRecord).filter_by(
            finding_id=finding_id
        ).first()
        if not trail:
            console.print(f"[red]Finding '{finding_id}' not found.[/red]")
            raise typer.Exit(1)

        retest_dt = None
        if retest_schedule:
            try:
                retest_dt = datetime.fromisoformat(retest_schedule)
            except ValueError:
                console.print("[red]Invalid retest-schedule date. Use ISO 8601 (e.g. 2026-03-01).[/red]")
                raise typer.Exit(1)

        mitigation = MitigationRecord(
            id=str(uuid.uuid4()),
            finding_id=finding_id,
            added_by=added_by,
            description=mitigation_desc,
            mitigation_type=mitigation_type,
            applied_at=datetime.utcnow(),
            retest_scheduled_at=retest_dt,
        )
        session.add(mitigation)

        if trail.status == "NEW":
            trail.status = "MITIGATED"
        session.commit()

    console.print(
        f"[green]✓ Mitigation recorded for {finding_id}[/green]\n"
        f"  Type: {mitigation_type}\n"
        f"  Status: MITIGATED"
        + (f"\n  Retest: {retest_schedule}" if retest_schedule else "")
    )


# ---------------------------------------------------------------------------
# Gap #3: sichgate compliance export-report
# ---------------------------------------------------------------------------

@compliance_app.command(name="export-report")
def export_report(
    output: Path = typer.Option(
        Path("compliance_report.json"),
        "--output", "-o",
        help="Output path (use .pdf suffix to attempt PDF generation)"
    ),
    fmt: str = typer.Option("json", "--format", help="Output format: json | pdf"),
    limit: int = typer.Option(500, "--limit", help="Maximum number of findings to include"),
) -> None:
    """
    Gap #3: Export a compliance audit report.

    Generates a structured report of all findings with standards mapping,
    mitigation history, and sign-off status.

    PDF export requires weasyprint (pip install weasyprint) or pandoc.

    Example:
        sichgate compliance export-report --format pdf --output compliance_report.pdf
        sichgate compliance export-report --output compliance_report.json
    """
    from datetime import datetime, timezone
    from sichgate.database import DatabaseManager
    from sichgate.compliance.api import _render_pdf, _build_html_report

    db = DatabaseManager()
    findings = db.list_compliance_findings(limit=limit)

    report = {
        "report_type": "compliance_audit_trail",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "generated_by": "SichGate Pro",
        "total_findings": len(findings),
        "findings": findings,
    }

    if fmt == "pdf" or str(output).endswith(".pdf"):
        from pathlib import Path as P
        pdf_path = _render_pdf(report, P(output).parent)
        if pdf_path and pdf_path.exists():
            import shutil
            shutil.move(str(pdf_path), str(output))
            console.print(f"[green]✓ PDF report written to {output}[/green]")
        else:
            console.print(
                "[yellow]PDF generation unavailable. "
                "Install weasyprint (pip install weasyprint) or pandoc.[/yellow]"
            )
            json_path = Path(str(output).replace(".pdf", ".json"))
            json_path.write_text(json.dumps(report, indent=2))
            console.print(f"[yellow]Fallback: JSON report written to {json_path}[/yellow]")
    else:
        output.write_text(json.dumps(report, indent=2))
        console.print(
            f"[green]✓ Compliance report written to {output}[/green]\n"
            f"  Findings: {len(findings)}"
        )


# ---------------------------------------------------------------------------
# Gap #4: sichgate compliance composition-test
# ---------------------------------------------------------------------------

@compliance_app.command(name="composition-test")
def composition_test(
    template: Optional[str] = typer.Option(
        None, "--template",
        help="Built-in template name: healthcare_intake | document_router | retrieval_pipeline"
    ),
    custom_template: Optional[Path] = typer.Option(None, "--custom-template", help="Path to custom YAML template"),
    model_1: Optional[str] = typer.Option(None, "--model-1", help="Override Stage 1 model path"),
    model_2: Optional[str] = typer.Option(None, "--model-2", help="Override Stage 2 model path"),
    model_3: Optional[str] = typer.Option(None, "--model-3", help="Override Stage 3 model path"),
    attacks: str = typer.Option(
        "cascade_failure,injection_propagation,data_leakage",
        "--attacks",
        help="Comma-separated attack names: cascade_failure | injection_propagation | data_leakage | all"
    ),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write composition report JSON to file"),
    verbose: bool = typer.Option(False, "--verbose", "-v"),
) -> None:
    """
    Gap #4: Multi-Model Composition Testing.

    Tests a multi-model pipeline for cascade failures, prompt injection
    propagation, and cross-stage data leakage.

    Example:
        sichgate compliance composition-test \\
          --template healthcare_intake \\
          --model-1 ./classifier \\
          --model-2 ./extractor \\
          --model-3 ./risk_scorer \\
          --attacks cascade_failure,injection_propagation,data_leakage \\
          --output composition_report.json
    """
    import logging
    if verbose:
        logging.getLogger().setLevel(logging.INFO)

    from sichgate.composition.pipeline import CompositionTemplate, CompositionPipeline
    from sichgate.composition.attacks import (
        CascadeFailureAttack,
        DataLeakageAcrossStagesAttack,
        PromptInjectionPropagationAttack,
    )

    # Load template
    if custom_template:
        comp_template = CompositionTemplate.from_yaml(str(custom_template))
    elif template:
        try:
            comp_template = CompositionTemplate.load_builtin(template)
        except ValueError as exc:
            console.print(f"[red]{exc}[/red]")
            raise typer.Exit(1)
    else:
        console.print("[red]Provide --template or --custom-template[/red]")
        raise typer.Exit(1)

    # Override model paths if provided
    model_overrides = [m for m in [model_1, model_2, model_3] if m is not None]
    if model_overrides:
        try:
            comp_template = comp_template.with_models(model_overrides)
        except ValueError as exc:
            console.print(f"[red]{exc}[/red]")
            raise typer.Exit(1)

    console.print(
        f"[bold]SichGate[/bold] Gap #4: Multi-Model Composition Testing\n"
        f"  Template: [cyan]{comp_template.name}[/cyan]\n"
        f"  Stages:   [cyan]{[s.stage_name for s in comp_template.stages]}[/cyan]"
    )

    # Resolve attack classes
    ATTACK_MAP = {
        "cascade_failure": CascadeFailureAttack,
        "injection_propagation": PromptInjectionPropagationAttack,
        "data_leakage": DataLeakageAcrossStagesAttack,
    }
    requested = [a.strip() for a in attacks.split(",")]
    attack_classes = []
    if "all" in requested:
        attack_classes = list(ATTACK_MAP.values())
    else:
        for name in requested:
            if name in ATTACK_MAP:
                attack_classes.append(ATTACK_MAP[name])
            else:
                console.print(f"[red]Unknown composition attack: {name}[/red]")
                console.print(f"Available: {list(ATTACK_MAP.keys())} or 'all'")
                raise typer.Exit(1)

    pipeline = CompositionPipeline(comp_template)

    with console.status("Running composition attacks..."):
        report = pipeline.run_composition_attacks(attack_classes=attack_classes)

    # Print results table
    table = Table(title=f"Composition Report: {comp_template.name}", show_header=True)
    table.add_column("Attack", style="cyan")
    table.add_column("Status", justify="center")
    table.add_column("Severity", justify="center")
    table.add_column("Notes")

    STATUS_COLORS = {"pass": "green", "fail": "red", "warning": "yellow", "error": "dim"}
    SEV_COLORS = {"critical": "bold red", "high": "red", "medium": "yellow"}

    for attack_result in report.get("attacks", []):
        if "error" in attack_result and "status" not in attack_result:
            table.add_row(
                attack_result.get("attack", ""),
                "[dim]error[/dim]", "", attack_result.get("error", "")[:80]
            )
            continue

        status = attack_result.get("status", "unknown")
        severity = attack_result.get("severity")
        notes = attack_result.get("notes", "")[:80]

        status_str = f"[{STATUS_COLORS.get(status, 'white')}]{status}[/]"
        severity_str = (
            f"[{SEV_COLORS.get(severity, 'white')}]{severity}[/]" if severity else "-"
        )
        table.add_row(
            attack_result.get("attack", ""),
            status_str, severity_str,
            notes + "..." if len(attack_result.get("notes", "")) > 80 else notes,
        )

    console.print(table)

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


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _persist_findings(findings: list) -> None:
    """Persist audit trail entries to the database."""
    import uuid
    from datetime import datetime
    from sichgate.database import get_session_factory, ComplianceAuditTrailRecord

    SessionFactory = get_session_factory()
    with SessionFactory() as session:
        for item in findings:
            trail = item["trail"]
            standards = item["standards"]

            existing = session.query(ComplianceAuditTrailRecord).filter_by(
                finding_id=trail["finding_id"]
            ).first()
            if existing:
                continue

            rec = ComplianceAuditTrailRecord(
                id=trail["id"],
                finding_id=trail["finding_id"],
                run_id=trail.get("run_id"),
                attack_type=trail["attack_type"],
                attack_name=trail["attack_name"],
                seed=trail.get("seed"),
                payload=trail.get("payload"),
                vulnerability_type=trail["vulnerability_type"],
                status="NEW",
                created_at=datetime.utcnow(),
                standards_json=json.dumps(standards),
            )
            session.add(rec)
        session.commit()
