"""LLM-as-judge evaluation — scores model responses on quality dimensions."""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass

import aiohttp

logger = logging.getLogger("ecate.calibration.judge")

JUDGE_PROMPT = """You are evaluating an AI model's response to a task.

## Task
{task_prompt}

## Response
{model_response}

## Expected Behavior (reference)
{expected_behavior}

Rate the response on three dimensions (1-5 each):
1. **Correctness**: Is the response factually/technically correct?
2. **Completeness**: Does it fully address the task?
3. **Instruction Following**: Does it respect all constraints and format requirements?

Respond ONLY with a JSON object, no other text:
{{"correctness": N, "completeness": N, "following": N, "reasoning": "brief explanation"}}"""


@dataclass
class JudgeScore:
    correctness: float
    completeness: float
    following: float
    composite: float
    reasoning: str


async def judge_response(
    session: aiohttp.ClientSession,
    judge_base_url: str,
    judge_api_key: str,
    judge_model: str,
    judge_provider_type: str,
    task_prompt: str,
    model_response: str,
    expected_behavior: str,
) -> JudgeScore:
    """Use an LLM to judge a model's response quality.

    Returns a JudgeScore with correctness, completeness, following, and composite scores.
    """
    prompt = JUDGE_PROMPT.format(
        task_prompt=task_prompt,
        model_response=model_response[:3000],  # Truncate very long responses
        expected_behavior=expected_behavior or "No specific expected behavior provided.",
    )

    # Build request based on provider type
    if judge_provider_type == "anthropic":
        url = f"{judge_base_url.rstrip('/')}/v1/messages"
        headers = {
            "Content-Type": "application/json",
            "x-api-key": judge_api_key,
            "anthropic-version": "2023-06-01",
        }
        body = {
            "model": judge_model,
            "max_tokens": 500,
            "temperature": 0,
            "messages": [{"role": "user", "content": prompt}],
        }
    else:
        base = judge_base_url.rstrip("/")
        url = f"{base}/chat/completions" if "/v1" in base else f"{base}/v1/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {judge_api_key}",
        }
        body = {
            "model": judge_model,
            "temperature": 0,
            "max_tokens": 500,
            "messages": [{"role": "user", "content": prompt}],
        }

    try:
        async with session.post(url, json=body, headers=headers) as resp:
            if resp.status != 200:
                error = await resp.text()
                logger.error(f"Judge API error {resp.status}: {error[:300]}")
                return JudgeScore(3.0, 3.0, 3.0, 3.0, f"Judge error: {resp.status}")

            data = await resp.json()

            # Extract text from response
            if judge_provider_type == "anthropic":
                content = data.get("content", [{}])
                text = content[0].get("text", "") if content else ""
            else:
                text = data.get("choices", [{}])[0].get("message", {}).get("content", "")

            # Parse JSON from response
            # Try to extract JSON even if there's surrounding text
            json_start = text.find("{")
            json_end = text.rfind("}") + 1
            if json_start >= 0 and json_end > json_start:
                json_str = text[json_start:json_end]
                scores = json.loads(json_str)
            else:
                logger.warning(f"Judge response is not valid JSON: {text[:200]}")
                return JudgeScore(3.0, 3.0, 3.0, 3.0, "Failed to parse judge response")

            correctness = float(scores.get("correctness", 3))
            completeness = float(scores.get("completeness", 3))
            following = float(scores.get("following", 3))
            reasoning = scores.get("reasoning", "")

            # Composite: weighted average
            composite = correctness * 0.5 + completeness * 0.3 + following * 0.2

            return JudgeScore(
                correctness=correctness,
                completeness=completeness,
                following=following,
                composite=round(composite, 2),
                reasoning=reasoning,
            )

    except json.JSONDecodeError as e:
        logger.error(f"Failed to parse judge JSON: {e}")
        return JudgeScore(3.0, 3.0, 3.0, 3.0, f"JSON parse error: {e}")
    except Exception as e:
        logger.error(f"Judge evaluation failed: {e}")
        return JudgeScore(3.0, 3.0, 3.0, 3.0, f"Error: {e}")
