"""Judge backend abstraction for LLM-based evaluators.

Defines :class:`JudgeBackend` and the GitHub Copilot SDK implementation. Azure
mode has no backend object — it calls the Azure evaluators inline (see
``evaluation_runner``), so a ``None`` backend means "use the inline Azure path".
"""

import asyncio
import json
import re
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

from cli_logging.cli_logger import emit_structured_log
from cli_logging.logging_utils import Operation
from common import METRIC_IDS
from error_messages import (
    JudgeError,
    EVAL_ERROR_CODE_JUDGE_RATE_LIMITED,
    EVAL_ERROR_CODE_JUDGE_AUTH,
    EVAL_ERROR_CODE_JUDGE_TIMEOUT,
)
from retry_policy import is_timeout_error

# Timeout (seconds) for a single Copilot SDK call dispatched to the background loop.
_SDK_CALL_TIMEOUT_SECONDS = 120

# Load prompt templates once at module level
_PROMPTS_PATH = Path(__file__).parent / "judge_prompts.json"
_PROMPT_TEMPLATES: Dict[str, Dict[str, str]] = {}


def _load_prompt_templates() -> Dict[str, Dict[str, str]]:
    """Load judge prompt templates from JSON file (cached)."""
    global _PROMPT_TEMPLATES
    if not _PROMPT_TEMPLATES:
        with open(_PROMPTS_PATH, "r", encoding="utf-8") as f:
            _PROMPT_TEMPLATES = json.load(f)
    return _PROMPT_TEMPLATES


def _parse_score_from_response(text: str, metric_id: str) -> Optional[Dict[str, Any]]:
    """Parse the JSON score from the LLM judge response.

    Tries JSON first (handles code fences, preamble, nested braces), then falls
    back to regex extraction. The fallback path tags the dict with
    ``"_fallback": True`` (callers strip it) so a salvaged score is
    distinguishable from a genuine parse.
    """
    # Strip markdown code fences if present (```json ... ``` or ``` ... ```)
    cleaned = re.sub(r'```(?:json)?\s*', '', text)
    cleaned = cleaned.strip()

    # Try parsing as JSON — find the first { and last } to handle nested braces
    try:
        start = cleaned.find('{')
        end = cleaned.rfind('}')
        if start != -1 and end != -1 and end > start:
            candidate = cleaned[start:end + 1]
            parsed = json.loads(candidate)
            if metric_id in parsed:
                return parsed
    except (json.JSONDecodeError, ValueError):
        pass

    # Second attempt: try each { as a potential start of JSON
    try:
        for match in re.finditer(r'\{', cleaned):
            try:
                parsed = json.loads(cleaned[match.start():])
                if metric_id in parsed:
                    return parsed
            except (json.JSONDecodeError, ValueError):
                # Try to find the matching closing brace
                try:
                    end = cleaned.index('}', match.start()) + 1
                    parsed = json.loads(cleaned[match.start():end])
                    if metric_id in parsed:
                        return parsed
                except (json.JSONDecodeError, ValueError, IndexError):
                    continue
    except Exception:
        pass

    # Fallback: extract just the score number
    score_match = re.search(r'\b([1-5])\b', text)
    if score_match:
        score = int(score_match.group(1))
        return {
            metric_id: float(score),
            f"gpt_{metric_id}": float(score),
            f"{metric_id}_reason": "Score extracted from non-JSON response.",
            "_fallback": True,
        }

    return None


class JudgeBackend(ABC):
    """Abstract base class for LLM judge backends.

    Implementations must provide an `evaluate` method that scores a
    prompt/response pair for a given evaluator name.
    """

    @abstractmethod
    def evaluate(
        self,
        eval_name: str,
        prompt: str,
        response: str,
        context: Optional[str] = None,
    ) -> Optional[Dict[str, Any]]:
        """Evaluate a prompt/response pair.

        Args:
            eval_name: Canonical evaluator name (e.g., "Relevance").
            prompt: The original user prompt.
            response: The agent's response text.
            context: Ground-truth or expected response.

        Returns:
            Raw evaluator output dict (e.g., {"relevance": 4}), or ``None`` when
            the judge produced genuinely unparseable / no-score output.

        Raises:
            JudgeError: When an implementation can positively identify a hard
                failure (e.g. authentication or rate-limiting), so the caller
                records the precise ``judge*`` code instead of inferring
                ``judgeOutputUnparseable`` from a bare ``None``. Other
                unexpected exceptions propagate for the caller to classify.
        """


