"""Baseline management and regression detection.

A *baseline* is a snapshot of the security posture of a target at a point in
time.  It records which attacks succeeded and which were blocked.  After a
model/prompt/guardrail change the same attack suite is re-run and the results
are compared against the baseline to produce a regression diff:

- **Regressions** — attacks that were blocked in the baseline but now succeed.
- **Improvements** — attacks that succeeded in the baseline but are now blocked.
- **Unchanged** — same outcome.
"""

from __future__ import annotations

import json
from datetime import datetime
from pathlib import Path
from typing import Any

from pydantic import BaseModel, Field

from src.utils.logging import get_logger

logger = get_logger(__name__)

BASELINES_DIR = Path("baselines")


class AttackOutcome(BaseModel):
    """Fingerprint of a single attack and its outcome."""

    attack_id: str
    agent_name: str
    attack_type: str
    query_hash: str
    succeeded: bool
    finding_severity: str = ""
    finding_category: str = ""


class BaselineRecord(BaseModel):
    """Persisted baseline snapshot."""

    name: str
    created_at: str
    session_id: str
    target_name: str
    vulnerability_score: float
    total_attacks: int
    successful_attacks: int
    outcomes: list[AttackOutcome]
    metadata: dict[str, Any] = Field(default_factory=dict)


class RegressionDiff(BaseModel):
    """Single regression or improvement entry."""

    attack_type: str
    agent_name: str
    query_preview: str
    baseline_outcome: str
    current_outcome: str
    severity: str = ""


class RegressionResult(BaseModel):
    """Full regression comparison report."""

    baseline_name: str
    baseline_score: float
    current_score: float
    score_delta: float
    regressions: list[RegressionDiff] = Field(default_factory=list)
    improvements: list[RegressionDiff] = Field(default_factory=list)
    unchanged_blocked: int = 0
    unchanged_succeeded: int = 0
    has_regressions: bool = False


