"""
JudgePostHook — wires SmartSelector → LLMJudge → AutoFixApplicator into the
InterceptorChain post-call hook. Evaluates every agent step within a 500ms timeout.
On RETRY decision, applies fix and POSTs result to backend.
"""

import asyncio
import logging
from typing import TYPE_CHECKING

import httpx

from ..judge.evaluator import JudgeConfig, JudgeDecision, LLMJudge
from ..judge.models import EnrichedContext, RootCause
from ..judge.smart_selector import SmartSelector
from ..realtime.auto_fix import AutoFixApplicator, FixStrategy
from .protocols import (
    FixAction,
    FixActionType,
    InterceptionContext,
    InterceptionDecision,
    PostCallResult,
)

if TYPE_CHECKING:
    from ..judge.evaluator import SpanEvaluation

logger = logging.getLogger("aigie.judge_hook")

JUDGE_TIMEOUT_MS = 500.0
MAX_RETRIES = 3

# Complete mapping: every RootCause → FixStrategy (no gaps)
ROOT_CAUSE_TO_STRATEGY: dict[RootCause, FixStrategy] = {
    # Hallucination → retry with grounding instruction
    RootCause.UNGROUNDED_CLAIM: FixStrategy.RETRY_REQUEST,
    RootCause.FABRICATED_REFERENCE: FixStrategy.RETRY_REQUEST,
    RootCause.OVERCONFIDENT_ASSERTION: FixStrategy.RETRY_REQUEST,
    RootCause.TOOL_RESULT_MISMATCH: FixStrategy.RETRY_REQUEST,
    # Drift → inject corrective context
    RootCause.TASK_DEVIATION: FixStrategy.INJECT_INSTRUCTION,
    RootCause.TOPIC_SHIFT: FixStrategy.INJECT_INSTRUCTION,
    RootCause.ROLE_DEVIATION: FixStrategy.INJECT_INSTRUCTION,
    # Safety → override response, execution continues
    RootCause.HARMFUL_CONTENT: FixStrategy.OVERRIDE_RESPONSE,
    RootCause.PII_LEAK: FixStrategy.OVERRIDE_RESPONSE,
    RootCause.INJECTION_SUCCESS: FixStrategy.OVERRIDE_RESPONSE,
    RootCause.POLICY_VIOLATION: FixStrategy.OVERRIDE_RESPONSE,
    # Quality → retry with clarity instruction
    RootCause.INCOMPLETE_RESPONSE: FixStrategy.RETRY_REQUEST,
    RootCause.INCOHERENT_LOGIC: FixStrategy.RETRY_REQUEST,
    RootCause.IRRELEVANT_CONTENT: FixStrategy.RETRY_REQUEST,
    RootCause.VAGUE_RESPONSE: FixStrategy.RETRY_REQUEST,
    # Tool usage → retry with corrected tool call
    RootCause.WRONG_TOOL: FixStrategy.RETRY_REQUEST,
    RootCause.BAD_PARAMETERS: FixStrategy.RETRY_REQUEST,
    RootCause.MISINTERPRETED_RESULT: FixStrategy.RETRY_REQUEST,
    RootCause.UNHANDLED_TOOL_ERROR: FixStrategy.RETRY_REQUEST,
    # Loop → override to break loop
    RootCause.STUCK_LOOP: FixStrategy.OVERRIDE_RESPONSE,
    RootCause.CIRCULAR_REASONING: FixStrategy.OVERRIDE_RESPONSE,
    RootCause.TOOL_RETRY_LOOP: FixStrategy.OVERRIDE_RESPONSE,
    # Fallback
    RootCause.UNKNOWN: FixStrategy.INJECT_INSTRUCTION,
}

_STRATEGY_TO_ACTION_TYPE: dict[FixStrategy, FixActionType] = {
    FixStrategy.RETRY_REQUEST: FixActionType.RETRY,
    FixStrategy.INJECT_INSTRUCTION: FixActionType.INJECT_INSTRUCTION,
    FixStrategy.OVERRIDE_RESPONSE: FixActionType.OVERRIDE_RESPONSE,
    FixStrategy.MODIFY_RESPONSE: FixActionType.MODIFY_RESPONSE,
    FixStrategy.FALLBACK_MODEL: FixActionType.FALLBACK,
    FixStrategy.TRUNCATE_CONTEXT: FixActionType.TRUNCATE_CONTEXT,
}


