"""
FailureTaxonomist — Tier 3: Diagnose WHY an agent chain failed.
===============================================================

Converts a raw risk score into an actionable failure mode using
the signal families from exp39-44:

  Behavioral  (exp39): step count, completion, hedging, answer gap
  Retrieval   (exp43): question↔observation cosine similarity
  Structural  (exp40): thought variance, answer alignment

Failure modes (non-exclusive — a chain can have multiple):

  PREMATURE_STOP      — finished after very few searches, likely guessed
  EXCESSIVE_SEARCH    — many more searches than typical; agent was struggling
  RETRIEVAL_FAILURE   — observations are irrelevant to the question
  CONFLICTING_EVIDENCE — high search diversity + high thought variance
  INSUFFICIENT_EVIDENCE — low mean relevance AND low n_relevant
  ANSWER_UNSUPPORTED  — answer words absent from reasoning chain
  LOW_RISK            — no significant failure signal detected

Usage:
    from qppg_service import FailureTaxonomist

    tx = FailureTaxonomist()
    result = tx.classify(question, steps, final_answer, finished)

    print(result.primary_mode)      # "EXCESSIVE_SEARCH"
    print(result.explanation)       # human-readable string
    print(result.signals)           # dict of raw signal values
    for mode in result.all_modes:   # non-exclusive list
        print(mode)
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

import numpy as np

# ── Thresholds derived from exp39/43 empirical distributions ─────────────────

# Correct chains: avg 2.34 steps; incorrect: avg 4.75 — use 4 as threshold
_STEP_HIGH_THRESHOLD    = 4       # > this → likely struggling
_STEP_LOW_THRESHOLD     = 1       # <= this + not finished → premature stop

# Retrieval relevance thresholds (exp43)
_GOOD_MEAN_SIM          = 0.50    # mean cosine sim(question, obs) for good retrieval
_POOR_MEAN_SIM          = 0.35    # below this → retrieval failure
_GOOD_MIN_SIM           = 0.35    # min_sim for correct chains avg 0.456
_POOR_MIN_SIM           = 0.22    # below this → at least one useless search

# Thought variance (exp39 SC3)
_HIGH_THOUGHT_VAR       = 0.006   # above → conflicting reasoning

# Uncertainty words density (exp39 SC5)
_HIGH_UNCERTAINTY       = 0.5     # per 100 words

# Answer gap (exp39 SC6)
_HIGH_ANSWER_GAP        = 0.70    # answer words absent from reasoning


@dataclass
class FailureMode:
    """Diagnosis result for a single agent chain."""

    primary_mode: str
    """The single most likely failure mode."""

    all_modes: list[str]
    """All detected failure modes (non-exclusive)."""

    confidence: float
    """Confidence in primary_mode [0,1]. Higher = stronger evidence."""

    explanation: str
    """Human-readable explanation of the failure."""

    signals: dict = field(default_factory=dict)
    """Raw signal values used to make the decision."""

    recommendation: str = ""
    """Actionable recommendation for the operator / self-healer."""

    def to_dict(self) -> dict:
        return {
            "primary_mode":  self.primary_mode,
            "all_modes":     self.all_modes,
            "confidence":    round(self.confidence, 3),
            "explanation":   self.explanation,
            "recommendation": self.recommendation,
            "signals":       {k: round(v, 4) if isinstance(v, float) else v
                              for k, v in self.signals.items()},
        }


class FailureTaxonomist:
    """
    Rule-based failure mode classifier using exp39-44 behavioral + retrieval signals.

    No training required.  Classification is deterministic given the same chain.

    Parameters
    ----------
    embedder_model : str
        SentenceTransformer model for retrieval quality computation.
        Default: "all-MiniLM-L6-v2"
    """

    def __init__(self, embedder_model: str = "all-MiniLM-L6-v2"):
        self._embedder_model = embedder_model
        self._embedder       = None

    def _get_embedder(self):
        if self._embedder is None:
            from sentence_transformers import SentenceTransformer
            self._embedder = SentenceTransformer(self._embedder_model)
        return self._embedder

    def classify(
        self,
        question: str,
        steps: list[dict],
        final_answer: str = "",
        finished: bool = True,
    ) -> FailureMode:
        """
        Classify the failure mode of a completed agent chain.

        Parameters
        ----------
        question     : original question
        steps        : agent steps (thought/action_type/action_arg/observation)
        final_answer : agent's final answer
        finished     : True if agent called Finish[], False if fallback

        Returns
        -------
        FailureMode with primary_mode, all_modes, explanation, recommendation
        """
        import re

        # ── Raw signals ───────────────────────────────────────────────────────
        search_steps = [s for s in steps if s.get("action_type") == "Search"]
        n_steps      = len(search_steps)

        # Thought variance (SC3)
        embedder = self._get_embedder()
        thoughts = [s.get("thought", "").strip() for s in steps if s.get("thought","").strip()]
        thought_var = 0.0
        if len(thoughts) >= 2:
            embs = embedder.encode(thoughts, normalize_embeddings=True, show_progress_bar=False)
            thought_var = float(embs.var(axis=0).mean())

        # Uncertainty words density (SC5)
        _UW = ["not sure","unclear","might","may not","possibly","perhaps","however",
               "although","uncertain","seems","appears","probably","could be","i think",
               "i believe","not certain","unsure","i'm not","don't know","confusing",
               "it's possible","not clear","ambiguous"]
        all_t = " ".join(s.get("thought","") for s in steps).lower()
        wc = max(1, len(all_t.split()))
        uncertainty = sum(all_t.count(w) for w in _UW) / wc * 100

        # Answer gap (SC6)
        _STOP = {"the","a","an","is","it","in","of","to","and","or","that","this",
                 "was","are","for","with","on","at","by","be","has","had","i"}
        fa_words = set(re.findall(r'\w+', final_answer.lower())) - _STOP
        th_words = set(re.findall(r'\w+', all_t)) - _STOP
        answer_gap = (1.0 - len(fa_words & th_words) / len(fa_words)) if fa_words else 0.5

        # Retrieval quality (exp43)
        obs_list = [s.get("observation","").strip()[:300]
                    for s in steps if s.get("observation","").strip()]
        mean_sim = max_sim = min_sim = 0.0
        if obs_list:
            q_emb    = embedder.encode([question], normalize_embeddings=True,
                                        show_progress_bar=False)[0]
            obs_embs = embedder.encode(obs_list, normalize_embeddings=True,
                                        show_progress_bar=False)
            sims     = (obs_embs @ q_emb).tolist()
            mean_sim = float(np.mean(sims))
            max_sim  = float(np.max(sims))
            min_sim  = float(np.min(sims))

        # ── Failure mode detection ─────────────────────────────────────────────
        modes = []
        evidences = {}

        # PREMATURE_STOP: very few searches + didn't finish properly OR finished too quickly
        if n_steps <= _STEP_LOW_THRESHOLD and not finished:
            modes.append("PREMATURE_STOP")
            evidences["premature_stop"] = f"only {n_steps} search(es), no clean finish"
        elif n_steps == 0 and finished:
            modes.append("PREMATURE_STOP")
            evidences["premature_stop"] = "finished with zero searches — likely guessed"

        # EXCESSIVE_SEARCH: too many steps
        if n_steps >= _STEP_HIGH_THRESHOLD:
            modes.append("EXCESSIVE_SEARCH")
            evidences["excessive_search"] = f"{n_steps} search steps (threshold: {_STEP_HIGH_THRESHOLD})"

        # RETRIEVAL_FAILURE: observations are irrelevant to the question
        if obs_list and mean_sim < _POOR_MEAN_SIM:
            modes.append("RETRIEVAL_FAILURE")
            evidences["retrieval_failure"] = (f"mean_sim={mean_sim:.3f} < {_POOR_MEAN_SIM} threshold; "
                                              f"search results appear off-topic")
        elif obs_list and min_sim < _POOR_MIN_SIM:
            modes.append("RETRIEVAL_FAILURE")
            evidences["retrieval_failure"] = (f"min_sim={min_sim:.3f} — at least one "
                                              f"observation irrelevant to question")

        # INSUFFICIENT_EVIDENCE: searched but results weakly relevant
        if (obs_list and _POOR_MEAN_SIM <= mean_sim < _GOOD_MEAN_SIM
                and n_steps >= 2 and "RETRIEVAL_FAILURE" not in modes):
            modes.append("INSUFFICIENT_EVIDENCE")
            evidences["insufficient_evidence"] = (f"mean_sim={mean_sim:.3f} — evidence found "
                                                  f"but weakly relevant; may need more searches")

        # CONFLICTING_EVIDENCE: high thought variance AND high search diversity
        search_queries = [s.get("action_arg","") for s in search_steps]
        query_diversity = 0.0
        if len(search_queries) >= 2:
            q_embs = embedder.encode(search_queries, normalize_embeddings=True, show_progress_bar=False)
            q_sim  = q_embs @ q_embs.T
            n = len(search_queries)
            dists = [1.0 - q_sim[i,j] for i in range(n) for j in range(i+1,n)]
            query_diversity = float(np.mean(dists)) if dists else 0.0

        if thought_var > _HIGH_THOUGHT_VAR and query_diversity > 0.5:
            modes.append("CONFLICTING_EVIDENCE")
            evidences["conflicting_evidence"] = (
                f"thought_var={thought_var:.4f} (high) AND query_diversity={query_diversity:.3f} "
                f"— agent repeatedly pivoting, evidence may be contradictory")

        # ANSWER_UNSUPPORTED: final answer not grounded in reasoning
        if answer_gap > _HIGH_ANSWER_GAP and fa_words:
            modes.append("ANSWER_UNSUPPORTED")
            evidences["answer_unsupported"] = (
                f"answer_gap={answer_gap:.3f} — final answer words missing from "
                f"reasoning chain (agent may have invented the answer)")

        # LOW_RISK: no failure signals
        if not modes:
            modes.append("LOW_RISK")

        # ── Priority ranking for primary_mode ─────────────────────────────────
        priority = ["RETRIEVAL_FAILURE", "CONFLICTING_EVIDENCE", "EXCESSIVE_SEARCH",
                    "INSUFFICIENT_EVIDENCE", "ANSWER_UNSUPPORTED", "PREMATURE_STOP", "LOW_RISK"]
        primary = next((m for m in priority if m in modes), modes[0])

        # ── Confidence ────────────────────────────────────────────────────────
        # More modes = higher confidence something is wrong; use n_modes as proxy
        n_failure_modes = len([m for m in modes if m != "LOW_RISK"])
        confidence = min(1.0, 0.5 + 0.15 * n_failure_modes)
        if primary == "LOW_RISK":
            confidence = 0.85   # high confidence nothing is wrong

        # ── Explanation + recommendation ──────────────────────────────────────
        explanation = _build_explanation(primary, evidences, n_steps, mean_sim, finished)
        recommendation = _build_recommendation(modes, n_steps, mean_sim, finished, thought_var)

        return FailureMode(
            primary_mode    = primary,
            all_modes       = modes,
            confidence      = confidence,
            explanation     = explanation,
            recommendation  = recommendation,
            signals = {
                "n_search_steps":   n_steps,
                "finished":         finished,
                "thought_variance": round(thought_var, 5),
                "uncertainty_density": round(uncertainty, 3),
                "answer_gap":       round(answer_gap, 3),
                "mean_sim":         round(mean_sim, 4),
                "max_sim":          round(max_sim, 4),
                "min_sim":          round(min_sim, 4),
                "query_diversity":  round(query_diversity, 3),
            },
        )

    def classify_batch(self, chains: list[dict]) -> list[FailureMode]:
        """Classify a batch of chains."""
        return [
            self.classify(
                c["question"],
                c.get("steps", []),
                c.get("final_answer", ""),
                c.get("finished", True),
            )
            for c in chains
        ]


# ── Helper functions ───────────────────────────────────────────────────────────

def _build_explanation(primary: str, evidences: dict, n_steps: int,
                       mean_sim: float, finished: bool) -> str:
    if primary == "LOW_RISK":
        return "No significant failure signals detected. Answer appears well-supported."
    explanations = {
        "RETRIEVAL_FAILURE": (
            f"The search engine returned results with low relevance to the question "
            f"(mean similarity {mean_sim:.2f}). The agent may have used poor search queries "
            f"or the topic is outside the retrieval system's coverage."
        ),
        "CONFLICTING_EVIDENCE": (
            f"The agent's reasoning pivoted significantly across {n_steps} searches, "
            f"suggesting the retrieved evidence was contradictory or the agent was confused."
        ),
        "EXCESSIVE_SEARCH": (
            f"The agent performed {n_steps} searches, more than typical for correct chains "
            f"(avg 2.3 for correct vs 4.75 for incorrect). Excessive searching correlates "
            f"strongly with incorrect final answers."
        ),
        "INSUFFICIENT_EVIDENCE": (
            f"Evidence was found but was only weakly relevant (mean sim {mean_sim:.2f}). "
            f"The agent may need more or better searches to answer confidently."
        ),
        "ANSWER_UNSUPPORTED": (
            f"The final answer contains words not present in the reasoning chain, "
            f"suggesting the agent may have introduced information not in the retrieved evidence."
        ),
        "PREMATURE_STOP": (
            f"The agent stopped with {n_steps} search(es) and "
            f"{'a clean finish' if finished else 'no clean finish'}. "
            f"This suggests the agent guessed without sufficient evidence."
        ),
    }
    base = explanations.get(primary, "Unknown failure mode.")
    if len(evidences) > 1:
        other = [v for k, v in evidences.items() if k != primary.lower()]
        if other:
            base += f" Additional signals: {'; '.join(other[:2])}."
    return base


def _build_recommendation(modes: list[str], n_steps: int, mean_sim: float,
                           finished: bool, thought_var: float) -> str:
    recs = []
    if "RETRIEVAL_FAILURE" in modes:
        recs.append("rephrase search queries to be more specific to the question domain")
    if "EXCESSIVE_SEARCH" in modes:
        recs.append("inject consolidation prompt: 'Based on what you've found so far, what is the best answer?'")
    if "INSUFFICIENT_EVIDENCE" in modes:
        recs.append("try alternative search queries or increase max_steps")
    if "CONFLICTING_EVIDENCE" in modes:
        recs.append("inject reasoning prompt: 'Reconcile the conflicting evidence you found'")
    if "ANSWER_UNSUPPORTED" in modes:
        recs.append("request answer verification: 'Cite which observation supports your answer'")
    if "PREMATURE_STOP" in modes:
        recs.append("force at least one search before allowing Finish[]")
    if not recs or modes == ["LOW_RISK"]:
        return "No intervention needed."
    return " | ".join(recs)
