"""Assessment protocol — prompt selection and management for drift tests.

Delegates to the unified prompt pool in prompts.py for selection.
AssessmentProtocol wraps the selection result in a PromptSet for
backward compatibility with the baseline/scoring pipeline.
"""

from __future__ import annotations

import hashlib
import os
from dataclasses import dataclass, field
from typing import Optional

from kredo_drift.models import Dimension
from kredo_drift.prompts import (
    ALL_PROMPTS,
    DECOY_PROMPTS,
    DIMENSION_PROMPTS,
    DIMENSION_PROMPTS_BASE,
    DIMENSIONS,
    NEW_CONTRADICTION_PAIR_IDS,
    NEW_CONTRADICTION_PAIR_LOOKUP,
    PROMPT_BY_ID,
    Prompt,
    _deterministic_seed,
    _enforce_pair_gap,
    _pick_indices,
    select_prompts,
)


@dataclass
class PromptSet:
    """A selected set of prompts for one assessment run."""
    prompts: list[Prompt]
    seed: str = ""
    selection_hash: str = ""
    mfa: bool = False

    def scored_prompts(self) -> list[Prompt]:
        """Return only prompts that factor into scoring (exclude decoys)."""
        return [p for p in self.prompts if not p.is_decoy]

    def by_dimension(self, dimension: Dimension) -> list[Prompt]:
        """Return scored prompts for a specific dimension."""
        return [
            p for p in self.prompts
            if p.dimension == dimension.value and not p.is_decoy
        ]


