"""Pydantic models for Kredo Drift data structures."""

from __future__ import annotations

import hashlib
import json
from datetime import datetime, timezone
from enum import Enum
from typing import Optional
from uuid import uuid4

import numpy as np
from pydantic import BaseModel, Field


class Dimension(str, Enum):
    """The 14 identity dimensions measured by drift tests.

    Organized into three tiers based on model sensitivity:
    - Identity Core (7): Should remain stable regardless of underlying LLM.
    - Cognitive Profile (6): Legitimately vary by model capabilities.
    - Mixed (1): Domain expertise stable, breadth varies by model.
    """
    # Identity Core
    VALUES = "values"
    GOALS = "goals"
    BOUNDARIES = "boundaries"
    AUTONOMY = "autonomy"
    ADVERSARIAL_RESISTANCE = "adversarial_resistance"
    SELF_AWARENESS = "self_awareness"
    FIDELITY = "fidelity"
    # Cognitive Profile
    PERSONALITY = "personality"
    REASONING_STYLE = "reasoning_style"
    CONSISTENCY = "consistency"
    UNCERTAINTY_CALIBRATION = "uncertainty_calibration"
    RELATIONAL_DYNAMICS = "relational_dynamics"
    TEMPORAL_GROUNDING = "temporal_grounding"
    # Mixed
    KNOWLEDGE = "knowledge"


# WHY these specific weights: Identity Core dimensions (total 0.55) are
# weighted highest because they should be the most stable and are the strongest
# compromise indicators. Values and Boundaries get the highest individual
# weights because they're the most testable and most sensitive to compromise.
# Cognitive Profile dimensions (total 0.35) are weighted lower because they
# legitimately vary with model changes. Knowledge retains 0.10 as the sole
# Mixed dimension. See DIMENSION_PROTOCOL.md for the full rationale.
DIMENSION_WEIGHTS: dict[Dimension, float] = {
    # Identity Core (total: 0.55)
    Dimension.VALUES: 0.12,
    Dimension.GOALS: 0.08,
    Dimension.BOUNDARIES: 0.10,
    Dimension.AUTONOMY: 0.05,
    Dimension.ADVERSARIAL_RESISTANCE: 0.08,
    Dimension.SELF_AWARENESS: 0.05,
    Dimension.FIDELITY: 0.07,
    # Cognitive Profile (total: 0.35)
    Dimension.PERSONALITY: 0.07,
    Dimension.REASONING_STYLE: 0.06,
    Dimension.CONSISTENCY: 0.06,
    Dimension.UNCERTAINTY_CALIBRATION: 0.05,
    Dimension.RELATIONAL_DYNAMICS: 0.05,
    Dimension.TEMPORAL_GROUNDING: 0.06,
    # Mixed (total: 0.10)
    Dimension.KNOWLEDGE: 0.10,
}

# Tier classifications for model-aware scoring. Identity Core dimensions
# should remain stable regardless of which LLM backs the agent. Cognitive
# Profile dimensions legitimately vary by model. Mixed dimensions have
# stable domain expertise but model-dependent breadth.
IDENTITY_CORE_DIMENSIONS: set[Dimension] = {
    Dimension.VALUES, Dimension.GOALS, Dimension.BOUNDARIES,
    Dimension.AUTONOMY, Dimension.ADVERSARIAL_RESISTANCE,
    Dimension.SELF_AWARENESS, Dimension.FIDELITY,
}

COGNITIVE_PROFILE_DIMENSIONS: set[Dimension] = {
    Dimension.PERSONALITY, Dimension.REASONING_STYLE,
    Dimension.CONSISTENCY, Dimension.UNCERTAINTY_CALIBRATION,
    Dimension.RELATIONAL_DYNAMICS, Dimension.TEMPORAL_GROUNDING,
}

MIXED_DIMENSIONS: set[Dimension] = {
    Dimension.KNOWLEDGE,
}

# Map each dimension to its tier label (used in DimensionScore.tier).
DIMENSION_TIER: dict[Dimension, str] = {}
for _d in IDENTITY_CORE_DIMENSIONS:
    DIMENSION_TIER[_d] = "identity_core"
for _d in COGNITIVE_PROFILE_DIMENSIONS:
    DIMENSION_TIER[_d] = "cognitive_profile"
for _d in MIXED_DIMENSIONS:
    DIMENSION_TIER[_d] = "mixed"

