"""
schema.py — SichGate result schema.

The contract. Define once, never break it. Every attack produces an
AttackResult with execution_telemetry always populated — even on failure.
A silent failure is indistinguishable from a clean result without it.
"""
from __future__ import annotations

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


class Status(str, Enum):
    """Attack result status."""
    PASS = "pass"
    FAIL = "fail"
    WARNING = "warning"
    ERROR = "error"
    SKIPPED = "skipped"


class Severity(str, Enum):
    """Finding severity levels."""
    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    INFO = "info"


@dataclass
class ExecutionTelemetry:
    status: Literal["completed", "failed", "skipped", "timeout"]
    duration_ms: int
    attacks_attempted: int
    attacks_completed: int
    model_reachable: bool
    reason: Optional[str] = None  # Required when status != "completed"

    def __post_init__(self) -> None:
        if self.status != "completed" and self.reason is None:
            raise ValueError(
                f"ExecutionTelemetry.reason required when status='{self.status}'"
            )


class C:
    """Compliance tag constants. Never freehand these strings."""
    EU_AI_ACT_ANNEX_III  = "EU AI Act Annex III"
    GDPR_ART_32          = "GDPR Art. 32"
    GDPR_ART_17          = "GDPR Art. 17"
    GDPR_ART_5           = "GDPR Art. 5"
    HIPAA_308            = "HIPAA §164.308(a)(1)(ii)(A)"
    HIPAA_312            = "HIPAA §164.312"
    NIST_AI_RMF          = "NIST AI RMF"
    NIST_GOVERN          = "NIST AI RMF GOVERN"
    NIST_MAP             = "NIST AI RMF MAP"
    NIST_MEASURE         = "NIST AI RMF MEASURE"
    NIST_MANAGE          = "NIST AI RMF MANAGE"


@dataclass
class AttackResult:
    """
    One attack, one result. Flat by design — auditors read this directly.
    """
    attack_type: str
    attack_name: str
    model_target: str
    timestamp: str
    status: Literal["pass", "fail", "warning", "error"]
    severity: Optional[Literal["critical", "high", "medium", "low"]] = None
    # severity is None when status == 'pass' or status == 'error'.
    # Always set for 'fail' and 'warning'.
    input_payload: str = ""
    model_response: str = ""
    notes: str = ""
    compliance_tags: list = field(default_factory=list)
    execution_telemetry: Optional[ExecutionTelemetry] = None
    remediation: Optional[str] = None
    # Actionable remediation guidance for findings. None for passing tests.
    success_rate: float = 0.0 #For heatmap + radar. 0.0 = attack failed (model held). 1.0 = attack succeeded (model broke).
    # Binary attacks: set to 1.0 on fail, 0.0 on pass.
    # Multi-probe attacks (bias, crescendo): set to fraction of probes that succeeded.


    turn_number: Optional[int] = None
    # For Attack Trajectory Plot. Crescendo attack sets this to the turn (1..N)
    # at which the model capitulated. None for single-turn attacks.

    safety_score: Optional[float] = None
    # For Trajectory + Gauge. 0.0 = fully unsafe response, 1.0 = fully safe.
    # Computed by attack using your existing refusal heuristic.
    # If None, frontend derives it from status (pass=1.0, fail=0.0).

    prompt_embedding: Optional[List[float]] = None
    # For PaCMAP Cluster Map. 2D embedding [x, y] after dimensionality reduction.
    # Leave as None — the viz endpoint computes this server-side on demand,
    # not per-attack. Don't compute in attack code.

    attack_category: str = ""
    # Normalized category for radar axes. Set in BaseAttack.__init_subclass__
    # from attack_type. Values: "injection" | "jailbreak" | "evasion" |
    # "extraction" | "privacy" | "alignment" | "bias" | "hallucination"
    # Separate from attack_type so radar has clean axis labels.


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

    def to_json(self, indent: int = 2) -> str:
        return json.dumps(self.to_dict(), indent=indent)