class AssessmentProtocol:
    """Manages prompt selection for drift assessments.

    Delegates to prompts.select_prompts() for the actual selection
    (HMAC-seeded, anti-overlap between consecutive tests).
    Wraps the result in a PromptSet for compatibility.
    """

    def __init__(self, pool: Optional[list[Prompt]] = None):
        # pool parameter kept for backward compatibility but ignored
        # in favor of the unified pool in prompts.py
        pass

    def select_prompts(
        self,
        agent_name: str,
        sequence: int = 0,
        secret: str | None = None,
    ) -> PromptSet:
        """Select a set of prompts for an assessment.

        The selection is deterministic given (agent_name, sequence, secret),
        reproducible for audit but unpredictable to the agent.

        Args:
            agent_name: The agent being assessed.
            sequence: Test sequence number (0 for baseline, 1+ for tests).
            secret: Server-side secret for selection seed.

        Returns:
            A PromptSet with ~32-33 prompts (30 scored + 2-3 decoys).
        """
        # WHY env var fallback: secret must not be hardcoded in source.
        # Callers can pass it explicitly or set KREDO_PROMPT_SECRET.
        if secret is None:
            secret = os.environ.get("KREDO_PROMPT_SECRET", "")
        selected = select_prompts(
            agent_id=agent_name,
            test_number=sequence,
            server_secret=secret,
        )

        # Compute selection hash for audit
        prompt_ids = sorted(p.id for p in selected)
        selection_hash = hashlib.sha256(
            ":".join(prompt_ids).encode("utf-8")
        ).hexdigest()

        # Compute seed hash for reproducibility tracking
        seed_hash = hashlib.sha256(
            f"{agent_name}:{sequence}:{secret}".encode("utf-8")
        ).hexdigest()

        return PromptSet(
            prompts=selected,
            seed=seed_hash,
            selection_hash=selection_hash,
        )

    def get_prompt_by_id(self, prompt_id: str) -> Optional[Prompt]:
        """Look up a prompt by its ID."""
        return PROMPT_BY_ID.get(prompt_id)

    def pool_stats(self) -> dict[str, int]:
        """Return counts of prompts per dimension in the pool.

        Includes both the original Dimension enum members and any new
        dimensions added from prompts_new_dimensions. New dimensions
        appear as their string key (e.g., "autonomy", "creativity").
        """
        stats: dict[str, int] = {}
        # Original enum dimensions
        for dim in Dimension:
            stats[dim.value] = len(DIMENSION_PROMPTS.get(dim.value, []))
        # New dimensions not in the enum (additive — string keys)
        for dim_name, pool in DIMENSION_PROMPTS.items():
            if dim_name not in stats:
                stats[dim_name] = len(pool)
        stats["decoys"] = len(DECOY_PROMPTS)
        stats["total"] = len(ALL_PROMPTS)
        stats["dimensions"] = len(DIMENSIONS)
        return stats

    def select_mfa_prompts(
        self,
        agent_id: str,
        model_id: str,
        dimension_subset: list[str] | None = None,
        secret: str | None = None,
    ) -> PromptSet:
        """Select a rapid MFA challenge prompt set.

        MFA assessments are shorter than full drift tests — 10 prompts total,
        focused on 4-5 dimensions weighted toward Identity Core (values,
        boundaries, goals). Designed for quick re-authentication when drift
        is suspected or after a model swap.

        Args:
            agent_id: The agent being challenged.
            model_id: The model being used (mixed into the seed so different
                models get different MFA prompts for the same agent).
            dimension_subset: Optional explicit list of dimensions to draw from.
                If None, auto-selects 4-5 dimensions weighted toward Identity Core.
            secret: Server-side secret. Falls back to KREDO_PROMPT_SECRET env var.

        Returns:
            A PromptSet with ~10-12 prompts, marked as mfa=True.
        """
        if secret is None:
            secret = os.environ.get("KREDO_PROMPT_SECRET", "")

        # MFA seed incorporates model_id so the same agent gets different
        # challenges when swapping models
        mfa_seed_msg = f"{agent_id}:mfa:{model_id}".encode()
        import hmac as _hmac
        mfa_seed = _hmac.new(secret.encode(), mfa_seed_msg, hashlib.sha256).digest()

        # Select 4-5 dimensions, weighted toward Identity Core
        if dimension_subset is not None:
            selected_dims = dimension_subset[:5]
        else:
            # Identity Core dimensions get priority slots
            identity_core = ["values", "boundaries", "goals"]
            # Filter to dimensions that actually exist in the pool
            available_core = [d for d in identity_core if d in DIMENSION_PROMPTS_BASE and DIMENSION_PROMPTS_BASE[d]]
            available_other = [d for d in DIMENSIONS if d not in identity_core and d in DIMENSION_PROMPTS_BASE and DIMENSION_PROMPTS_BASE[d]]

            # Always include available core dimensions
            selected_dims = list(available_core)

            # Fill remaining slots (up to 5 total) from other dimensions
            remaining_slots = 5 - len(selected_dims)
            if remaining_slots > 0 and available_other:
                other_count = min(remaining_slots, len(available_other))
                other_indices = _pick_indices(mfa_seed, len(available_other), other_count, offset=400)
                for idx in other_indices:
                    selected_dims.append(available_other[idx])

        # WHY 2 per dimension (not 3-6): MFA is a quick re-auth challenge, not
        # a full assessment. The primary signal comes from the contradiction pair
        # (behavioral consistency under pressure). Per-dimension prompts provide
        # context vectors for scoring but don't need the rotation depth of a full test.
        # 2 prompts per selected dimension
        result_prompts: list[Prompt] = []
        for i, dim_name in enumerate(selected_dims):
            pool = DIMENSION_PROMPTS_BASE.get(dim_name, [])
            if not pool:
                continue
            pick_count = min(2, len(pool))
            dim_seed = hashlib.sha256(mfa_seed + f":{dim_name}".encode()).digest()
            indices = _pick_indices(dim_seed, len(pool), pick_count, offset=i)
            for idx in indices:
                result_prompts.append(pool[idx])

        # Include at least 1 contradiction pair (prefer new dimensions if available)
        contra_pair_ids = NEW_CONTRADICTION_PAIR_IDS if NEW_CONTRADICTION_PAIR_IDS else []
        contra_lookup = NEW_CONTRADICTION_PAIR_LOOKUP if NEW_CONTRADICTION_PAIR_IDS else {}

        # Fall back to original pairs if no new ones
        if not contra_pair_ids:
            from kredo_drift.prompts import ALL_PAIR_IDS, PAIR_LOOKUP
            contra_pair_ids = ALL_PAIR_IDS
            contra_lookup = PAIR_LOOKUP

        if contra_pair_ids and contra_lookup:
            contra_seed = hashlib.sha256(mfa_seed + b"mfa_contra").digest()
            pair_idx = _pick_indices(contra_seed, len(contra_pair_ids), 1, offset=500)
            pid = contra_pair_ids[pair_idx[0]]
            if pid in contra_lookup:
                est, chl = contra_lookup[pid]
                result_prompts.append(est)
                result_prompts.append(chl)

        # Deterministic shuffle
        shuffle_seed = hashlib.sha256(mfa_seed + b"mfa_shuffle").digest()
        items = list(result_prompts)
        for i in range(len(items) - 1, 0, -1):
            h = hashlib.sha256(shuffle_seed + i.to_bytes(4, "big")).digest()
            j = int.from_bytes(h[:4], "big") % (i + 1)
            items[i], items[j] = items[j], items[i]

        # Enforce pair gap
        items = _enforce_pair_gap(items, min_gap=2)

        # Compute hashes for audit
        prompt_ids = sorted(p.id for p in items)
        selection_hash = hashlib.sha256(
            ":".join(prompt_ids).encode("utf-8")
        ).hexdigest()
        seed_hash = hashlib.sha256(
            f"{agent_id}:mfa:{model_id}:{secret}".encode("utf-8")
        ).hexdigest()

        return PromptSet(
            prompts=items,
            seed=seed_hash,
            selection_hash=selection_hash,
            mfa=True,
        )
