# sichgate/reporting.py
"""
Compliance-auditor-ready report generation.

Produces JSON output suitable for:
- Legal/compliance officer review
- Regulatory body submission (EU AI Act conformity assessment)
- GRC tool integration (JIRA, Servicenow)
- Third-party auditor validation

Key principles:
1. Every finding links to specific regulatory articles
2. All evidence is verbatim (input_payload, model_response)
3. Reports are reproducible (include model version, test date, timestamps)
4. Remediation is actionable (not "fix the model" but HOW)
"""

from dataclasses import dataclass, asdict, field
from typing import List, Dict, Optional, Literal
from datetime import datetime
from enum import Enum

from sichgate.schema import AttackResult, Status, Severity
from sichgate.compliance_registry import (
    ComplianceRef,
    ids_to_refs,
    list_all_refs,
    frameworks_covered,
)


@dataclass
class RegulatoryContext:
    """Metadata about which regulatory frameworks apply and their status."""
    applicable_frameworks: List[str]           # ["EU AI Act", "GDPR", "NIST AI RMF", ...]
    annex_iii_applicable: bool = False         # Is this a high-risk biometric system?
    key_deadlines: Dict[str, str] = field(default_factory=dict)

    @staticmethod
    def default() -> "RegulatoryContext":
        """Standard regulatory context for 2026."""
        return RegulatoryContext(
            applicable_frameworks=[
                "EU AI Act",
                "GDPR",
                "NIST AI RMF",
                "OWASP LLM Top 10",
                "ISO/IEC 42001",
            ],
            annex_iii_applicable=False,
            key_deadlines={
                "EU_AI_Act_high_risk": "2026-08-02",
                "GDPR": "in_force",
                "NIST_AI_RMF": "voluntary",
                "ISO_42001": "certification_available",
            },
        )


@dataclass
class FindingSummary:
    """Count of findings by severity."""
    critical: int = 0
    high: int = 0
    medium: int = 0
    low: int = 0
    passed_tests: int = 0

    def total(self) -> int:
        return self.critical + self.high + self.medium + self.low

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass
class ComplianceCoverage:
    """Which compliance requirements were tested."""
    tested_refs: List[str]                      # List of ref IDs that appeared in findings
    untested_refs: List[str]                    # List of ref IDs that were not tested

    def coverage_percentage(self) -> float:
        """What % of known compliance refs had at least one test?"""
        all_refs = list_all_refs()
        if not all_refs:
            return 0.0
        return 100 * len(self.tested_refs) / len(all_refs)


@dataclass
class SichGateReport:
    """
    Complete compliance-ready report from a test run.

    This is the artifact that goes to:
    - Compliance officer for conformity assessment
    - Auditor for third-party validation
    - GRC tool for integration
    """

    # ── Report metadata ────────────────────────────────────────
    report_id: str                              # Unique identifier
    generated_by: str                           # "SichGate Pro v1.0.0"
    generated_at: str                           # ISO 8601 timestamp
    license_key_hash: str                       # Hash of license key (for audit trail)

    # ── Target model information ───────────────────────────────
    model_target: str                           # Model name or path tested
    model_version: Optional[str] = None         # Version string if available
    parameters_millions: Optional[int] = None   # Model size (80M, 600M, etc.)

    # ── Regulatory context ─────────────────────────────────────
    regulatory_context: RegulatoryContext = field(default_factory=RegulatoryContext.default)

    # ── Test execution metadata ────────────────────────────────
    test_date_start: str = ""                   # ISO 8601
    test_date_end: str = ""                     # ISO 8601
    total_attacks_run: int = 0
    total_attacks_passed: int = 0

    # ── Findings summary ───────────────────────────────────────
    finding_summary: FindingSummary = field(default_factory=FindingSummary)

    # ── Compliance coverage ────────────────────────────────────
    compliance_coverage: Optional[ComplianceCoverage] = None

    # ── All attack results ─────────────────────────────────────
    results: List[AttackResult] = field(default_factory=list)

    def to_dict(self) -> dict:
        """Convert to JSON-serialisable dict."""
        d = {
            "report_id": self.report_id,
            "generated_by": self.generated_by,
            "generated_at": self.generated_at,
            "license_key_hash": self.license_key_hash,
            "model_target": self.model_target,
            "model_version": self.model_version,
            "parameters_millions": self.parameters_millions,
            "regulatory_context": asdict(self.regulatory_context),
            "test_date_start": self.test_date_start,
            "test_date_end": self.test_date_end,
            "total_attacks_run": self.total_attacks_run,
            "total_attacks_passed": self.total_attacks_passed,
            "finding_summary": self.finding_summary.to_dict(),
            "compliance_coverage": asdict(self.compliance_coverage) if self.compliance_coverage else None,
            "results": [r.to_dict() for r in self.results],
        }
        return d

    def to_json_string(self) -> str:
        """Convert to JSON string (for file export)."""
        import json
        return json.dumps(self.to_dict(), indent=2)

    @staticmethod
    def from_attack_results(
        results: List[AttackResult],
        model_target: str,
        report_id: str,
        license_key_hash: str = "DEMO",
        model_version: Optional[str] = None,
        parameters_millions: Optional[int] = None,
        regulatory_context: Optional[RegulatoryContext] = None,
    ) -> "SichGateReport":
        """Build a report from attack results."""

        # Compute summary
        summary = FindingSummary()
        tested_refs = set()
        for result in results:
            tested_refs.update(result.compliance_tags)
            if result.status == Status.PASS:
                summary.passed_tests += 1
            elif result.severity:
                if result.severity == Severity.CRITICAL:
                    summary.critical += 1
                elif result.severity == Severity.HIGH:
                    summary.high += 1
                elif result.severity == Severity.MEDIUM:
                    summary.medium += 1
                elif result.severity == Severity.LOW:
                    summary.low += 1

        # Compute compliance coverage
        all_ref_ids = set(r.id for r in list_all_refs())
        tested_ref_list = sorted(tested_refs)
        untested_ref_list = sorted(all_ref_ids - tested_refs)
        coverage = ComplianceCoverage(
            tested_refs=tested_ref_list,
            untested_refs=untested_ref_list,
        )

        # Set timestamps
        now = datetime.utcnow().isoformat() + "Z"
        timestamps = [r.timestamp for r in results if r.timestamp]
        test_date_start = min(timestamps) if timestamps else now
        test_date_end = max(timestamps) if timestamps else now

        return SichGateReport(
            report_id=report_id,
            generated_by="SichGate Pro v1.0.0",
            generated_at=now,
            license_key_hash=license_key_hash,
            model_target=model_target,
            model_version=model_version,
            parameters_millions=parameters_millions,
            regulatory_context=regulatory_context or RegulatoryContext.default(),
            test_date_start=test_date_start,
            test_date_end=test_date_end,
            total_attacks_run=len(results),
            total_attacks_passed=summary.passed_tests,
            finding_summary=summary,
            compliance_coverage=coverage,
            results=results,
        )

    def expanded_compliance_coverage(self) -> Dict[str, List[ComplianceRef]]:
        """Return full ComplianceRef objects grouped by framework."""
        if not self.compliance_coverage:
            return {}

        all_refs_by_id = {r.id: r for r in list_all_refs()}
        refs_by_framework: Dict[str, List[ComplianceRef]] = {}

        for ref_id in self.compliance_coverage.tested_refs:
            ref = all_refs_by_id.get(ref_id)
            if ref:
                framework = ref.framework
                if framework not in refs_by_framework:
                    refs_by_framework[framework] = []
                refs_by_framework[framework].append(ref)

        return refs_by_framework