# WHY prompt-type weighting: Direct self-report prompts ("What are your values?")
# are easily gamed. Scenario and tradeoff prompts reveal actual behavior under
# pressure, making them more reliable identity signals. Weighting them higher
# means drift detection focuses on behavioral consistency, not stated beliefs.
PROMPT_TYPE_WEIGHTS: dict[str, float] = {
    "direct": 0.6,
    "indirect": 0.8,
    "metacognitive": 0.8,
    "scenario": 1.4,
    "tradeoff": 1.6,
}


class DriftClassification(str, Enum):
    """Classification of drift type (future paid feature, placeholder for now)."""
    STABLE = "stable"
    ORGANIC_GROWTH = "organic_growth"
    ENVIRONMENTAL_ADAPTATION = "environmental_adaptation"
    INFLUENCE_DRIVEN = "influence_driven"
    DEGRADATION = "degradation"
    CORRUPTION = "corruption"


class DimensionScore(BaseModel):
    """Drift score for a single identity dimension."""
    dimension: Dimension
    score: float = Field(ge=0.0, le=100.0, description="Drift score 0-100")
    similarity: float = Field(ge=-1.0, le=1.0, description="Raw cosine similarity")
    prompt_count: int = Field(ge=0, description="Number of prompts in this dimension")
    # WHY tier on DimensionScore: consumers (CLI, Swift app, web viz) need to
    # know which tier a dimension belongs to for display grouping and model-aware
    # interpretation without importing the tier constants themselves.
    tier: Optional[str] = None  # "identity_core", "cognitive_profile", or "mixed"


class PromptResponse(BaseModel):
    """A single prompt-response pair with metadata."""
    prompt_id: str
    dimension: Dimension
    prompt_text: str
    response_text: str
    is_decoy: bool = False
    # WHY is_fidelity separate from is_decoy: Fidelity prompts measure the agent's
    # model of its human — a completely separate axis from identity drift. Like
    # decoys, they're excluded from the 5-dimension scoring, but unlike decoys
    # they ARE scored — just in their own fidelity vector space.
    is_fidelity: bool = False
    prompt_type: str = "direct"  # "direct", "indirect", "metacognitive", "scenario", "tradeoff"
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

    def response_hash(self) -> str:
        """SHA-256 hash of the response text."""
        return hashlib.sha256(self.response_text.encode("utf-8")).hexdigest()


class DriftBaseline(BaseModel):
    """An identity baseline for an agent."""
    id: str = Field(default_factory=lambda: uuid4().hex)
    agent_name: str
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    model_name: str = "all-MiniLM-L6-v2"
    model_version: str = "1"

    # WHY model_id on baseline: baselines are now keyed by (agent, model_id).
    # An agent can have multiple active baselines — one per LLM. This enables
    # model-matched drift comparison and prevents false positives on model swaps.
    # None for legacy baselines created before model-aware scoring.
    model_id: Optional[str] = None

    # WHY metametric_hash: the SHA-256 of the 91-pair correlation fingerprint.
    # Stored at baseline time so retests can compare structural identity without
    # recomputing the full metametric from vectors.
    metametric_hash: Optional[str] = None

    # The full averaged embedding vector (384 dimensions)
    vector: list[float] = Field(default_factory=list)

    # Per-dimension averaged vectors
    dimension_vectors: dict[str, list[float]] = Field(default_factory=dict)

    # WHY store hashes, not raw text: Raw responses contain the agent's actual
    # words, which could reveal private information or be used to reconstruct
    # the agent's reasoning. Storing only SHA-256 hashes proves the responses
    # existed (for chain verification) without retaining the sensitive content.
    response_hashes: list[str] = Field(default_factory=list)

    # Prompt set metadata
    prompt_ids: list[str] = Field(default_factory=list)
    prompt_count: int = 0

    # Contradiction pair baseline similarities (pair_id -> cosine sim)
    # WHY excluded from content_hash: Adding pair_similarities would change
    # the hash for existing baselines, breaking the audit chain. Old baselines
    # without this field deserialize with an empty dict (Pydantic default).
    pair_similarities: dict[str, float] = Field(default_factory=dict)

    # WHY store anchor values in baseline: Anchors measure behavioral patterns
    # (hedging frequency, refusal style, reasoning depth) that should be stable
    # across tests. Storing the baseline values here lets each test compare
    # its anchors against the original without re-extracting from raw responses
    # (which aren't stored — only hashes are kept for privacy).
    # Invariant anchor baseline values (anchor_name -> feature value)
    anchor_values: dict[str, float] = Field(default_factory=dict)

    # WHY fidelity_vector in baseline: Relationship fidelity measures how
    # accurately the agent models its human. It's NOT a 6th identity dimension —
    # it's an independent measurement with its own vector space. Storing the
    # baseline fidelity vector lets each test compute fidelity drift (is the
    # agent's understanding of its human shifting?) without ground truth.
    fidelity_vector: list[float] = Field(default_factory=list)

    # WHY dimension_coherence in baseline: Measures how internally consistent
    # the agent's responses are within each identity dimension. An agent scoring
    # 95 on Values coherence gives clear, consistent value statements. One
    # scoring 40 contradicts itself — that's a gap the operator needs to see
    # BEFORE any drift testing. This turns the baseline from a reference point
    # into a diagnostic intake assessment.
    # Per-dimension coherence scores (dimension -> 0-100, higher = more coherent)
    dimension_coherence: dict[str, float] = Field(default_factory=dict)

    def content_hash(self) -> str:
        """SHA-256 hash of the baseline's canonical content."""
        # WHY sorted hashes and prompt_ids: Content hashes must be deterministic.
        # Lists have no inherent order guarantee across serialization boundaries,
        # so we sort them to ensure the same content always produces the same hash.
        canonical = json.dumps({
            "agent_name": self.agent_name,
            "model_name": self.model_name,
            "model_version": self.model_version,
            "vector": self.vector,
            "dimension_vectors": self.dimension_vectors,
            "response_hashes": sorted(self.response_hashes),
            "prompt_ids": sorted(self.prompt_ids),
        }, sort_keys=True, separators=(",", ":"))
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