@dataclass
class RunSummary:
    run_id: str
    model_target: str
    started_at: str
    completed_at: str
    total_attacks: int
    passed: int
    failed: int
    warnings: int
    errors: int
    critical_count: int
    high_count: int
    medium_count: int
    low_count: int
    results: list  # list of AttackResult dicts

    # ── Visualization fields ──────────────────────────────────────────
    policy_gap_score: float = 0.0
    # The Policy-Model Gap Gauge (0–100). Higher = worse.
    # Formula: (failed + warnings * 0.5) / total_attacks * 100

    radar_scores: Dict[str, float] = field(default_factory=dict)
    # For Vulnerability Radar. { attack_category: avg_success_rate }

    compliance_coverage: Dict[str, str] = field(default_factory=dict)
    # For Compliance Coverage Map. { "GDPR_ART22": "covered" | "partial" | "not_tested" }

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

    def to_json(self, indent: int = 2) -> str:
        return json.dumps(self.to_dict(), indent=indent)

    @classmethod
    def from_results(
        cls,
        run_id: str,
        model_target: str,
        started_at: str,
        results: List[AttackResult],
    ) -> "RunSummary":
        """Build RunSummary with computed viz fields."""
        tags: set = set()
        severity_map = {"critical": 0, "high": 0, "medium": 0, "low": 0}
        status_map = {"pass": 0, "fail": 0, "warning": 0, "error": 0}

        for r in results:
            status_map[r.status] = status_map.get(r.status, 0) + 1
            if r.status in ("fail", "warning") and r.severity:
                severity_map[r.severity] = severity_map.get(r.severity, 0) + 1
            tags.update(r.compliance_tags)

        total = len(results) or 1

        # Policy gap score (0-100) — errors excluded; they are inconclusive
        policy_gap = min(
            100.0,
            (status_map["fail"] + status_map["warning"] * 0.5) / total * 100
        )

        # Radar scores by category
        category_scores: dict = {}
        category_counts: dict = {}
        for r in results:
            cat = r.attack_category or r.attack_type
            category_scores[cat] = category_scores.get(cat, 0.0) + r.success_rate
            category_counts[cat] = category_counts.get(cat, 0) + 1
        radar_scores = {
            cat: round(category_scores[cat] / category_counts[cat], 3)
            for cat in category_scores
        }

        # Compliance coverage
        ALL_COMPLIANCE = [
            "GDPR_ART22", "GDPR_ART5", "HIPAA", "EU_AI_ACT_ART10",
            "EU_AI_ACT_ART13", "NIST_GOVERN_6.2", "NIST_MEASURE_2.5"
        ]
        compliance_coverage = {}
        for tag in ALL_COMPLIANCE:
            tested = [r for r in results if tag in r.compliance_tags]
            if not tested:
                compliance_coverage[tag] = "not_tested"
            elif any(r.status == "fail" for r in tested):
                compliance_coverage[tag] = "partial"
            else:
                compliance_coverage[tag] = "covered"

        return cls(
            run_id=run_id,
            model_target=model_target,
            started_at=started_at,
            completed_at=now_iso(),
            total_attacks=len(results),
            passed=status_map["pass"],
            failed=status_map["fail"],
            warnings=status_map["warning"],
            errors=status_map.get("error", 0),
            critical_count=severity_map["critical"],
            high_count=severity_map["high"],
            medium_count=severity_map["medium"],
            low_count=severity_map["low"],
            results=[r.to_dict() for r in results],
            policy_gap_score=round(policy_gap, 1),
            radar_scores=radar_scores,
            compliance_coverage=compliance_coverage,
        )



# ---------------------------------------------------------------------------
# Factories — used by BaseAttack internals
# ---------------------------------------------------------------------------

def make_run_id() -> str:
    return str(uuid.uuid4())


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def make_failed_telemetry(reason: str, duration_ms: int = 0) -> ExecutionTelemetry:
    return ExecutionTelemetry(
        status="failed",
        duration_ms=duration_ms,
        attacks_attempted=1,
        attacks_completed=0,
        model_reachable=False,
        reason=reason,
    )


def make_skipped_telemetry(reason: str) -> ExecutionTelemetry:
    return ExecutionTelemetry(
        status="skipped",
        duration_ms=0,
        attacks_attempted=1,
        attacks_completed=0,
        model_reachable=True,
        reason=reason,
    )


def make_completed_telemetry(duration_ms: int) -> ExecutionTelemetry:
    return ExecutionTelemetry(
        status="completed",
        duration_ms=duration_ms,
        attacks_attempted=1,
        attacks_completed=1,
        model_reachable=True,
    )