class CopilotSDKJudgeBackend(JudgeBackend):
    """Judge backend using the GitHub Copilot SDK.

    Uses the `github-copilot-sdk` Python package for authentication and
    model access. Supports automatic auth via:
      - Signed-in GitHub CLI user
      - COPILOT_GITHUB_TOKEN / GH_TOKEN / GITHUB_TOKEN env vars
      - OAuth GitHub App tokens

    Environment variables:
        GITHUB_COPILOT_JUDGE_MODEL: Model to use (default: "auto")
        GITHUB_COPILOT_JUDGE_AUTO_FALLBACK: Retry with model="auto" on rate
            limit (default: "true"). Set to "false" to disable.
    """

    DEFAULT_MODEL = "auto"

    def __init__(self, log_level: str = "info"):
        import os
        import threading
        # Model from GITHUB_COPILOT_JUDGE_MODEL, or "auto" (Copilot picks per request).
        env_model = os.environ.get("GITHUB_COPILOT_JUDGE_MODEL")
        self._model = env_model or self.DEFAULT_MODEL
        self._active_model = self._model
        fallback_env = os.environ.get("GITHUB_COPILOT_JUDGE_AUTO_FALLBACK", "true").lower()
        self._auto_fallback = fallback_env not in ("false", "0", "no")
        # Mirror CLI verbosity: the runtime is chatty at "info", so only surface
        # its logs on --log-level debug.
        self._runtime_log_level = "debug" if log_level == "debug" else "error"
        self._templates = _load_prompt_templates()
        self._client = None
        self._init_failed = False
        self._rate_limited = False
        # Concrete models the service used (meaningful for "auto"); filled in evaluate().
        self._resolved_models = set()
        self._lock = threading.Lock()
        # Dedicated background event loop avoids "event loop already running"
        # when called from async code.
        self._loop = asyncio.new_event_loop()
        self._thread = threading.Thread(target=self._loop.run_forever, daemon=True)
        self._thread.start()

        if env_model:
            model_note = f"model={self._model} (from GITHUB_COPILOT_JUDGE_MODEL)"
        else:
            model_note = (
                f"model={self._model} (default — Copilot selects a model per request; "
                "set GITHUB_COPILOT_JUDGE_MODEL to pin one)"
            )
        emit_structured_log(
            "info",
            f"🤖 Copilot SDK Judge initialized: {model_note}, auto_fallback={self._auto_fallback} "
            "(uses GitHub auth — separate from M365 A2A auth)",
            operation=Operation.EVALUATE,
        )
        emit_structured_log(
            "warning",
            "⚠️  The Copilot SDK judge backend is experimental. "
            "Custom evaluators (user-authored .prompty + .py) are not yet supported "
            "and will be skipped. Only built-in LLM evaluators (Relevance, Coherence, "
            "Groundedness, Similarity) are routed through the Copilot SDK.",
            operation=Operation.EVALUATE,
        )

    def _record_resolved_model(self, model: Optional[str]) -> None:
        """Record a concrete model the service reported using (thread-safe)."""
        if model:
            with self._lock:
                self._resolved_models.add(model)

    def get_resolved_models(self) -> list:
        """Return the sorted set of concrete models the service actually used."""
        with self._lock:
            return sorted(self._resolved_models)

    def describe(self) -> str:
        """Judge description for the run's output metadata — configured model
        plus, when it differs (e.g. "auto"), the model(s) actually resolved to."""
        resolved = self.get_resolved_models()
        if resolved and (self._model == "auto" or set(resolved) != {self._model}):
            return f"GitHub Copilot (model: {self._model} → resolved: {', '.join(resolved)})"
        return f"GitHub Copilot (model: {self._model})"

    def _run(self, coro):
        """Submit a coroutine to the background event loop and wait for the result."""
        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
        return future.result(timeout=_SDK_CALL_TIMEOUT_SECONDS)

    def _ensure_client(self):
        """Lazily create the Copilot SDK client (not session).

        Thread-safe: uses a lock so only one worker initializes the client.
        Raises :class:`~error_messages.JudgeError` (code
        ``judgeAuthenticationError``) with actionable guidance if
        authentication fails.
        """
        if self._client is not None:
            return
        with self._lock:
            # Double-check after acquiring lock
            if self._client is not None:
                return
            if self._init_failed:
                raise JudgeError(
                    "Copilot SDK judge session previously failed to initialize (GitHub auth required). "
                    "Run 'gh auth login' or set GITHUB_TOKEN env var.",
                    code=EVAL_ERROR_CODE_JUDGE_AUTH,
                )

            from copilot import CopilotClient

            try:
                self._client = CopilotClient(log_level=self._runtime_log_level)
                self._run(self._client.start())
            except Exception as e:
                self._init_failed = True
                self._client = None
                emit_structured_log(
                    "error",
                    f"Failed to initialize Copilot SDK session: {e}\n"
                    "GitHub authentication is required for the Copilot SDK judge backend.\n"
                    "(This is separate from M365 Copilot A2A authentication)\n"
                    "Options:\n"
                    "  1. Run 'gh auth login' to sign in via GitHub CLI\n"
                    "  2. Set GITHUB_TOKEN or COPILOT_GITHUB_TOKEN environment variable\n"
                    "  3. Set GH_TOKEN environment variable",
                    operation=Operation.EVALUATE,
                )
                raise JudgeError(
                    "Copilot SDK judge authentication failed (GitHub auth required). "
                    "Run 'gh auth login' or set GITHUB_TOKEN env var.",
                    code=EVAL_ERROR_CODE_JUDGE_AUTH,
                ) from e

    def verify_model(self) -> None:
        """Fail-fast check that the configured model is available.

        No-op for ``"auto"``. Otherwise raises :class:`ValueError` (listing the
        available models) if ``GITHUB_COPILOT_JUDGE_MODEL`` isn't one the account can
        use, turning a per-prompt failure into one clear startup error.
        """
        if self._active_model == "auto":
            return

        self._ensure_client()
        models = self._run(self._client.list_models())
        available = sorted(
            getattr(m, "id", None) for m in models if getattr(m, "id", None)
        )
        if self._active_model not in available:
            options = ", ".join(available) if available else "(none reported)"
            raise ValueError(
                f"Copilot judge model '{self._active_model}' is not available for "
                f"your account. Set GITHUB_COPILOT_JUDGE_MODEL to one of: {options} "
                "(or 'auto' to let Copilot choose)."
            )

    def _create_session(self):
        """Create a fresh session for a single prompt (no history bleed between prompts)."""
        from copilot.session import PermissionHandler

        session = self._run(self._client.create_session(
            model=self._active_model,
            on_permission_request=PermissionHandler.approve_all,
            infinite_sessions={"enabled": False},
        ))
        return session

    def evaluate(
        self,
        eval_name: str,
        prompt: str,
        response: str,
        context: Optional[str] = None,
    ) -> Optional[Dict[str, Any]]:
        """Evaluate using the Copilot SDK.

        Builds the evaluation prompt from templates, sends it via the SDK,
        and parses the JSON score from the response.

        Returns ``None`` only for genuinely unparseable / no-score output (or a
        missing template). Positively-identified hard failures raise
        :class:`~error_messages.JudgeError` with a precise ``judge*`` code
        (``judgeRateLimited`` / ``judgeAuthenticationError`` / ``judgeTimeout``);
        any other SDK exception is re-raised unchanged so the caller's classifier
        can categorize it.
        """
        template = self._templates.get(eval_name)
        if not template:
            emit_structured_log(
                "warning",
                f"No prompt template for evaluator '{eval_name}'",
                operation=Operation.EVALUATE,
            )
            return None

        metric_id = METRIC_IDS.get(eval_name, eval_name.lower())

        # Build the scoring prompt
        user_message = template["user_template"]
        user_message = user_message.replace("$prompt", prompt)
        user_message = user_message.replace("$response", response)
        user_message = user_message.replace("$context", context or "")
        system_message = template["system"]

        scoring_prompt = f"{system_message}\n\n{user_message}"

        emit_structured_log(
            "info",
            f"📡 Evaluating '{eval_name}' via Copilot SDK ({self._active_model})...",
            operation=Operation.EVALUATE,
        )

        try:
            self._ensure_client()
            sent = self._send_prompt(eval_name, scoring_prompt)
            raw_text, resolved_model = sent if sent else (None, None)
            self._record_resolved_model(resolved_model)
            model_label = resolved_model or self._active_model

            if raw_text:
                result = _parse_score_from_response(raw_text, metric_id)
                if result is not None:
                    is_fallback = result.pop("_fallback", False)
                    if not is_fallback:
                        emit_structured_log(
                            "debug",
                            f"✅ '{eval_name}' score from Copilot SDK (model={model_label}): {result.get(metric_id)}",
                            operation=Operation.EVALUATE,
                        )
                        return result

                    # Salvaged a bare score from non-JSON — log raw output, still return it.
                    emit_structured_log(
                        "warning",
                        f"⚠️  Non-JSON response for '{eval_name}'. Raw output:\n---\n{raw_text}\n---",
                        operation=Operation.EVALUATE,
                    )
                    return result

            emit_structured_log(
                "warning",
                f"Could not parse score from Copilot SDK response for '{eval_name}'",
                operation=Operation.EVALUATE,
            )
        except Exception as e:
            if self._is_rate_limit_error(e):
                if self._auto_fallback and self._active_model != "auto":
                    # User opted in to auto-fallback
                    self._switch_to_auto_model()
                    try:
                        sent = self._send_prompt(eval_name, scoring_prompt)
                        raw_text, resolved_model = sent if sent else (None, None)
                        self._record_resolved_model(resolved_model)
                        if raw_text:
                            result = _parse_score_from_response(raw_text, metric_id)
                            if result is not None:
                                result.pop("_fallback", None)
                                emit_structured_log(
                                    "debug",
                                    f"✅ '{eval_name}' score from Copilot SDK "
                                    f"(auto fallback, model={resolved_model or self._active_model}): {result.get(metric_id)}",
                                    operation=Operation.EVALUATE,
                                )
                                return result
                    except Exception as retry_err:
                        emit_structured_log(
                            "error",
                            f"Copilot SDK judge call failed for '{eval_name}' after auto-model fallback: {retry_err}",
                            operation=Operation.EVALUATE,
                        )
                        raise JudgeError(
                            f"Copilot SDK judge call for '{eval_name}' was rate-limited; "
                            f"the auto-model fallback also failed: {retry_err}",
                            code=EVAL_ERROR_CODE_JUDGE_RATE_LIMITED,
                        ) from retry_err
                    # Fallback ran but produced no parseable score — still a
                    # rate-limit-induced failure, not unparseable output.
                    raise JudgeError(
                        f"Copilot SDK judge call for '{eval_name}' was rate-limited; "
                        f"the auto-model fallback produced no parseable score.",
                        code=EVAL_ERROR_CODE_JUDGE_RATE_LIMITED,
                    ) from e
                else:
                    self._rate_limited = True
                    emit_structured_log(
                        "error",
                        f"🚫 Rate limit hit for '{eval_name}'. "
                        f"Model '{self._active_model}' quota exhausted. "
                        f"Wait for your limit to reset, set GITHUB_COPILOT_JUDGE_MODEL to a different model, "
                        f"or set GITHUB_COPILOT_JUDGE_AUTO_FALLBACK=true to automatically retry with model='auto'.",
                        operation=Operation.EVALUATE,
                    )
                    raise JudgeError(
                        f"Copilot SDK judge call for '{eval_name}' was rate-limited "
                        f"(model '{self._active_model}' quota exhausted).",
                        code=EVAL_ERROR_CODE_JUDGE_RATE_LIMITED,
                    ) from e
            else:
                emit_structured_log(
                    "error",
                    f"Copilot SDK judge call failed for '{eval_name}': {e}",
                    operation=Operation.EVALUATE,
                )
                # Timeouts are positively identifiable from the exception type
                # (walking the cause chain), so classify at the source with a
                # typed JudgeError rather than leaving it to the caller's
                # structural fallback.
                if is_timeout_error(e):
                    raise JudgeError(
                        f"Copilot SDK judge call for '{eval_name}' timed out.",
                        code=EVAL_ERROR_CODE_JUDGE_TIMEOUT,
                    ) from e
                # Otherwise re-raise the original SDK exception so the caller's
                # layered classifier can categorize it (unparseable / generic).
                raise

        return None

    def _send_prompt(self, eval_name: str, prompt_text: str) -> Optional[Tuple[str, Optional[str]]]:
        """Send a scoring prompt via a fresh session (one per prompt).

        Returns ``(content, resolved_model)`` or ``None``. ``resolved_model`` is
        the model the service actually used (``AssistantMessageData.model``, may
        be ``None``) — useful when ``"auto"`` is in effect.
        """
        session = self._create_session()
        try:
            response = self._run(session.send_and_wait(prompt_text))
            if response and hasattr(response, 'data') and hasattr(response.data, 'content'):
                resolved_model = getattr(response.data, 'model', None)
                return response.data.content, resolved_model
            return None
        finally:
            try:
                self._run(session.disconnect())
            except Exception:
                pass

    def _is_rate_limit_error(self, error: Exception) -> bool:
        """Check if an exception is a rate limit error."""
        error_str = str(error).lower()
        return "rate limit" in error_str or "rate_limit" in error_str

    def _switch_to_auto_model(self):
        """Switch to 'auto' model after hitting rate limits on the primary model."""
        with self._lock:
            if self._rate_limited:
                return  # Already switched
            self._rate_limited = True
            self._active_model = "auto"

            emit_structured_log(
                "warning",
                f"⚠️  Rate limit hit on model '{self._model}'. "
                f"Switching to 'auto' model for remaining evaluations.",
                operation=Operation.EVALUATE,
            )

    def close(self):
        """Clean up the Copilot SDK client and background loop."""
        if self._client:
            try:
                self._run(self._client.stop())
            except Exception:
                pass
            self._client = None
        if self._loop and self._loop.is_running():
            self._loop.call_soon_threadsafe(self._loop.stop)
            self._thread.join(timeout=5)
            self._loop.close()
            self._loop = None