"""
Async Judge Runner — Tier 2 parallel evaluation (500-3000ms).

Fires LLM judges in background via asyncio.create_task(). The agent
NEVER waits for results. When a judge completes and finds an issue,
the fix is queued in InterventionDispatcher for the next step's pre-call.

Designed for high scale: each judge task is independent, uses bounded
concurrency, and never blocks the agent's hot path.
"""

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from ..judge.context_aggregator import ContextAggregator
    from ..judge.evaluator import LLMJudge
    from ..judge.selector import JudgeSelector
    from ..runtime.remediation_loop import RemediationLoop

logger = logging.getLogger("aigie.runtime.async_judge")


@dataclass
class PendingFix:
    """A fix queued for the next step's pre-call hook."""

    trace_id: str
    span_id: str
    fix_type: str  # "corrective_context", "retry_instruction", "response_modification"
    fix_content: str  # The corrective text or instruction
    judge_category: str  # Which judge detected the issue
    judge_score: float  # Score that triggered the fix
    confidence: float
    created_at: float = field(default_factory=time.monotonic)


class AsyncJudgeRunner:
    """Runs LLM judges in background, queues fixes for next step.

    Usage in the hot path:
        # After LLM response is captured and returned to agent:
        runner.fire_and_forget(trace_id, span_id, output, messages, metadata)
        # Agent continues immediately. Judge runs in background.

    On next step's pre-call:
        fix = runner.pop_pending_fix(trace_id)
        if fix:
            messages.insert(0, {"role": "system", "content": fix.fix_content})
    """

    def __init__(
        self,
        judge: Optional["LLMJudge"] = None,
        context_aggregator: Optional["ContextAggregator"] = None,
        judge_selector: Optional["JudgeSelector"] = None,
        remediation_loop: Optional["RemediationLoop"] = None,
        max_concurrent_judges: int = 4,
        score_threshold: float = 0.5,
        api_url: str | None = None,
        api_key: str | None = None,
    ):
        self._judge = judge
        self._context_aggregator = context_aggregator
        self._selector = judge_selector
        self._remediation_loop = remediation_loop
        self._max_concurrent = asyncio.Semaphore(max_concurrent_judges)
        self._score_threshold = score_threshold
        self._api_url = api_url
        self._api_key = api_key
        self._http_client = None  # Lazy persistent HTTP client

        # Pending fixes per trace (lock-free: dict ops are atomic in CPython)
        self._pending_fixes: Dict[str, List[PendingFix]] = {}

        # Active tasks for cleanup
        self._active_tasks: set = set()

        # Stats
        self._stats = {
            "judges_fired": 0,
            "judges_completed": 0,
            "judges_failed": 0,
            "fixes_queued": 0,
            "fixes_applied": 0,
        }

    def fire_and_forget(
        self,
        trace_id: str,
        span_id: str,
        span_type: str,
        output_content: str,
        input_messages: List[Dict[str, Any]],
        metadata: Dict[str, Any] | None = None,
        error: str | None = None,
        signal_hints: List[str] | None = None,
    ) -> None:
        """Fire async judge evaluation. Returns immediately.

        This is the ONLY method called from the hot path.
        Everything else runs in background tasks.

        Args:
            signal_hints: Lightweight tags from FastDetector (e.g., ["refusal",
                "hallucination_marker"]) that help JudgeSelector pick the right
                platform judges. If None, selection falls back to span_type only.
        """
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            logger.debug("No event loop — skipping async judge")
            return

        task = loop.create_task(
            self._run_evaluation(
                trace_id,
                span_id,
                span_type,
                output_content,
                input_messages,
                metadata or {},
                error,
                signal_hints or [],
            )
        )
        self._active_tasks.add(task)
        task.add_done_callback(self._active_tasks.discard)

    async def _run_evaluation(
        self,
        trace_id: str,
        span_id: str,
        span_type: str,
        output_content: str,
        input_messages: List[Dict[str, Any]],
        metadata: Dict[str, Any],
        error: str | None,
        signal_hints: List[str] | None = None,
    ) -> None:
        """Background evaluation task. Agent does NOT wait for this."""
        try:
            self._stats["judges_fired"] += 1

            # 1. Feed context aggregator
            if self._context_aggregator:
                self._context_aggregator.complete_span(
                    span_id=span_id,
                    output_content=output_content,
                    input_tokens=metadata.get("prompt_tokens", 0),
                    output_tokens=metadata.get("completion_tokens", 0),
                    cost=metadata.get("total_cost", 0.0),
                    status="error" if error else "success",
                    error=error,
                )

            # 2. Select applicable judges — use signal-aware selection if available
            judge_categories: List[str] = []
            if self._selector:
                if signal_hints and hasattr(self._selector, "select_with_signals"):
                    judge_categories = self._selector.select_with_signals(
                        span_type,
                        signal_hints=signal_hints,
                        metadata=metadata,
                    )
                    # Filter to LLM-only judges (exclude heuristic judges)
                    if hasattr(self._selector, "select_heuristic_only"):
                        heuristic_set = set(self._selector.select_heuristic_only(span_type))
                        judge_categories = [j for j in judge_categories if j not in heuristic_set]
                else:
                    judge_categories = self._selector.select_llm_only(span_type)
            elif self._judge:
                judge_categories = ["quality"]  # Default

            if not judge_categories:
                return

            # 3. Run judges with bounded concurrency
            async with self._max_concurrent:
                evaluation = await self._evaluate_span(
                    trace_id,
                    span_id,
                    output_content,
                    input_messages,
                    metadata,
                    judge_categories,
                )

            if not evaluation:
                self._stats["judges_completed"] += 1
                return

            # 4. Check if fix needed
            score = evaluation.get("score", 1.0)
            if score < self._score_threshold:
                fix = self._build_fix(
                    trace_id,
                    span_id,
                    evaluation,
                    judge_categories,
                )
                if fix:
                    self._pending_fixes.setdefault(trace_id, []).append(fix)
                    self._stats["fixes_queued"] += 1
                    logger.debug(
                        f"Async judge queued fix for trace={trace_id}: "
                        f"score={score:.2f}, type={fix.fix_type}"
                    )

            # 5. Feed remediation loop for learning
            if self._remediation_loop:
                try:
                    await self._remediation_loop.process_span(
                        span_id=span_id,
                        trace_id=trace_id,
                        input_messages=input_messages,
                        output_content=output_content,
                        error=Exception(error) if error else None,
                    )
                except Exception as e:
                    logger.debug(f"RemediationLoop.process_span failed: {e}")

            self._stats["judges_completed"] += 1

        except Exception as e:
            self._stats["judges_failed"] += 1
            logger.debug(f"Async judge evaluation failed (non-fatal): {e}")

    async def _evaluate_span(
        self,
        trace_id: str,
        span_id: str,
        output_content: str,
        input_messages: List[Dict[str, Any]],
        metadata: Dict[str, Any],
        judge_categories: List[str],
    ) -> Dict[str, Any] | None:
        """Run evaluation via LLMJudge or platform API."""
        # Try local LLMJudge first
        if self._judge:
            try:
                context = None
                if self._context_aggregator:
                    context = self._context_aggregator.get_or_create_context(trace_id)

                evaluation = await asyncio.wait_for(
                    self._judge.evaluate_span(
                        span_id=span_id,
                        input_messages=input_messages,
                        output_content=output_content,
                        context=context,
                    ),
                    timeout=5.0,
                )
                return {
                    "score": evaluation.score if hasattr(evaluation, "score") else 1.0,
                    "issues": evaluation.issues if hasattr(evaluation, "issues") else [],
                    "decision": str(evaluation.decision)
                    if hasattr(evaluation, "decision")
                    else "continue",
                    "categories": judge_categories,
                }
            except asyncio.TimeoutError:
                logger.debug("LLMJudge evaluation timed out")
            except Exception as e:
                logger.debug(f"LLMJudge evaluation failed: {e}")

        # Fallback: platform sync-evaluate API
        if self._api_url and self._api_key:
            try:
                return await self._evaluate_via_platform(
                    trace_id,
                    span_id,
                    output_content,
                    input_messages,
                    judge_categories,
                )
            except Exception as e:
                logger.debug(f"Platform evaluation failed: {e}")

        return None

    async def _get_http_client(self):
        """Lazy init persistent HTTP client."""
        if self._http_client is None:
            import httpx

            self._http_client = httpx.AsyncClient(
                timeout=5.0,
                headers={"X-API-Key": self._api_key, "Content-Type": "application/json"},
            )
        return self._http_client

    async def _evaluate_via_platform(
        self,
        trace_id: str,
        span_id: str,
        output_content: str,
        input_messages: List[Dict[str, Any]],
        judge_categories: List[str],
    ) -> Dict[str, Any] | None:
        """Call platform POST /evals/sync-evaluate."""
        client = await self._get_http_client()
        url = f"{self._api_url}/evals/sync-evaluate"
        resp = await client.post(
            url,
            json={
                "traceId": trace_id,
                "spanId": span_id,
                "output": output_content,
                "input": str(input_messages[-1].get("content", "")) if input_messages else "",
                "categories": judge_categories,
                "fastMode": True,
            },
        )
        if resp.status_code == 200:
            data = resp.json()
            return {
                "score": data.get("score", 1.0),
                "issues": data.get("issues", []),
                "decision": data.get("decision", "continue"),
                "categories": judge_categories,
            }
        return None

    def _build_fix(
        self,
        trace_id: str,
        span_id: str,
        evaluation: Dict[str, Any],
        categories: List[str],
    ) -> PendingFix | None:
        """Build a corrective fix from judge evaluation."""
        issues = evaluation.get("issues", [])
        score = evaluation.get("score", 1.0)
        category = categories[0] if categories else "unknown"

        if not issues:
            # Generic quality fix
            return PendingFix(
                trace_id=trace_id,
                span_id=span_id,
                fix_type="corrective_context",
                fix_content=(
                    "IMPORTANT: The quality of your previous response was below acceptable threshold. "
                    "Please ensure accuracy and completeness in your next response."
                ),
                judge_category=category,
                judge_score=score,
                confidence=0.6,
            )

        # Build specific fix based on issues
        issue_text = "; ".join(str(i) for i in issues[:3])
        return PendingFix(
            trace_id=trace_id,
            span_id=span_id,
            fix_type="corrective_context",
            fix_content=f"IMPORTANT: Issues detected in previous response: {issue_text}. Please address these in your next response.",
            judge_category=category,
            judge_score=score,
            confidence=0.7,
        )

    # ------------------------------------------------------------------
    # Public API for pre-call hook
    # ------------------------------------------------------------------

    def pop_pending_fix(self, trace_id: str) -> PendingFix | None:
        """Pop the highest-priority pending fix for a trace. Non-blocking."""
        fixes = self._pending_fixes.get(trace_id)
        if not fixes:
            return None
        # Pop first (oldest = should be applied first)
        fix = fixes.pop(0)
        if not fixes:
            del self._pending_fixes[trace_id]
        self._stats["fixes_applied"] += 1
        return fix

    def has_pending_fix(self, trace_id: str) -> bool:
        """Check if there are pending fixes for a trace."""
        return bool(self._pending_fixes.get(trace_id))

    def clear_trace(self, trace_id: str) -> None:
        """Clean up state for a completed trace."""
        self._pending_fixes.pop(trace_id, None)

    async def shutdown(self) -> None:
        """Wait for all active tasks to complete and close HTTP client."""
        if self._active_tasks:
            await asyncio.gather(*self._active_tasks, return_exceptions=True)
            self._active_tasks.clear()
        if self._http_client:
            await self._http_client.aclose()
            self._http_client = None

    def get_stats(self) -> Dict[str, Any]:
        return dict(self._stats)
