"""Layer 4: Discovery count and error prevention scoring."""

from __future__ import annotations

import re
from dataclasses import dataclass

from .llm_client import JudgeLLM
from .prompts import (
    DISCOVERY_SYSTEM_PROMPT,
    DISCOVERY_USER_PROMPT,
    ERROR_SYSTEM_PROMPT,
    ERROR_USER_PROMPT,
)


def _parse_score(response: str) -> int:
    m = re.search(r"Score:\s*(\d)", response)
    if m:
        return max(1, min(5, int(m.group(1))))
    m = re.search(r"\d", response)
    return max(1, min(5, int(m.group(0)))) if m else 3


def _parse_count(response: str) -> int:
    m = re.search(r"Count:\s*(\d+)", response)
    if m:
        return int(m.group(1))
    m = re.search(r"\d+", response)
    return int(m.group(0)) if m else 0


@dataclass
class Layer4Scores:
    discovery_count: float  # avg count of non-obvious insights
    error_prevention: float  # avg 1-5 score


async def evaluate_layer4(
    judge: JudgeLLM,
    output: str,
    input_text: str,
    repeat: int = 3,
) -> Layer4Scores:
    """Score discovery and error prevention."""
    discovery_scores: list[int] = []
    error_scores: list[int] = []

    for _ in range(repeat):
        d_response = await judge.judge(
            DISCOVERY_SYSTEM_PROMPT,
            DISCOVERY_USER_PROMPT.format(
                input_text=input_text[:2000],
                output=output[:4000],
            ),
        )
        discovery_scores.append(_parse_count(d_response))

        e_response = await judge.judge(
            ERROR_SYSTEM_PROMPT,
            ERROR_USER_PROMPT.format(
                input_text=input_text[:2000],
                output=output[:4000],
            ),
        )
        error_scores.append(_parse_score(e_response))

    return Layer4Scores(
        discovery_count=sum(discovery_scores) / len(discovery_scores),
        error_prevention=sum(error_scores) / len(error_scores),
    )