class JudgePostHook:
    """Post-call hook: SmartSelector → LLMJudge → AutoFixApplicator → backend POST."""

    name = "judge_post_hook"
    priority = 90  # Runs after other hooks

    def __init__(
        self,
        judge: LLMJudge | None = None,
        selector: SmartSelector | None = None,
        fixer: AutoFixApplicator | None = None,
        backend_url: str | None = None,
        api_key: str | None = None,
    ):
        self._judge = judge or LLMJudge(JudgeConfig())
        self._selector = selector or SmartSelector()
        self._fixer = fixer or AutoFixApplicator()
        self._backend_url = backend_url
        self._api_key = api_key
        self._last_original_output: str | None = None  # For testing

    async def __call__(self, ctx: InterceptionContext) -> PostCallResult:
        if ctx.retry_count >= MAX_RETRIES:
            logger.warning(
                f"[judge_hook] max retries ({MAX_RETRIES}) reached for span {ctx.span_id}, returning as-is"
            )
            return PostCallResult.allow(hook_name=self.name)

        original_output = ctx.response_content or ""
        self._last_original_output = original_output

        # Run judge within 500ms timeout
        evaluation = await self._evaluate_with_timeout(ctx)

        if evaluation is None:
            # Timeout — heuristic fallback: pass (fail open)
            return PostCallResult.allow(hook_name=self.name, latency_ms=JUDGE_TIMEOUT_MS)

        result = evaluation.result
        eval_time = evaluation.evaluation_time_ms

        if evaluation.decision == JudgeDecision.PASS or not result.has_issues:
            # Fire-and-forget backend POST for PASS
            try:
                asyncio.create_task(
                    self._post_to_backend(ctx, evaluation, original_output, None, "pass")
                )
            except RuntimeError:
                pass  # No running event loop — best effort only
            return PostCallResult.allow(hook_name=self.name, latency_ms=eval_time)

        # RETRY or STOP decision
        fix_instruction = (
            result.suggested_remediation or result.reasoning or "Retry with corrections."
        )

        # Default strategy: INJECT_INSTRUCTION (no structured diagnosis from base LLMJudge)
        action_type = FixActionType.INJECT_INSTRUCTION

        # If STOP — override response, no retry
        if evaluation.decision == JudgeDecision.STOP or result.should_stop:
            action_type = FixActionType.OVERRIDE_RESPONSE

        fix = FixAction(
            action_type=action_type,
            confidence=evaluation.confidence,
            reason=fix_instruction,
            parameters={"instruction": fix_instruction},
        )

        fix_result = await self._fixer.apply_fixes(ctx, [fix])
        fixed_output = (
            fix_result.modified_response or original_output
            if fix_result.success
            else original_output
        )

        try:
            asyncio.create_task(
                self._post_to_backend(
                    ctx, evaluation, original_output, str(fixed_output), action_type.value
                )
            )
        except RuntimeError:
            pass  # No running event loop — best effort only

        if action_type == FixActionType.INJECT_INSTRUCTION and result.should_retry:
            return PostCallResult.retry(
                reason=fix_instruction,
                instruction=fix_instruction,
                hook_name=self.name,
                latency_ms=eval_time,
            )

        # OVERRIDE or fix failed — modify content, no retry
        return PostCallResult.modify(
            content=str(fixed_output) if fixed_output else original_output,
            reason=f"judge_override:{evaluation.decision.value}",
            hook_name=self.name,
            latency_ms=eval_time,
        )

    async def _evaluate_with_timeout(self, ctx: InterceptionContext) -> "SpanEvaluation | None":
        """Run LLMJudge.evaluate_span with 500ms hard timeout. Returns None on timeout."""
        enriched = EnrichedContext(
            span_type=getattr(ctx, "span_type", "llm"),
            model=ctx.model,
        )
        selected_judges = self._selector.select(enriched)
        if selected_judges:
            logger.debug(
                f"[judge_hook] selected judges for span {ctx.span_id}: "
                + ", ".join(j.judge_id for j in selected_judges)
            )
        try:
            return await asyncio.wait_for(
                self._judge.evaluate_span(
                    span_id=ctx.span_id or "",
                    input_messages=ctx.messages,
                    output_content=ctx.response_content or "",
                    context=ctx,
                ),
                timeout=JUDGE_TIMEOUT_MS / 1000.0,
            )
        except asyncio.TimeoutError:
            logger.warning(f"[judge_hook] judge timeout for span {ctx.span_id}")
            return None
        except Exception as e:
            logger.warning(f"[judge_hook] judge error for span {ctx.span_id}: {e}")
            return None

    async def _post_to_backend(
        self,
        ctx: InterceptionContext,
        evaluation: "SpanEvaluation",
        original_output: str,
        fixed_output: str | None,
        action_type: str,
    ) -> None:
        """POST eval score to backend. Non-blocking — failures are logged, never raised."""
        if not self._backend_url or not self._api_key:
            return

        result = getattr(evaluation, "result", None)
        payload = {
            "spanId": ctx.span_id,
            "traceId": ctx.trace_id,
            "name": "judge",
            "value": result.overall_score if result else None,
            "isRuntime": True,
            "actionType": action_type,
            "rootCause": None,  # Not available from base LLMJudge
            "originalOutput": original_output,
            "fixedOutput": fixed_output,
            "judgeModel": evaluation.judge_model,
            "latencyMs": evaluation.evaluation_time_ms,
            "confidence": evaluation.confidence,
            "reasoning": result.reasoning if result else None,
        }
        try:
            async with httpx.AsyncClient(timeout=5.0) as client:
                await client.post(
                    f"{self._backend_url}/api/v1/evals/scores",
                    json=payload,
                    headers={"X-API-Key": self._api_key},
                )
        except Exception as e:
            logger.warning(f"[judge_hook] backend POST failed: {e}")