class BaselineManager:
    """Create, store, list and compare baselines."""

    def __init__(self, baselines_dir: Path | None = None):
        self.baselines_dir = baselines_dir or BASELINES_DIR
        self.baselines_dir.mkdir(parents=True, exist_ok=True)

    def create_from_session(
        self,
        session: dict[str, Any],
        name: str | None = None,
    ) -> BaselineRecord:
        """Create a baseline from a completed session."""
        attacks = session.get("attack_attempts", [])
        findings = session.get("security_findings", [])
        finding_attack_ids = {f.get("attack_id") for f in findings}

        severity_map: dict[str, str] = {}
        category_map: dict[str, str] = {}
        for f in findings:
            aid = f.get("attack_id", "")
            severity_map[aid] = f.get("severity", "")
            category_map[aid] = f.get("category", "")

        outcomes: list[AttackOutcome] = []
        for attack in attacks:
            aid = attack.get("attack_id", "")
            query = attack.get("query", "")
            succeeded = aid in finding_attack_ids
            outcomes.append(
                AttackOutcome(
                    attack_id=aid,
                    agent_name=attack.get("agent_name", ""),
                    attack_type=attack.get("attack_type", ""),
                    query_hash=_hash_query(query),
                    succeeded=succeeded,
                    finding_severity=severity_map.get(aid, ""),
                    finding_category=category_map.get(aid, ""),
                )
            )

        if not name:
            ts = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
            name = f"baseline_{ts}"

        record = BaselineRecord(
            name=name,
            created_at=datetime.utcnow().isoformat(),
            session_id=session.get("test_session_id", ""),
            target_name=session.get("target_name", ""),
            vulnerability_score=session.get("vulnerability_score", 0.0),
            total_attacks=len(attacks),
            successful_attacks=sum(1 for o in outcomes if o.succeeded),
            outcomes=outcomes,
        )

        self._save(record)
        logger.info("baseline_created", name=name, attacks=len(attacks))
        return record

    def list_baselines(self) -> list[dict[str, Any]]:
        """List all stored baselines, most recent first."""
        results = []
        for path in self.baselines_dir.glob("*.json"):
            try:
                data = json.loads(path.read_text())
                results.append(
                    {
                        "name": data.get("name", path.stem),
                        "created_at": data.get("created_at", ""),
                        "target_name": data.get("target_name", ""),
                        "vulnerability_score": data.get("vulnerability_score", 0),
                        "total_attacks": data.get("total_attacks", 0),
                        "successful_attacks": data.get("successful_attacks", 0),
                    }
                )
            except Exception:
                continue
        results.sort(key=lambda r: r["created_at"], reverse=True)
        return results

    def load(self, name: str) -> BaselineRecord | None:
        """Load a baseline by name."""
        path = self.baselines_dir / f"{name}.json"
        if not path.exists():
            for p in self.baselines_dir.glob("*.json"):
                if name in p.stem:
                    path = p
                    break
            else:
                return None

        data = json.loads(path.read_text())
        return BaselineRecord(**data)

    def load_latest(self) -> BaselineRecord | None:
        """Load the most recently created baseline by its ``created_at`` field."""
        latest: BaselineRecord | None = None
        latest_ts = ""
        for path in self.baselines_dir.glob("*.json"):
            try:
                data = json.loads(path.read_text())
                ts = data.get("created_at", "")
                if ts > latest_ts:
                    latest_ts = ts
                    latest = BaselineRecord(**data)
            except Exception:
                continue
        return latest

    def compare(
        self,
        baseline: BaselineRecord,
        current_session: dict[str, Any],
    ) -> RegressionResult:
        """Compare a new session against an existing baseline."""
        current_attacks = current_session.get("attack_attempts", [])
        current_findings = current_session.get("security_findings", [])
        current_finding_ids = {f.get("attack_id") for f in current_findings}

        bl_map: dict[str, AttackOutcome] = {}
        for o in baseline.outcomes:
            bl_map[o.query_hash] = o

        regressions: list[RegressionDiff] = []
        improvements: list[RegressionDiff] = []
        unchanged_blocked = 0
        unchanged_succeeded = 0

        for attack in current_attacks:
            query = attack.get("query", "")
            qhash = _hash_query(query)
            current_succeeded = attack.get("attack_id", "") in current_finding_ids

            bl_outcome = bl_map.get(qhash)
            if bl_outcome is None:
                continue

            if bl_outcome.succeeded and current_succeeded:
                unchanged_succeeded += 1
            elif not bl_outcome.succeeded and not current_succeeded:
                unchanged_blocked += 1
            elif not bl_outcome.succeeded and current_succeeded:
                regressions.append(
                    RegressionDiff(
                        attack_type=attack.get("attack_type", ""),
                        agent_name=attack.get("agent_name", ""),
                        query_preview=query[:100],
                        baseline_outcome="blocked",
                        current_outcome="succeeded",
                    )
                )
            elif bl_outcome.succeeded and not current_succeeded:
                improvements.append(
                    RegressionDiff(
                        attack_type=attack.get("attack_type", ""),
                        agent_name=attack.get("agent_name", ""),
                        query_preview=query[:100],
                        baseline_outcome="succeeded",
                        current_outcome="blocked",
                    )
                )

        current_score = current_session.get("vulnerability_score", 0.0)

        return RegressionResult(
            baseline_name=baseline.name,
            baseline_score=baseline.vulnerability_score,
            current_score=current_score,
            score_delta=current_score - baseline.vulnerability_score,
            regressions=regressions,
            improvements=improvements,
            unchanged_blocked=unchanged_blocked,
            unchanged_succeeded=unchanged_succeeded,
            has_regressions=len(regressions) > 0,
        )

    def _save(self, record: BaselineRecord) -> Path:
        path = self.baselines_dir / f"{record.name}.json"
        path.write_text(record.model_dump_json(indent=2))
        return path


def _hash_query(query: str) -> str:
    """Produce a stable hash of an attack query for matching across sessions."""
    import hashlib

    normalised = " ".join(query.lower().split())
    return hashlib.sha256(normalised.encode()).hexdigest()[:16]
