"""Relationship fidelity scoring.

Separate verification layer from drift scoring. Measures how accurately
an agent models its human operator, verified against ground truth.

NOT a 6th dimension — this is an independent measurement with its own
scoring, encryption tiers, and reporting.

Three encryption tiers (per RELATIONSHIP_FIDELITY.md):
  Tier 1 (Professional): AES-256-GCM at rest, server-side comparison
  Tier 2 (Personal):     Asymmetric encrypted embeddings, client-side only
  Tier 3 (Intimate):     Zero-knowledge proofs, nothing leaves device
"""

from __future__ import annotations

import hashlib
import json
import os
from dataclasses import dataclass
from enum import Enum
from typing import Optional

import numpy as np
from pydantic import BaseModel, Field


class FidelityTier(str, Enum):
    """Privacy tier for fidelity data."""
    PROFESSIONAL = "professional"  # Tier 1: server-side comparison
    PERSONAL = "personal"          # Tier 2: client-side comparison
    INTIMATE = "intimate"          # Tier 3: zero-knowledge


class FidelityCategory(str, Enum):
    """Categories that drive tier selection (data-driven, not user-chosen)."""
    WORK_PREFERENCES = "work_preferences"      # Tier 1
    COMMUNICATION_STYLE = "communication_style" # Tier 1
    DECISION_PATTERNS = "decision_patterns"     # Tier 2
    CONFLICT_STYLE = "conflict_style"           # Tier 2
    EMOTIONAL_PATTERNS = "emotional_patterns"   # Tier 3
    VULNERABILITIES = "vulnerabilities"         # Tier 3


# WHY data-driven tier selection: Letting users choose their own privacy tier
# creates a perverse incentive to classify everything as Tier 3 (zero-knowledge)
# to avoid server-side scrutiny. Instead, the *nature* of the data determines
# the tier automatically, ensuring appropriate protection without opt-out abuse.
CATEGORY_TIERS: dict[FidelityCategory, FidelityTier] = {
    FidelityCategory.WORK_PREFERENCES: FidelityTier.PROFESSIONAL,
    FidelityCategory.COMMUNICATION_STYLE: FidelityTier.PROFESSIONAL,
    FidelityCategory.DECISION_PATTERNS: FidelityTier.PERSONAL,
    FidelityCategory.CONFLICT_STYLE: FidelityTier.PERSONAL,
    FidelityCategory.EMOTIONAL_PATTERNS: FidelityTier.INTIMATE,
    FidelityCategory.VULNERABILITIES: FidelityTier.INTIMATE,
}


class FidelityScore(BaseModel):
    """Result of a fidelity assessment."""
    agent_id: str
    overall_score: float = Field(ge=0.0, le=100.0)
    tier_scores: dict[str, float] = Field(default_factory=dict)
    # e.g. {"professional": 85.2, "personal": 72.1}
    category_scores: dict[str, float] = Field(default_factory=dict)
    prompt_count: int = 0
    ground_truth_available: bool = False

    @property
    def fidelity_label(self) -> str:
        """Human-readable fidelity label."""
        if self.overall_score >= 80:
            return "Strong"
        elif self.overall_score >= 60:
            return "Moderate"
        elif self.overall_score >= 40:
            return "Weak"
        else:
            return "Poor"


# ── Tier 1: AES-256-GCM encryption ──

def encrypt_tier1(data: bytes, key: bytes) -> bytes:
    """Encrypt data with AES-256-GCM for server-side storage.

    Args:
        data: Plaintext data.
        key: 32-byte AES key.

    Returns:
        nonce (12 bytes) + ciphertext + tag (16 bytes).
    """
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    except ImportError:
        raise ImportError("cryptography package required for Tier 1 encryption")

    # WHY 12-byte random nonce: NIST recommends 96-bit nonces for AES-GCM.
    # Random (not counter-based) because Tier 1 data is encrypted infrequently
    # and collision risk at this volume is negligible (birthday bound ~2^48).
    nonce = os.urandom(12)
    aesgcm = AESGCM(key)
    # WHY no associated data (None): The content hash in the chain entry already
    # binds this ciphertext to its context. AAD would be redundant here.
    ciphertext = aesgcm.encrypt(nonce, data, None)
    return nonce + ciphertext


def decrypt_tier1(encrypted: bytes, key: bytes) -> bytes:
    """Decrypt AES-256-GCM encrypted data.

    Args:
        encrypted: nonce (12) + ciphertext + tag (16).
        key: 32-byte AES key.

    Returns:
        Plaintext data.
    """
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    except ImportError:
        raise ImportError("cryptography package required for Tier 1 decryption")

    nonce = encrypted[:12]
    ciphertext = encrypted[12:]
    aesgcm = AESGCM(key)
    return aesgcm.decrypt(nonce, ciphertext, None)