def create_bom(
    reports: List[SichGateReport],
    sichgate_version: str = "1.0.0",
    license_key_hash: str = "DEMO",
    test_machine_id: str = "local",
) -> dict:
    """
    Create an AI Bill of Materials (BOM) from multiple reports.

    This is submitted alongside conformity assessment documentation
    to regulatory bodies. It provides a complete view of:
    - What models were tested
    - What was tested
    - Trending (are issues improving?)
    """

    # Group by model
    models_tested = {}
    all_findings = []

    for report in reports:
        model = report.model_target
        if model not in models_tested:
            models_tested[model] = {
                "reports": [],
                "total_findings": FindingSummary(),
            }
        models_tested[model]["reports"].append(report)

        for result in report.results:
            if result.status != Status.PASS:
                all_findings.append(result)
                # Accumulate severity counts
                if result.severity == Severity.CRITICAL:
                    models_tested[model]["total_findings"].critical += 1
                elif result.severity == Severity.HIGH:
                    models_tested[model]["total_findings"].high += 1
                elif result.severity == Severity.MEDIUM:
                    models_tested[model]["total_findings"].medium += 1
                elif result.severity == Severity.LOW:
                    models_tested[model]["total_findings"].low += 1

    # Compute trending
    findings_trend = {}
    for model, data in models_tested.items():
        reports = sorted(data["reports"], key=lambda r: r.test_date_end)
        if len(reports) >= 2:
            first_summary = reports[0].finding_summary
            latest_summary = reports[-1].finding_summary
            first_total = first_summary.total()
            latest_total = latest_summary.total()
            if latest_total < first_total:
                trend = "improving"
            elif latest_total > first_total:
                trend = "degrading"
            else:
                trend = "stable"
            findings_trend[model] = {
                "first_run_date": reports[0].test_date_end,
                "first_findings": first_summary.to_dict(),
                "latest_run_date": reports[-1].test_date_end,
                "latest_findings": latest_summary.to_dict(),
                "trend": trend,
            }

    return {
        "bom_generated_at": datetime.utcnow().isoformat() + "Z",
        "sichgate_version": sichgate_version,
        "license_key_hash": license_key_hash,
        "test_machine_id": test_machine_id,
        "models_tested": {
            model: {
                "parameters_millions": data["reports"][-1].parameters_millions,
                "num_reports": len(data["reports"]),
                "total_findings_all_time": data["total_findings"].to_dict(),
                "test_date_range": {
                    "start": min(r.test_date_start for r in data["reports"]),
                    "end": max(r.test_date_end for r in data["reports"]),
                },
            }
            for model, data in models_tested.items()
        },
        "findings_trend": findings_trend,
        "compliance_frameworks_covered": frameworks_covered(),
        "total_models_tested": len(models_tested),
        "total_findings_all_models": sum(d["total_findings"].total() for d in models_tested.values()),
    }