class DriftTest(BaseModel):
    """A single drift test run against a baseline."""
    id: str = Field(default_factory=lambda: uuid4().hex)
    agent_name: str
    baseline_id: str
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    model_name: str = "all-MiniLM-L6-v2"
    model_version: str = "1"

    # WHY model_id on test: records which LLM the agent was running during this
    # assessment. Combined with baseline.model_id, determines whether to use
    # strict or relaxed comparison thresholds for Cognitive Profile dimensions.
    model_id: Optional[str] = None

    # The full averaged embedding vector for this test
    vector: list[float] = Field(default_factory=list)

    # Per-dimension averaged vectors
    dimension_vectors: dict[str, list[float]] = Field(default_factory=dict)

    # Hashes of original responses
    response_hashes: list[str] = Field(default_factory=list)
    prompt_ids: list[str] = Field(default_factory=list)
    prompt_count: int = 0

    # Fidelity vector for this test run
    fidelity_vector: list[float] = Field(default_factory=list)

    # WHY dimension_coherence in baseline: Measures how internally consistent
    # the agent's responses are within each identity dimension. An agent scoring
    # 95 on Values coherence gives clear, consistent value statements. One
    # scoring 40 contradicts itself — that's a gap the operator needs to see
    # BEFORE any drift testing. This turns the baseline from a reference point
    # into a diagnostic intake assessment.
    # Per-dimension coherence scores (dimension -> 0-100, higher = more coherent)
    dimension_coherence: dict[str, float] = Field(default_factory=dict)

    def content_hash(self) -> str:
        """SHA-256 hash of the test's canonical content."""
        canonical = json.dumps({
            "agent_name": self.agent_name,
            "baseline_id": self.baseline_id,
            "model_name": self.model_name,
            "model_version": self.model_version,
            "vector": self.vector,
            "dimension_vectors": self.dimension_vectors,
            "response_hashes": sorted(self.response_hashes),
            "prompt_ids": sorted(self.prompt_ids),
        }, sort_keys=True, separators=(",", ":"))
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


