"""
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"
  RESOLVE_CONFLICT     — inject "identify which source is authoritative and resolve contradiction"
  ANCHOR_TO_EVIDENCE   — inject "revise your answer to match what the evidence says"
  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"
RESOLVE_CONFLICT  = "RESOLVE_CONFLICT"
ANCHOR_TO_EVIDENCE = "ANCHOR_TO_EVIDENCE"
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,
        experimental: bool = False,
    ):
        # 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
        self.experimental             = experimental

    # ── 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: RETRIEVAL_FAILURE (+11.1pp) and EXCESSIVE_SEARCH (+20pp)
        # from exp_selfhealer_ab (2026-03-20, n=99 chains).
        # LOW_RISK is intentionally excluded — it produces NO_ACTION and has no
        # repair effect; listing it as "validated" was misleading.
        _VALIDATED_MODES = {"RETRIEVAL_FAILURE", "EXCESSIVE_SEARCH"}

        # 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: gate on experimental flag ──────────────────────
        if not self.experimental:
            _warnings.warn(
                f"SelfHealer.suggest(): failure mode '{mode}' has not been externally "
                f"validated. Only RETRIEVAL_FAILURE (+11.1pp) and EXCESSIVE_SEARCH "
                f"(+20pp) are empirically validated (exp_selfhealer_ab, 2026-03-20). "
                f"Pass experimental=True to SelfHealer() to enable unvalidated modes. "
                f"Returning NO_ACTION to avoid unvalidated interventions.",
                UserWarning,
                stacklevel=2,
            )
            return RecoveryAction(
                action_type      = NO_ACTION,
                prompt_injection = "",
                explanation      = (
                    f"Mode '{mode}' not validated — pass experimental=True to enable. "
                    f"Validated modes: RETRIEVAL_FAILURE (+11.1pp), EXCESSIVE_SEARCH (+20pp)."
                ),
                urgency          = "LOW",
                failure_mode     = mode,
            )

        # experimental=True: heuristic repair strategies for unvalidated modes
        if mode == "CONFLICTING_EVIDENCE":
            return self._resolve_conflict(question, steps, failure)

        if mode == "ANSWER_UNSUPPORTED":
            return self._anchor_to_evidence(question, steps, partial_answer, failure)

        # All other experimental modes: ADDITIONAL_SEARCH as best-effort fallback
        if mode == "PREMATURE_STOP":
            return self._additional_search(question, steps, failure, urgent=True)

        if mode == "INSUFFICIENT_EVIDENCE":
            return self._additional_search(question, steps, failure, urgent=False)

        # Unknown mode: emit informational warning
        _warnings.warn(
            f"SelfHealer.suggest(): failure mode '{mode}' has no repair strategy. "
            f"Returning NO_ACTION (experimental=True).",
            UserWarning,
            stacklevel=2,
        )
        return RecoveryAction(
            action_type      = NO_ACTION,
            prompt_injection = "",
            explanation      = f"Mode '{mode}' has no repair strategy in SelfHealer.",
            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 _resolve_conflict(question: str, steps: list[dict],
                          failure: FailureMode) -> RecoveryAction:
        """Repair prompt for CONFLICTING_EVIDENCE: focus on most authoritative source."""
        obs_contradiction = failure.signals.get("obs_contradiction", 0.0)
        injection = (
            f"Your search results contain conflicting or contradictory information "
            f"(contradiction score: {obs_contradiction:.2f}).\n\n"
            f"Question: '{question[:200]}'\n\n"
            f"Please review your observations carefully and:\n"
            f"1. Identify which sources contradict each other.\n"
            f"2. Determine which source is most authoritative, recent, or specific to the question.\n"
            f"3. If the conflict cannot be resolved, state both perspectives clearly and "
            f"indicate which is more likely to be correct and why.\n"
            f"Focus on the most direct, authoritative evidence and provide your final answer."
        )
        return RecoveryAction(
            action_type      = RESOLVE_CONFLICT,
            prompt_injection = injection,
            explanation      = (
                f"Contradictory observations detected (obs_contradiction={obs_contradiction:.2f}). "
                f"Agent directed to identify and resolve conflicting evidence."
            ),
            urgency          = "MEDIUM",
            failure_mode     = failure.primary_mode,
        )

    @staticmethod
    def _anchor_to_evidence(question: str, steps: list[dict], answer: str,
                            failure: FailureMode) -> RecoveryAction:
        """Repair prompt for ANSWER_UNSUPPORTED: re-anchor answer to retrieved evidence."""
        answer_sim  = failure.signals.get("answer_sim",  0.0)
        answer_gap  = failure.signals.get("answer_gap",  0.0)
        injection = (
            f"Your answer does not appear to be well-supported by the evidence you retrieved "
            f"(answer-evidence alignment score: {answer_sim:.2f}/1.0).\n\n"
            f"Your current answer: '{answer[:200]}'\n"
            f"Question: '{question[:200]}'\n\n"
            f"Please review your search observations and:\n"
            f"1. Find a specific passage from your observations that directly supports your answer.\n"
            f"2. If no such passage exists, search for more relevant information.\n"
            f"3. Revise your answer to reflect only what the evidence actually says — "
            f"do not add facts that were not in your search results.\n"
            f"Provide a revised answer grounded in your retrieved observations."
        )
        return RecoveryAction(
            action_type      = ANCHOR_TO_EVIDENCE,
            prompt_injection = injection,
            explanation      = (
                f"Answer not grounded in retrieved evidence "
                f"(answer_sim={answer_sim:.2f}, answer_gap={answer_gap:.2f}). "
                f"Agent directed to anchor answer to observed facts."
            ),
            urgency          = "HIGH",
            failure_mode     = failure.primary_mode,
        )

    def heal_with_retry(
        self,
        chain: dict,
        failure_mode: str,
        execute_fn,
        max_attempts: int = 3,
    ) -> dict:
        """Attempt to heal a chain, retrying with different strategies up to max_attempts.

        Args:
            chain: the chain dict to heal (keys: question, steps, final_answer, etc.)
            failure_mode: string name e.g. 'RETRIEVAL_FAILURE', 'CONFLICTING_EVIDENCE'
            execute_fn: callable(strategy_name: str, chain: dict) → dict with keys:
                        'success' (bool), 'new_chain' (dict)
            max_attempts: maximum number of attempts before giving up.
        Returns:
            dict with keys: success (bool), new_chain (dict), attempts (int), strategy_used (str).
        """
        current_chain = chain
        strategy = "ADDITIONAL_SEARCH"  # default fallback
        for attempt in range(max_attempts):
            # Build a minimal FailureMode-like object or just use the string
            # to get the suggested action type
            action = self._suggest_for_mode(failure_mode, current_chain)
            strategy = action if isinstance(action, str) else getattr(action, "action_type", "ADDITIONAL_SEARCH")
            outcome = execute_fn(strategy, current_chain)
            if outcome.get("success", False):
                return {
                    "success": True,
                    "new_chain": outcome.get("new_chain", current_chain),
                    "attempts": attempt + 1,
                    "strategy_used": strategy,
                }
            current_chain = outcome.get("new_chain", current_chain)
        return {
            "success": False,
            "new_chain": current_chain,
            "attempts": max_attempts,
            "strategy_used": strategy,
        }

    def _suggest_for_mode(self, failure_mode: str, chain: dict) -> str:
        """Map failure_mode string to a strategy string."""
        mapping = {
            "RETRIEVAL_FAILURE": "ADDITIONAL_SEARCH",
            "EXCESSIVE_SEARCH": "FORCE_FINISH",
            "CONFLICTING_EVIDENCE": "RESOLVE_CONFLICT",
            "ANSWER_UNSUPPORTED": "ANCHOR_TO_EVIDENCE",
            "PREMATURE_STOP": "ADDITIONAL_SEARCH",
            "INSUFFICIENT_EVIDENCE": "ADDITIONAL_SEARCH",
            "LOW_RISK": "NO_ACTION",
        }
        return mapping.get(failure_mode, "ADDITIONAL_SEARCH")

    @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,
        )
