"""
SelfHealer — Tier 4: Automatically trigger recovery actions.
=============================================================

Closes the detect → diagnose → repair loop.  Takes a FailureMode from
FailureTaxonomist and returns a RecoveryAction — either a prompt injection
to continue the agent run, or a structural directive to the agent harness.

Recovery action types:

  REPHRASE_QUERY       — inject a new search query for the failed search step
  CONSOLIDATE          — inject "summarise what you know and give your best answer"
  VERIFY_ANSWER        — inject "cite which observation supports your answer"
  ADDITIONAL_SEARCH    — force at least one more search before finishing
  FORCE_FINISH         — inject "you've searched enough, give your best answer now"
  NO_ACTION            — risk too low to warrant intervention

Usage:
    from qppg_service import SelfHealer, FailureTaxonomist

    tx     = FailureTaxonomist()
    healer = SelfHealer()

    failure = tx.classify(question, steps, final_answer, finished)
    action  = healer.suggest(failure, question, steps)

    if action.action_type != "NO_ACTION":
        # Inject action.prompt_injection into your agent harness
        new_response = agent.continue_with(action.prompt_injection)

Integration with a ReAct agent:
    # After any Search step:
    if monitor.track(question, steps_so_far, partial_answer):
        failure = tx.classify(question, steps_so_far, partial_answer, finished=False)
        action  = healer.suggest(failure, question, steps_so_far)
        # Pass action.prompt_injection as next user turn
"""

from __future__ import annotations

import warnings
from dataclasses import dataclass, field
from typing import Optional

from .failure_taxonomy import FailureMode


# ── Action types ──────────────────────────────────────────────────────────────

REPHRASE_QUERY    = "REPHRASE_QUERY"
CONSOLIDATE       = "CONSOLIDATE"
VERIFY_ANSWER     = "VERIFY_ANSWER"
ADDITIONAL_SEARCH = "ADDITIONAL_SEARCH"
FORCE_FINISH      = "FORCE_FINISH"
NO_ACTION         = "NO_ACTION"


@dataclass
class RecoveryAction:
    """A single recovery action to inject into the agent harness."""

    action_type: str
    """One of: REPHRASE_QUERY, CONSOLIDATE, VERIFY_ANSWER, ADDITIONAL_SEARCH,
    FORCE_FINISH, NO_ACTION."""

    prompt_injection: str
    """The text to inject into the agent conversation as the next system / user
    message.  Empty string if action_type == NO_ACTION."""

    explanation: str
    """Why this action was chosen."""

    urgency: str = "LOW"
    """LOW / MEDIUM / HIGH — how strongly to enforce this action."""

    failure_mode: str = ""
    """The FailureMode.primary_mode that triggered this action."""

    metadata: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "action_type":       self.action_type,
            "prompt_injection":  self.prompt_injection,
            "explanation":       self.explanation,
            "urgency":           self.urgency,
            "failure_mode":      self.failure_mode,
        }


