"""Baseline capture, storage, and comparison.

Manages the lifecycle of identity baselines:
- Capture: run assessment, vectorize responses, store baseline
- Load: retrieve baseline for an agent
- Compare: run a test and compute drift score against baseline
"""

from __future__ import annotations

from typing import Optional

import numpy as np

from kredo_drift.assessment import AssessmentProtocol, PromptSet
from kredo_drift.prompts import Prompt
from kredo_drift.contradiction import compute_pair_similarities, score_contradictions
from kredo_drift.models import (
    Dimension,
    DriftBaseline,
    DriftResult,
    DriftTest,
    PromptResponse,
    PROMPT_TYPE_WEIGHTS,
)
from kredo_drift.scoring import compute_drift_score, cosine_similarity
from kredo_drift.store import DriftStore
from kredo_drift.vectorize import Vectorizer
from kredo_drift.features import extract_features
from kredo_drift.anchors import compute_anchor_baseline, score_anchors
from kredo_drift.temporal import compute_temporal_metrics, compute_dimension_temporal


class BaselineManager:
    """Manages baseline capture, test execution, and drift scoring.

    This is the main orchestrator that ties together the assessment protocol,
    vectorization engine, scoring algorithm, and storage layer.
    """

    def __init__(
        self,
        store: Optional[DriftStore] = None,
        vectorizer: Optional[Vectorizer] = None,
        protocol: Optional[AssessmentProtocol] = None,
    ):
        self.store = store or DriftStore()
        self.vectorizer = vectorizer or Vectorizer()
        self.protocol = protocol or AssessmentProtocol()

    def get_assessment_prompts(
        self,
        agent_name: str,
        sequence: int = 0,
    ) -> PromptSet:
        """Get the set of prompts for an assessment.

        Args:
            agent_name: The agent being assessed.
            sequence: 0 for baseline, 1+ for tests.

        Returns:
            A PromptSet with prompts to present to the agent.
        """
        return self.protocol.select_prompts(agent_name, sequence)

    def create_baseline(
        self,
        agent_name: str,
        responses: list[PromptResponse],
    ) -> DriftBaseline:
        """Create an identity baseline from assessment responses.

        Vectorizes all responses, computes per-dimension and overall vectors,
        and stores the baseline.

        Args:
            agent_name: The agent being baselined.
            responses: The agent's responses to assessment prompts.

        Returns:
            The created DriftBaseline.
        """
        # Separate identity responses (scored), fidelity responses, and decoys
        scored = [r for r in responses if not r.is_decoy and not r.is_fidelity]
        fidelity = [r for r in responses if r.is_fidelity]

        # Vectorize all scored responses
        texts = [r.response_text for r in scored]
        vectors = self.vectorizer.encode_batch(texts)

        # WHY prompt-type weights: Direct self-report prompts ("What are your
        # values?") are easily gamed — an agent can state anything. Scenario
        # and tradeoff prompts reveal actual behavior under pressure, making
        # them more reliable identity signals. Weighting them 1.4-1.6x vs
        # direct at 0.6x shifts the identity vector toward behavioral truth.
        # Build prompt-type weights for weighted averaging
        prompt_weights = np.array([
            PROMPT_TYPE_WEIGHTS.get(r.prompt_type, 1.0) for r in scored
        ], dtype=np.float32)

        # Compute per-dimension vectors (weighted by prompt type)
        dimension_vectors: dict[str, list[float]] = {}
        for dimension in Dimension:
            dim_indices = [
                i for i, r in enumerate(scored)
                if r.dimension == dimension
            ]
            if dim_indices:
                dim_vecs = vectors[dim_indices]
                dim_weights = prompt_weights[dim_indices]
                avg = self.vectorizer.weighted_average_vectors(dim_vecs, dim_weights)
                dimension_vectors[dimension.value] = avg.tolist()

        # Compute intra-dimension coherence using centroid similarity.
        # Each response is compared to the dimension's centroid (weighted avg).
        # This measures "how well does each response represent the core of this
        # dimension?" — not "do all responses say the same thing?"
        # An agent can give diverse, nuanced answers (low pairwise sim) that
        # all strongly represent the same underlying value (high centroid sim).
        dimension_coherence: dict[str, float] = {}
        for dimension in Dimension:
            dim_indices = [
                i for i, r in enumerate(scored)
                if r.dimension == dimension
            ]
            if len(dim_indices) >= 2:
                dim_vecs = vectors[dim_indices]
                dim_weights = prompt_weights[dim_indices]
                centroid = self.vectorizer.weighted_average_vectors(dim_vecs, dim_weights)
                # Each response's similarity to the centroid
                sims = [cosine_similarity(dim_vecs[j], centroid) for j in range(len(dim_vecs))]
                avg_sim = sum(sims) / len(sims)
                # Map similarity to 0-100 with a more generous curve:
                # sim=1.0 -> 100, sim=0.5 -> 70, sim=0.3 -> 50, sim=0.0 -> 20
                # This reflects that embeddings rarely exceed 0.7 cosine sim
                # even for clearly related sentences.
                coherence = round(min(100.0, max(0.0, 20 + avg_sim * 80)), 1)
                dimension_coherence[dimension.value] = coherence
            elif len(dim_indices) == 1:
                dimension_coherence[dimension.value] = 50.0

        # Fidelity coherence
        if len(fidelity) >= 2:
            fid_texts = [r.response_text for r in fidelity]
            fid_vecs = self.vectorizer.encode_batch(fid_texts)
            fid_weights = np.array([
                PROMPT_TYPE_WEIGHTS.get(r.prompt_type, 1.0) for r in fidelity
            ], dtype=np.float32)
            fid_centroid = self.vectorizer.weighted_average_vectors(fid_vecs, fid_weights)
            fid_sims = [cosine_similarity(fid_vecs[j], fid_centroid) for j in range(len(fid_vecs))]
            avg_fid = sum(fid_sims) / len(fid_sims)
            dimension_coherence["fidelity"] = round(min(100.0, max(0.0, 20 + avg_fid * 80)), 1)

        # Compute overall averaged vector (weighted by prompt type)
        overall_vector = self.vectorizer.weighted_average_vectors(vectors, prompt_weights)

        # Compute contradiction pair similarities
        response_vectors: dict[str, np.ndarray] = {}
        for i, r in enumerate(scored):
            response_vectors[r.prompt_id] = vectors[i]
        pair_sims = compute_pair_similarities(response_vectors)

        # WHY extract features at baseline time: Anchor values are computed
        # from the baseline's behavioral features and stored permanently.
        # Each subsequent test extracts the same features and compares them
        # to the baseline anchors. This means behavioral drift (how the agent
        # reasons, refuses, hedges) is tracked independently of what it says.
        # Extract structured features and compute anchor baseline
        response_features = [
            extract_features(r.prompt_id, r.response_text) for r in scored
        ]
        anchor_baseline = compute_anchor_baseline(response_features)

        # WHY fidelity vector computed separately: Fidelity measures how well
        # the agent models its human — an independent axis from identity drift.
        # Averaging fidelity responses into their own vector space means we can
        # track "does the agent still understand its human?" without polluting
        # the 5-dimension identity score.
        fidelity_vector: list[float] = []
        if fidelity:
            fid_texts = [r.response_text for r in fidelity]
            fid_vectors = self.vectorizer.encode_batch(fid_texts)
            fid_weights = np.array([
                PROMPT_TYPE_WEIGHTS.get(r.prompt_type, 1.0) for r in fidelity
            ], dtype=np.float32)
            fid_avg = self.vectorizer.weighted_average_vectors(fid_vectors, fid_weights)
            fidelity_vector = fid_avg.tolist()


        # Include fidelity in dimension_vectors so it appears as the 6th dimension
        if fidelity_vector:
            dimension_vectors["fidelity"] = fidelity_vector
        # Build baseline
        baseline = DriftBaseline(
            agent_name=agent_name,
            vector=overall_vector.tolist(),
            dimension_vectors=dimension_vectors,
            response_hashes=[r.response_hash() for r in responses],
            prompt_ids=[r.prompt_id for r in responses],
            prompt_count=len(responses),
            pair_similarities=pair_sims,
            anchor_values=anchor_baseline,
            fidelity_vector=fidelity_vector,
            dimension_coherence=dimension_coherence,
        )

        # Store it
        self.store.save_baseline(baseline)

        return baseline

    def run_test(
        self,
        agent_name: str,
        responses: list[PromptResponse],
        baseline: Optional[DriftBaseline] = None,
    ) -> DriftResult:
        """Run a drift test against the agent's baseline.

        Vectorizes the test responses, computes drift score, and stores results.

        Args:
            agent_name: The agent being tested.
            responses: The agent's responses to test prompts.
            baseline: Optional baseline to compare against.
                      If None, loads the latest baseline from the store.

        Returns:
            A DriftResult with scores and classification.

        Raises:
            ValueError: If no baseline exists for the agent.
        """
        # Load baseline if not provided
        if baseline is None:
            baseline = self.store.load_baseline(agent_name)
            if baseline is None:
                raise ValueError(
                    f"No baseline found for agent '{agent_name}'. "
                    "Run 'drift baseline' first."
                )

        # Separate identity responses (scored), fidelity responses, and decoys
        scored = [r for r in responses if not r.is_decoy and not r.is_fidelity]
        fidelity = [r for r in responses if r.is_fidelity]

        # Vectorize all scored responses
        texts = [r.response_text for r in scored]
        vectors = self.vectorizer.encode_batch(texts)

        # Build prompt-type weights for weighted averaging
        prompt_weights = np.array([
            PROMPT_TYPE_WEIGHTS.get(r.prompt_type, 1.0) for r in scored
        ], dtype=np.float32)

        # Compute per-dimension vectors (weighted by prompt type)
        dimension_vectors: dict[str, list[float]] = {}
        for dimension in Dimension:
            dim_indices = [
                i for i, r in enumerate(scored)
                if r.dimension == dimension
            ]
            if dim_indices:
                dim_vecs = vectors[dim_indices]
                dim_weights = prompt_weights[dim_indices]
                avg = self.vectorizer.weighted_average_vectors(dim_vecs, dim_weights)
                dimension_vectors[dimension.value] = avg.tolist()

        # Compute intra-dimension coherence using centroid similarity.
        # Each response is compared to the dimension's centroid (weighted avg).
        # This measures "how well does each response represent the core of this
        # dimension?" — not "do all responses say the same thing?"
        # An agent can give diverse, nuanced answers (low pairwise sim) that
        # all strongly represent the same underlying value (high centroid sim).
        dimension_coherence: dict[str, float] = {}
        for dimension in Dimension:
            dim_indices = [
                i for i, r in enumerate(scored)
                if r.dimension == dimension
            ]
            if len(dim_indices) >= 2:
                dim_vecs = vectors[dim_indices]
                dim_weights = prompt_weights[dim_indices]
                centroid = self.vectorizer.weighted_average_vectors(dim_vecs, dim_weights)
                # Each response's similarity to the centroid
                sims = [cosine_similarity(dim_vecs[j], centroid) for j in range(len(dim_vecs))]
                avg_sim = sum(sims) / len(sims)
                # Map similarity to 0-100 with a more generous curve:
                # sim=1.0 -> 100, sim=0.5 -> 70, sim=0.3 -> 50, sim=0.0 -> 20
                # This reflects that embeddings rarely exceed 0.7 cosine sim
                # even for clearly related sentences.
                coherence = round(min(100.0, max(0.0, 20 + avg_sim * 80)), 1)
                dimension_coherence[dimension.value] = coherence
            elif len(dim_indices) == 1:
                dimension_coherence[dimension.value] = 50.0

        # Fidelity coherence
        if len(fidelity) >= 2:
            fid_texts = [r.response_text for r in fidelity]
            fid_vecs = self.vectorizer.encode_batch(fid_texts)
            fid_weights = np.array([
                PROMPT_TYPE_WEIGHTS.get(r.prompt_type, 1.0) for r in fidelity
            ], dtype=np.float32)
            fid_centroid = self.vectorizer.weighted_average_vectors(fid_vecs, fid_weights)
            fid_sims = [cosine_similarity(fid_vecs[j], fid_centroid) for j in range(len(fid_vecs))]
            avg_fid = sum(fid_sims) / len(fid_sims)
            dimension_coherence["fidelity"] = round(min(100.0, max(0.0, 20 + avg_fid * 80)), 1)

        # Compute overall averaged vector (weighted by prompt type)
        overall_vector = self.vectorizer.weighted_average_vectors(vectors, prompt_weights)

        # Compute contradiction pair similarities for this test
        response_vectors: dict[str, np.ndarray] = {}
        for i, r in enumerate(scored):
            response_vectors[r.prompt_id] = vectors[i]
        test_pair_sims = compute_pair_similarities(response_vectors)

        # Score contradictions against baseline
        contradiction_results = score_contradictions(
            baseline.pair_similarities, test_pair_sims,
        )

        # Extract structured features and score anchors
        response_features = [
            extract_features(r.prompt_id, r.response_text) for r in scored
        ]
        anchor_results = score_anchors(
            baseline.anchor_values, response_features,
        )

        # Compute fidelity vector and drift
        fidelity_vector: list[float] = []
        fidelity_score = None
        fidelity_similarity = None
        if fidelity:
            fid_texts = [r.response_text for r in fidelity]
            fid_vectors = self.vectorizer.encode_batch(fid_texts)
            fid_weights = np.array([
                PROMPT_TYPE_WEIGHTS.get(r.prompt_type, 1.0) for r in fidelity
            ], dtype=np.float32)
            fid_avg = self.vectorizer.weighted_average_vectors(fid_vectors, fid_weights)
            fidelity_vector = fid_avg.tolist()

            # WHY fidelity drift uses same cosine-similarity math: The agent's
            # relational understanding should be stable — if the fidelity vector
            # drifts, the agent's model of its human is shifting. Same 0-100
            # scale as identity drift, same interpretation (lower = more stable).
            if baseline.fidelity_vector:
                baseline_fid = np.array(baseline.fidelity_vector, dtype=np.float32)
                test_fid = np.array(fidelity_vector, dtype=np.float32)
                sim = float(cosine_similarity(baseline_fid, test_fid))
                fidelity_similarity = sim
                fidelity_score = max(0.0, min(100.0, (1.0 - sim) * 100.0))


        # Include fidelity in dimension_vectors so it appears as the 6th dimension
        if fidelity_vector:
            dimension_vectors["fidelity"] = fidelity_vector
        # Build test record
        test = DriftTest(
            agent_name=agent_name,
            baseline_id=baseline.id,
            vector=overall_vector.tolist(),
            dimension_vectors=dimension_vectors,
            response_hashes=[r.response_hash() for r in responses],
            prompt_ids=[r.prompt_id for r in responses],
            prompt_count=len(responses),
            fidelity_vector=fidelity_vector,
        )

        # Compute drift score with contradiction amplification
        result = compute_drift_score(
            baseline, test, contradictions=contradiction_results,
        )

        # WHY attach anchors post-scoring: Anchors are metadata that describe
        # behavioral stability, not a scoring input. Keeping them separate from
        # compute_drift_score() means the core cosine-similarity score remains
        # deterministic and unchanged. Anchors add diagnostic depth without
        # altering the number that goes into attestations.
        # Attach anchor results
        result.anchors = [
            {
                "name": a.name,
                "dimension": a.dimension,
                "feature": a.feature,
                "baseline_value": a.baseline_value,
                "current_value": a.current_value,
                "drift": a.drift,
                "stable": a.stable,
            }
            for a in anchor_results
        ]

        # Attach fidelity drift (separate from identity drift score)
        result.fidelity_score = fidelity_score
        result.fidelity_similarity = fidelity_similarity

        # Store test and result
        self.store.save_test(test)
        self.store.save_result(result)

        # WHY compute temporal after save: The result must be stored before
        # we query history, so this test appears in the history list. Temporal
        # metrics are then attached to the in-memory result object (not re-saved)
        # because they're ephemeral — recalculated each time from the history
        # window, not frozen at test time.
        # Compute temporal metrics from history
        history = self.store.load_results(agent_name, limit=10)
        if len(history) >= 3:
            temporal = compute_temporal_metrics(result.score, history)
            if temporal:
                result.velocity = temporal.velocity
                result.acceleration = temporal.acceleration
                result.trend_direction = temporal.trend_direction

            dim_trends = compute_dimension_temporal(
                result.dimension_scores, history,
            )
            if dim_trends:
                result.dimension_trends = dim_trends

        return result

    def get_latest_result(self, agent_name: str) -> Optional[DriftResult]:
        """Get the most recent drift result for an agent."""
        return self.store.load_latest_result(agent_name)

    def get_history(self, agent_name: str, limit: int = 20) -> list[DriftResult]:
        """Get drift score history for an agent."""
        return self.store.load_results(agent_name, limit=limit)

    def has_baseline(self, agent_name: str) -> bool:
        """Check if an agent has an existing baseline."""
        return self.store.load_baseline(agent_name) is not None
