"""
L2 Pro Premium Judge — stronger model, more dimensions, Pro-tier only.

Why a SEPARATE module from L1?
  - L1 is fast + cheap (deepseek-chat, 4 dimensions, 5 samples)
  - L2 is slow + expensive (gpt-4o, 8 dimensions, 15 samples)
  - Different defaults, different config, different metrics
  - L2 is gated to the Pro plan; L1 is universal

L2 dimensions (8):
  1. accurate       — meaning preserved
  2. fluent         — natural in target language
  3. correct_lang   — in the target language (not source leakage)
  4. no_leaks       — no prompt artifacts
  5. terminology    — domain terms correctly handled
  6. style          — appropriate register (formal/informal/technical)
  7. completeness   — no content dropped or added
  8. formatting     — codes, numbers, units preserved

L1 was a binary pass/fail. L2 returns per-dimension scores (0/1) so the
caller can decide which dimensions matter for a given job.

DESIGN CONSTRAINTS (same as L1):
  - 100% API-based. No local models, no GPU.
  - Async, with a hard timeout.
  - Defensive: never raises.
  - Output is structured JSON.
"""

from __future__ import annotations

import asyncio
import json
import re
import time
from dataclasses import dataclass, field, asdict
from typing import List, Optional, Tuple, Dict, Any

from core.logging import get_logger

logger = get_logger(__name__)


# ---------- Result dataclasses ----------

@dataclass
class L2DimensionVerdict:
    """8-dimension verdict for a single chunk."""
    accurate: bool = False
    fluent: bool = False
    correct_lang: bool = False
    no_leaks: bool = False
    terminology: bool = False
    style: bool = False
    completeness: bool = False
    formatting: bool = False
    reason: str = ""

    @property
    def passed_count(self) -> int:
        return sum([
            self.accurate, self.fluent, self.correct_lang, self.no_leaks,
            self.terminology, self.style, self.completeness, self.formatting,
        ])

    @property
    def total(self) -> int:
        return 8

    @property
    def score(self) -> float:
        return self.passed_count / self.total

    @property
    def passed(self) -> bool:
        # L2 is conservative: any single fail = chunk fails
        return self.passed_count == self.total

    def to_log_dict(self) -> dict:
        return asdict(self)


@dataclass
class L2Result:
    """Aggregate result of an L2 check on a sample of chunks."""
    verdict: str                  # "pass", "fail", "skip"
    chunks_evaluated: int = 0
    chunks_passed: int = 0
    chunks_failed: int = 0
    failure_rate: float = 0.0
    average_score: float = 0.0    # mean of per-chunk scores (0.0 to 1.0)
    dimension_pass_rates: Dict[str, float] = field(default_factory=dict)
    samples: List[dict] = field(default_factory=list)
    model_used: str = ""
    elapsed_ms: float = 0.0
    cost_estimate_usd: float = 0.0
    error: str = ""

    def to_log_dict(self) -> dict:
        return asdict(self)


# ---------- Prompt template (8 dimensions) ----------

L2_JUDGE_SYSTEM_PROMPT = """You are an expert translation quality evaluator using MQM-inspired criteria.

For each (SOURCE, TRANSLATION) pair, check these 8 criteria (yes/no for each):

1. ACCURATE     — Does the translation preserve the meaning of the source?
2. FLUENT       — Is the translation natural and grammatical in {target_lang_name}?
3. CORRECT_LANG — Is the translation actually in {target_lang_name} (ISO: {target_lang})?
4. NO_LEAKS     — Is the translation free of prompt artifacts, source-language text, or meta-commentary?
5. TERMINOLOGY  — Are domain-specific terms (technical, legal, medical, etc.) correctly translated?
6. STYLE        — Is the register/tone appropriate (formal/informal/technical matching the source)?
7. COMPLETENESS — Is all content present, with nothing added or dropped?
8. FORMATTING   — Are codes, numbers, dates, and units preserved exactly?

A translation FAILS if ANY criterion is "no". The "reason" must be in English and ≤ 20 words.

Respond with a JSON array, one object per pair, in the same order. NO other text, NO markdown fences:
[
  {{"accurate": "yes"|"no", "fluent": "yes"|"no", "correct_lang": "yes"|"no", "no_leaks": "yes"|"no", "terminology": "yes"|"no", "style": "yes"|"no", "completeness": "yes"|"no", "formatting": "yes"|"no", "reason": "short justification"}}
]
"""


# ---------- LLM client ----------