class SelfHealer:
    """
    Maps FailureMode → RecoveryAction.

    Rule-based (no ML required). Provides both:
    - A prompt_injection (ready to pass to the LLM)
    - A structured directive (for programmatic agent harness control)

    .. note::
        **Validated on 2 failure modes (exp_selfhealer_ab, 2026-03-20).**

        A/B test (99 HotpotQA wrong chains) confirmed ≥+10pp pass@1 improvement
        on two failure modes:

        - **RETRIEVAL_FAILURE**: +11.1pp (n=45, baseline 0.0%)
        - **EXCESSIVE_SEARCH**: +20.0pp (n=5, baseline 0.0%)

        Overall healed pass@1 = 6.1% vs 0.0% retry-only (Δ+6.1pp).

        **Caveats**: effect sizes are modest; ANSWER_UNSUPPORTED and
        INSUFFICIENT_EVIDENCE modes showed no improvement (n=2 each, too small).
        The LOW_RISK mode (n=45) showed no improvement, as expected (no injection
        needed). Results are specific to HotpotQA-style multi-hop chains — validate
        on your domain before relying on recovery actions in production.

    Parameters
    ----------
    enable_force_finish : bool
        If True, will inject a FORCE_FINISH prompt when the agent has
        searched too many times. Default True.
    max_steps_before_force : int
        Number of search steps above which FORCE_FINISH is triggered
        for EXCESSIVE_SEARCH mode. Default 5.
    """

    def __init__(
        self,
        enable_force_finish: bool = True,
        max_steps_before_force: int = 5,
    ):
        # exp_selfhealer_ab (2026-03-20): GATE PASSED — RETRIEVAL_FAILURE +11.1pp,
        # EXCESSIVE_SEARCH +20.0pp. Warning downgraded from EXPERIMENTAL to informational.
        # Effect is modest (6.1% overall healed) — see class docstring for caveats.
        self.enable_force_finish      = enable_force_finish
        self.max_steps_before_force   = max_steps_before_force

    # ── Public API ────────────────────────────────────────────────────────────

    def suggest(
        self,
        failure: FailureMode,
        question: str,
        steps: list[dict],
        partial_answer: str = "",
    ) -> RecoveryAction:
        """
        Suggest a recovery action given a failure diagnosis.

        Parameters
        ----------
        failure        : FailureMode from FailureTaxonomist.classify()
        question       : original question
        steps          : agent steps so far
        partial_answer : current best answer (may be empty if not yet finished)

        Returns
        -------
        RecoveryAction
        """
        import warnings as _warnings
        mode = failure.primary_mode
        n_steps = failure.signals.get("n_search_steps", 0)
        mean_sim = failure.signals.get("mean_sim", 0.5)

        # Validated modes only: RETRIEVAL_FAILURE (+11.1pp) and EXCESSIVE_SEARCH (+20pp)
        # from exp_selfhealer_ab (2026-03-20, n=99 chains).
        _VALIDATED_MODES = {"RETRIEVAL_FAILURE", "EXCESSIVE_SEARCH", "LOW_RISK"}

        # Priority order: most critical first
        if mode == "LOW_RISK":
            return RecoveryAction(
                action_type      = NO_ACTION,
                prompt_injection = "",
                explanation      = "Risk is low — no intervention needed.",
                urgency          = "LOW",
                failure_mode     = mode,
            )

        if mode == "RETRIEVAL_FAILURE":
            return self._rephrase_query(question, steps, failure)

        if mode == "EXCESSIVE_SEARCH":
            if self.enable_force_finish and n_steps >= self.max_steps_before_force:
                return self._force_finish(question, steps, failure)
            return self._consolidate(question, steps, failure)

        # ── Unvalidated modes: emit warning and return NO_ACTION ─────────────
        _warnings.warn(
            f"SelfHealer.suggest(): failure mode '{mode}' has no validated repair "
            f"strategy. Only RETRIEVAL_FAILURE (+11.1pp) and EXCESSIVE_SEARCH "
            f"(+20pp) are empirically validated (exp_selfhealer_ab, 2026-03-20). "
            f"Returning NO_ACTION to avoid unvalidated interventions. To suppress "
            f"this warning, filter chains to validated modes before calling suggest().",
            UserWarning,
            stacklevel=2,
        )
        return RecoveryAction(
            action_type      = NO_ACTION,
            prompt_injection = "",
            explanation      = (
                f"Mode '{mode}' has no validated repair strategy in SelfHealer. "
                f"Validated modes: RETRIEVAL_FAILURE, EXCESSIVE_SEARCH."
            ),
            urgency          = "LOW",
            failure_mode     = mode,
        )

    def suggest_batch(
        self,
        failures: list[FailureMode],
        chains: list[dict],
    ) -> list[RecoveryAction]:
        """Suggest recovery for a batch of failures."""
        return [
            self.suggest(
                f,
                c.get("question", ""),
                c.get("steps", []),
                c.get("final_answer", ""),
            )
            for f, c in zip(failures, chains)
        ]

    # ── Recovery action builders ──────────────────────────────────────────────

    @staticmethod
    def _rephrase_query(question: str, steps: list[dict],
                        failure: FailureMode) -> RecoveryAction:
        last_query = next(
            (s.get("action_arg","") for s in reversed(steps)
             if s.get("action_type") == "Search"), ""
        )
        mean_sim = failure.signals.get("mean_sim", 0.0)
        injection = (
            f"Your last search results were not very relevant to the question "
            f"(relevance score: {mean_sim:.2f}/1.0). "
            f"The question is: '{question[:200]}'\n"
            f"Your last query was: '{last_query[:150]}'\n\n"
            f"Please rephrase your search query to be more specific and directly "
            f"address the key facts needed to answer the question. "
            f"Try using different keywords or a more targeted query."
        )
        return RecoveryAction(
            action_type      = REPHRASE_QUERY,
            prompt_injection = injection,
            explanation      = (f"Search results had low relevance (mean_sim={mean_sim:.2f}). "
                                f"Rephrasing the query may improve retrieval."),
            urgency          = "HIGH",
            failure_mode     = failure.primary_mode,
        )

    @staticmethod
    def _consolidate(question: str, steps: list[dict],
                     failure: FailureMode) -> RecoveryAction:
        n_steps = failure.signals.get("n_search_steps", 0)
        injection = (
            f"You have performed {n_steps} searches. "
            f"Before searching further, take a moment to consolidate what you've found.\n\n"
            f"Question: '{question[:200]}'\n\n"
            f"Based on ALL the evidence you have retrieved so far, what is your best answer? "
            f"If the evidence is contradictory, explain which source is most credible and why. "
            f"Provide your final answer now."
        )
        return RecoveryAction(
            action_type      = CONSOLIDATE,
            prompt_injection = injection,
            explanation      = (f"Agent has done {n_steps} searches with conflicting or "
                                f"excessive results. Consolidation prompt to resolve confusion."),
            urgency          = "MEDIUM",
            failure_mode     = failure.primary_mode,
        )

    @staticmethod
    def _verify_answer(question: str, steps: list[dict], answer: str,
                       failure: FailureMode) -> RecoveryAction:
        injection = (
            f"You answered: '{answer[:200]}'\n\n"
            f"Before submitting this answer, please verify it against your retrieved evidence:\n"
            f"1. Which specific observation (search result) supports this answer?\n"
            f"2. Quote the relevant passage from your search results.\n"
            f"3. If your search results do NOT support this answer, search again for better evidence "
            f"or revise your answer to match what the evidence actually says.\n\n"
            f"Question: '{question[:200]}'"
        )
        return RecoveryAction(
            action_type      = VERIFY_ANSWER,
            prompt_injection = injection,
            explanation      = ("Answer words are absent from the reasoning chain — "
                                "agent may have hallucinated. Verification prompt added."),
            urgency          = "HIGH",
            failure_mode     = failure.primary_mode,
        )

    @staticmethod
    def _additional_search(question: str, steps: list[dict],
                           failure: FailureMode, urgent: bool = False) -> RecoveryAction:
        n_steps = failure.signals.get("n_search_steps", 0)
        if urgent:
            injection = (
                f"You have not gathered enough evidence to answer reliably (only {n_steps} search(es)). "
                f"You MUST search for more information before answering.\n\n"
                f"Question: '{question[:200]}'\n\n"
                f"Please perform at least one more targeted search to gather the key facts needed."
            )
        else:
            injection = (
                f"You have {n_steps} search results but the evidence may be incomplete. "
                f"Consider searching for one or two more specific facts related to:\n"
                f"'{question[:200]}'\n\n"
                f"If you have sufficient evidence, provide your final answer."
            )
        return RecoveryAction(
            action_type      = ADDITIONAL_SEARCH,
            prompt_injection = injection,
            explanation      = (f"Agent stopped too early ({n_steps} searches). "
                                f"Forcing additional evidence gathering."),
            urgency          = "HIGH" if urgent else "MEDIUM",
            failure_mode     = failure.primary_mode,
        )

    @staticmethod
    def _force_finish(question: str, steps: list[dict],
                      failure: FailureMode) -> RecoveryAction:
        n_steps = failure.signals.get("n_search_steps", 0)
        injection = (
            f"You have now performed {n_steps} searches. This is sufficient. "
            f"STOP searching and provide your final answer NOW.\n\n"
            f"Question: '{question[:200]}'\n\n"
            f"Based on everything you have found, give your best answer. "
            f"If you are uncertain, state your best guess with a brief explanation of why. "
            f"Do NOT search again — use Action: Finish[<your answer>]."
        )
        return RecoveryAction(
            action_type      = FORCE_FINISH,
            prompt_injection = injection,
            explanation      = (f"Agent has done {n_steps} searches — above max threshold. "
                                f"Forcing completion to prevent infinite search loops."),
            urgency          = "HIGH",
            failure_mode     = failure.primary_mode,
        )