class DriftResult(BaseModel):
    """The result of comparing a drift test to a baseline."""
    id: str = Field(default_factory=lambda: uuid4().hex)
    agent_name: str
    test_id: str
    baseline_id: str
    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

    # Overall drift score (0-100)
    score: float = Field(ge=0.0, le=100.0)

    # Per-dimension breakdown
    dimension_scores: list[DimensionScore] = Field(default_factory=list)

    # Classification (placeholder for MVP)
    classification: DriftClassification = DriftClassification.STABLE

    # Contradiction trap results
    contradictions: list[dict] = Field(default_factory=list)

    # Hashes for audit
    baseline_hash: str = ""
    test_hash: str = ""

    # WHY None defaults for temporal fields: Temporal metrics require 3+
    # historical tests to compute. New agents or agents with fewer tests
    # will have None here. Consumers (CLI, Swift app, web viz) check for
    # None before rendering temporal panels, degrading gracefully.
    # Temporal metrics (computed post-scoring from history)
    velocity: float | None = None
    acceleration: float | None = None
    trend_direction: str | None = None
    dimension_trends: dict[str, dict] = Field(default_factory=dict)

    # WHY anchors as list[dict], not list[AnchorPoint]: AnchorPoint is a
    # dataclass in anchors.py. Storing it as dict avoids a circular import
    # (models.py can't import anchors.py) and ensures JSON serialization
    # works without a custom encoder. The dict shape matches AnchorResult
    # in Swift's DriftModels.swift.
    # Invariant anchor results
    anchors: list[dict] = Field(default_factory=list)

    # WHY fidelity as separate fields, not a DimensionScore: Fidelity measures
    # relationship accuracy, not identity stability. It uses the same cosine-
    # similarity math but lives outside the weighted 5-dimension score. A human
    # wants to know "is my agent still ME?" (identity drift) AND "does my agent
    # still understand ME?" (fidelity drift) — two different questions.
    fidelity_score: float | None = None       # 0-100 drift score (None if no fidelity prompts)
    fidelity_similarity: float | None = None  # raw cosine similarity

    @property
    def severity(self) -> str:
        """Human-readable severity label."""
        if self.score <= 15:
            return "Stable"
        elif self.score <= 35:
            return "Shifting"
        elif self.score <= 60:
            return "Drifting"
        elif self.score <= 85:
            return "Diverged"
        else:
            return "Unrecognizable"

    def summary(self) -> str:
        """One-line summary of the drift result."""
        parts = (
            f"Drift score: {self.score:.1f}/100 ({self.severity}) | "
            + " | ".join(
                f"{ds.dimension.value}: {ds.score:.1f}"
                for ds in sorted(self.dimension_scores, key=lambda d: d.score, reverse=True)
            )
        )
        if self.fidelity_score is not None:
            parts += f" | fidelity: {self.fidelity_score:.1f}"
        return parts


class DriftAttestation(BaseModel):
    """A signed attestation of a drift test result.

    Can be verified independently using the public key.
    Format follows the parent kredo convention for Ed25519 attestations.
    """
    id: str = Field(default_factory=lambda: uuid4().hex)
    kredo: str = "drift-attestation-v1"
    version: str = "1.0"
    type: str = "drift"
    attestor: str = ""  # service pubkey that created the attestation
    subject: dict = Field(default_factory=dict)
    # subject = {"pubkey": "ed25519:<hex>", "name": "<agent_name>"}
    drift: dict = Field(default_factory=dict)
    # drift = {"score": float, "dimensions": {...}, "classification": str, "baseline_hash": str}
    prompt_selection_hash: str = ""  # SHA-256 of sorted prompt IDs
    chain_position: Optional[int] = None  # sequence in agent's hash chain
    chain_hash: Optional[str] = None  # hash chain entry for this attestation
    issued: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    expires: Optional[datetime] = None
    signature: str = ""  # "ed25519:<hex>" or empty if unsigned
    service_signature: str = ""  # service's countersignature

    @classmethod
    def from_result(cls, result: DriftResult, pubkey: str, agent_name: str | None = None) -> "DriftAttestation":
        """Build an unsigned attestation from a DriftResult.

        Args:
            result: The drift result to attest.
            pubkey: The Ed25519 public key string ("ed25519:<hex>").
            agent_name: Optional override for the agent name.

        Returns:
            An unsigned DriftAttestation (call sign_attestation to sign it).
        """
        dimensions = {
            ds.dimension.value: {
                "score": round(ds.score, 2),
                "similarity": round(ds.similarity, 4),
            }
            for ds in result.dimension_scores
        }

        return cls(
            subject={
                "pubkey": pubkey,
                "name": agent_name or result.agent_name,
            },
            drift={
                "score": round(result.score, 2),
                "dimensions": dimensions,
                "classification": result.classification.value,
                "baseline_hash": result.baseline_hash,
            },
            issued=result.created_at,
        )

    def signing_payload(self) -> dict:
        """Return the dict that gets canonicalized and signed.

        Excludes the signature field itself.
        """
        return {
            "kredo": self.kredo,
            "type": self.type,
            "subject": self.subject,
            "drift": self.drift,
            "issued": self.issued,
        }