class L2ProJudge:
    """
    Calls a STRONG LLM via the OpenAI-compatible API to judge translation
    quality across 8 dimensions. Pro-tier only.

    Default model: gpt-4o (strongest general judge we can afford at scale).
    For sub-$0.01/job cost, we limit to 15 samples per job and 8 dimensions.
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.openai.com/v1",
        model: str = "gpt-4o",
        timeout_seconds: float = 20.0,
        max_retries: int = 1,
    ):
        if not api_key:
            raise ValueError("api_key is required for L2ProJudge")
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._model = model
        self._timeout = timeout_seconds
        self._max_retries = max_retries
        self._client = None

    def _get_client(self):
        if self._client is None:
            try:
                import openai
                self._client = openai.AsyncOpenAI(
                    api_key=self._api_key,
                    base_url=self._base_url,
                    timeout=self._timeout,
                )
            except Exception as e:
                logger.warning("l2_judge_client_init_failed", error=str(e))
                return None
        return self._client

    async def judge_batch(
        self,
        pairs: List[Tuple[str, str]],
        target_lang: str,
        target_lang_name: str = "",
    ) -> L2Result:
        """
        Judge a batch of (source, translation) pairs across 8 dimensions.

        Returns an L2Result with verdict="skip" on any internal error.
        Never raises.
        """
        start = time.time()
        empty = L2Result(verdict="skip", error="not_run")

        if not pairs:
            return L2Result(verdict="skip", error="empty pairs",
                            elapsed_ms=round((time.time() - start) * 1000, 2))

        client = self._get_client()
        if client is None:
            return L2Result(verdict="skip", error="client unavailable",
                            elapsed_ms=round((time.time() - start) * 1000, 2))

        user_lines = [f"Target language: {target_lang} ({target_lang_name})\n"]
        for i, (src, trans) in enumerate(pairs, 1):
            user_lines.append(f"\n--- Pair {i} ---")
            user_lines.append(f"SOURCE: {src}")
            user_lines.append(f"TRANSLATION: {trans}")
        user_msg = "\n".join(user_lines)

        system_prompt = L2_JUDGE_SYSTEM_PROMPT.format(
            target_lang=target_lang,
            target_lang_name=target_lang_name or target_lang,
        )

        try:
            response = await asyncio.wait_for(
                self._call_with_retries(client, system_prompt, user_msg),
                timeout=self._timeout + 5.0,
            )
        except asyncio.TimeoutError:
            elapsed_ms = round((time.time() - start) * 1000, 2)
            logger.warning("l2_judge_timeout",
                           timeout_s=self._timeout, elapsed_ms=elapsed_ms)
            return L2Result(verdict="skip", error="timeout", elapsed_ms=elapsed_ms)
        except Exception as e:
            elapsed_ms = round((time.time() - start) * 1000, 2)
            logger.warning("l2_judge_error",
                           error=str(e)[:200], elapsed_ms=elapsed_ms)
            return L2Result(verdict="skip", error=str(e)[:200],
                            elapsed_ms=elapsed_ms)

        verdicts = self._parse_response(response, len(pairs))

        if not verdicts:
            elapsed_ms = round((time.time() - start) * 1000, 2)
            return L2Result(verdict="skip", error="parse_failed",
                            elapsed_ms=elapsed_ms)

        # Aggregate
        passed = sum(1 for v in verdicts if v.passed)
        failed = len(verdicts) - passed
        failure_rate = failed / len(verdicts) if verdicts else 0.0
        average_score = sum(v.score for v in verdicts) / len(verdicts)

        # Per-dimension pass rates
        dimensions = [
            "accurate", "fluent", "correct_lang", "no_leaks",
            "terminology", "style", "completeness", "formatting",
        ]
        dim_pass_rates = {}
        for dim in dimensions:
            count = sum(1 for v in verdicts if getattr(v, dim))
            dim_pass_rates[dim] = round(count / len(verdicts), 3) if verdicts else 0.0

        # L2 verdict: strict — any chunk fail = overall fail
        verdict = "pass" if failed == 0 else "fail"

        elapsed_ms = round((time.time() - start) * 1000, 2)
        cost_estimate = self._estimate_cost(len(pairs))

        return L2Result(
            verdict=verdict,
            chunks_evaluated=len(verdicts),
            chunks_passed=passed,
            chunks_failed=failed,
            failure_rate=round(failure_rate, 3),
            average_score=round(average_score, 3),
            dimension_pass_rates=dim_pass_rates,
            samples=[v.to_log_dict() for v in verdicts],
            model_used=self._model,
            elapsed_ms=elapsed_ms,
            cost_estimate_usd=cost_estimate,
        )

    async def _call_with_retries(self, client, system_prompt: str, user_msg: str):
        """Call the LLM with retry on transient errors."""
        last_exc = None
        for attempt in range(self._max_retries + 1):
            try:
                response = await client.chat.completions.create(
                    model=self._model,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_msg},
                    ],
                    temperature=0.0,
                    max_tokens=1200,          # larger than L1 (more dimensions)
                    response_format={"type": "json_object"},
                )
                return response
            except Exception as e:
                last_exc = e
                if attempt < self._max_retries:
                    await asyncio.sleep(0.8)
        raise last_exc

    def _parse_response(self, response, expected_count: int) -> List[L2DimensionVerdict]:
        """Parse the LLM response into a list of 8-dimension verdicts."""
        try:
            content = response.choices[0].message.content or ""
        except (AttributeError, IndexError) as e:
            logger.warning("l2_judge_bad_response", error=str(e))
            return []

        content = content.strip()
        if content.startswith("```"):
            content = re.sub(r"^```(?:json)?\s*\n?", "", content)
            content = re.sub(r"\n?```\s*$", "", content)

        try:
            data = json.loads(content)
        except json.JSONDecodeError as e:
            logger.warning("l2_judge_json_parse_error",
                           error=str(e), content_preview=content[:200])
            return []

        if isinstance(data, dict):
            items = None
            for key in ("verdicts", "results", "translations", "data"):
                if key in data and isinstance(data[key], list):
                    items = data[key]
                    break
            if items is None:
                for v in data.values():
                    if isinstance(v, list):
                        items = v
                        break
            if items is None:
                logger.warning("l2_judge_no_list_in_response")
                return []
        elif isinstance(data, list):
            items = data
        else:
            logger.warning("l2_judge_unexpected_response_type",
                           type_=type(data).__name__)
            return []

        verdicts: List[L2DimensionVerdict] = []
        for item in items:
            try:
                v = L2DimensionVerdict(
                    accurate=str(item.get("accurate", "")).lower() == "yes",
                    fluent=str(item.get("fluent", "")).lower() == "yes",
                    correct_lang=str(item.get("correct_lang", "")).lower() == "yes",
                    no_leaks=str(item.get("no_leaks", "")).lower() == "yes",
                    terminology=str(item.get("terminology", "")).lower() == "yes",
                    style=str(item.get("style", "")).lower() == "yes",
                    completeness=str(item.get("completeness", "")).lower() == "yes",
                    formatting=str(item.get("formatting", "")).lower() == "yes",
                    reason=str(item.get("reason", ""))[:300],
                )
                verdicts.append(v)
            except Exception as e:
                logger.warning("l2_judge_item_parse_error",
                               error=str(e), item=str(item)[:200])

        return verdicts

    def _estimate_cost(self, num_pairs: int) -> float:
        """Rough USD cost estimate for the call."""
        # L2 has more dimensions = longer output
        input_tokens = 250 + (num_pairs * 280)
        output_tokens = num_pairs * 110
        model_lower = self._model.lower()
        # IMPORTANT: check 'mini' BEFORE full 'gpt-4o' because
        # 'gpt-4o-mini' contains 'gpt-4o'.
        if "gpt-4o-mini" in model_lower:
            input_cost = input_tokens / 1_000_000 * 0.15
            output_cost = output_tokens / 1_000_000 * 0.60
        elif "gpt-4o" in model_lower:
            input_cost = input_tokens / 1_000_000 * 2.50
            output_cost = output_tokens / 1_000_000 * 10.00
        elif "claude" in model_lower:
            input_cost = input_tokens / 1_000_000 * 3.00
            output_cost = output_tokens / 1_000_000 * 15.00
        else:
            # Generic conservative
            input_cost = input_tokens / 1_000_000 * 1.00
            output_cost = output_tokens / 1_000_000 * 3.00
        return round(input_cost + output_cost, 6)


# ---------- Convenience factory ----------

def make_l2_judge_from_env() -> Optional[L2ProJudge]:
    """
    Build an L2ProJudge from environment variables. Returns None if
    no API key is configured.

    Reads:
      - L2_JUDGE_API_KEY       (required)
      - L2_JUDGE_BASE_URL      (default: OpenAI)
      - L2_JUDGE_MODEL         (default: gpt-4o)
      - L2_JUDGE_TIMEOUT       (default: 20.0)
    """
    import os
    api_key = os.getenv("L2_JUDGE_API_KEY", "").strip()
    if not api_key:
        return None
    return L2ProJudge(
        api_key=api_key,
        base_url=os.getenv("L2_JUDGE_BASE_URL", "https://api.openai.com/v1"),
        model=os.getenv("L2_JUDGE_MODEL", "gpt-4o"),
        timeout_seconds=float(os.getenv("L2_JUDGE_TIMEOUT", "20.0")),
    )