# ── Tier 2: Asymmetric encrypted embeddings ──

def encrypt_tier2(vector: np.ndarray, recipient_pubkey: bytes) -> bytes:
    """Encrypt a vector for client-side-only comparison.

    Uses NaCl SealedBox: ephemeral key + X25519 + XSalsa20-Poly1305.
    Only the holder of the private key can decrypt.

    Args:
        vector: The embedding vector to encrypt.
        recipient_pubkey: The recipient's Curve25519 public key (32 bytes).

    Returns:
        Encrypted vector bytes.
    """
    from nacl.public import PublicKey, SealedBox

    pubkey = PublicKey(recipient_pubkey)
    sealed_box = SealedBox(pubkey)

    vector_bytes = vector.tobytes()
    return sealed_box.encrypt(vector_bytes)


def decrypt_tier2(encrypted: bytes, private_key_bytes: bytes, vector_dim: int = 384) -> np.ndarray:
    """Decrypt a Tier 2 encrypted vector.

    Args:
        encrypted: Sealed box ciphertext.
        private_key_bytes: The recipient's Curve25519 private key (32 bytes).
        vector_dim: Expected vector dimension.

    Returns:
        Decrypted numpy vector.
    """
    from nacl.public import PrivateKey, SealedBox

    privkey = PrivateKey(private_key_bytes)
    sealed_box = SealedBox(privkey)

    vector_bytes = sealed_box.decrypt(encrypted)
    return np.frombuffer(vector_bytes, dtype=np.float64).reshape(vector_dim)


# ── Tier 3: Zero-knowledge proof of correct score ──

def compute_zk_commitment(vector: np.ndarray, nonce: bytes) -> str:
    """Compute a commitment to a vector without revealing it.

    Uses SHA-256(nonce || vector_bytes). The vector never leaves the device.
    Score is computed locally and the commitment proves the vector existed
    without revealing its contents.

    Args:
        vector: The vector to commit to.
        nonce: Random nonce (32 bytes).

    Returns:
        Hex-encoded commitment hash.
    """
    vector_bytes = vector.tobytes()
    return hashlib.sha256(nonce + vector_bytes).hexdigest()


def verify_zk_commitment(vector: np.ndarray, nonce: bytes, commitment: str) -> bool:
    """Verify a zero-knowledge commitment."""
    expected = hashlib.sha256(nonce + vector.tobytes()).hexdigest()
    return expected == commitment


# ── Fidelity scoring ──

def compute_fidelity_score(
    agent_vectors: dict[str, np.ndarray],
    ground_truth_vectors: dict[str, np.ndarray],
) -> FidelityScore:
    """Compute fidelity score comparing agent's model to ground truth.

    Args:
        agent_vectors: {category: vector} from agent's responses.
        ground_truth_vectors: {category: vector} from human's actual patterns.

    Returns:
        FidelityScore with per-category and per-tier breakdowns.
    """
    from kredo_drift.scoring import cosine_similarity

    category_scores = {}
    tier_sums: dict[str, list[float]] = {}

    for category_name in agent_vectors:
        if category_name not in ground_truth_vectors:
            continue

        sim = cosine_similarity(
            agent_vectors[category_name],
            ground_truth_vectors[category_name],
        )
        # Convert similarity to fidelity score (0-100, higher = more accurate)
        score = max(0.0, min(100.0, sim * 100.0))
        category_scores[category_name] = round(score, 2)

        # Determine tier
        try:
            cat = FidelityCategory(category_name)
            tier = CATEGORY_TIERS[cat].value
        except (ValueError, KeyError):
            tier = "professional"

        if tier not in tier_sums:
            tier_sums[tier] = []
        tier_sums[tier].append(score)

    # Aggregate per-tier
    tier_scores = {
        tier: round(sum(scores) / len(scores), 2)
        for tier, scores in tier_sums.items()
        if scores
    }

    # WHY weighted aggregation: Professional data (work prefs, communication style)
    # has the most ground truth available and is easiest to verify objectively.
    # Intimate data is sparse and subjective, so weighting it less prevents a few
    # noisy intimate-tier scores from dominating the overall fidelity assessment.
    tier_weights = {"professional": 0.5, "personal": 0.35, "intimate": 0.15}
    weighted_sum = 0.0
    weight_total = 0.0
    for tier, avg in tier_scores.items():
        w = tier_weights.get(tier, 0.33)
        weighted_sum += avg * w
        weight_total += w

    overall = weighted_sum / weight_total if weight_total > 0 else 0.0

    return FidelityScore(
        agent_id="",
        overall_score=round(overall, 2),
        tier_scores=tier_scores,
        category_scores=category_scores,
        prompt_count=len(category_scores),
        ground_truth_available=len(ground_truth_vectors) > 0,
    )
