from typing import Any

from solar_eval.evaluators.base import BaseEvaluator
from solar_eval.evaluators.scoring import (
    error_count_to_severity,
    score_to_severity,
    severity_to_score,
    weighted_average,
)
from solar_eval.models.sample import EvalSample
from solar_eval.providers.base import BaseProvider


class LLMJudgeEvaluator(BaseEvaluator):
    """LLM-as-Judge evaluator with weighted categories.

    Ported from chosun-projects body evaluation.
    Categories have weights and are scored by error count -> severity -> score.
    """

    required_fields = frozenset({"input", "output", "reference"})

    def __init__(
        self,
        categories: list[dict[str, Any]],
        judge_prompts: dict[str, str] | None = None,
    ) -> None:
        self.categories = categories  # [{name, weight}]
        self.weights = {c["name"]: c["weight"] for c in categories}
        self.judge_prompts = judge_prompts or {}

    async def evaluate(
        self,
        sample: EvalSample,
        provider: BaseProvider | None = None,
        judge_model: str = "gpt-4o",
    ) -> dict[str, Any]:
        if provider is None:
            raise ValueError("LLMJudgeEvaluator requires a provider")

        scores = {}
        details = {}

        for category in self.categories:
            cat_name = category["name"]
            system_prompt = self.judge_prompts.get(cat_name, self._default_prompt(cat_name))

            user_content = (
                f"Korean (Source):\n{sample.input.get('original', '')}\n\n"
                f"Golden Reference:\n{sample.reference}\n\n"
                f"Model Output:\n{sample.output}"
            )

            # F1: judge 호출 실패를 삼켜 만점(perfect/1.0)으로 채점하지 않는다 -- 여기서
            # 잡지 않고 그대로 전파한다. 카테고리 하나가 죽으면 이 샘플의 judge 결과
            # 전체가 무효라, 남은 카테고리를 계속 도는 대신 즉시 중단한다. 실패를
            # eval_results 밖으로 빼는 것(eval_failures)은 호출자(runner/CLI)의 책임이다.
            result = await provider.chat_completion_json(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_content},
                ],
                model=judge_model,
            )
            count_key = self._get_count_key(cat_name)
            error_count = result.get(count_key, result.get("error_count", 0))
            severity = error_count_to_severity(error_count)
            score = severity_to_score(severity)

            scores[cat_name] = score
            details[cat_name] = {
                "error_count": error_count,
                "severity": severity,
                "score": score,
            }

        overall_score = weighted_average(scores, self.weights)
        overall_severity = score_to_severity(overall_score, tiers=7)

        return {
            "score": overall_score,
            "severity": overall_severity,
            "category_scores": scores,
            "details": details,
        }

    def aggregate(self, results: list[dict[str, Any]]) -> dict[str, Any]:
        if not results:
            return {"overall_score": 0.0, "scores": {}}

        # Average per category
        cat_scores: dict[str, list[float]] = {c["name"]: [] for c in self.categories}
        overall_scores: list[float] = []

        for r in results:
            overall_scores.append(r["score"])
            for cat_name, score in r.get("category_scores", {}).items():
                if cat_name in cat_scores:
                    cat_scores[cat_name].append(score)

        avg_scores = {k: sum(v) / len(v) if v else 0.0 for k, v in cat_scores.items()}
        overall = sum(overall_scores) / len(overall_scores)

        return {
            "overall_score": overall,
            "scores": avg_scores,
            "severity": score_to_severity(overall, tiers=7),
            "num_samples": len(results),
        }

    def _get_count_key(self, category: str) -> str:
        key_map = {
            "factual_accuracy": "error_count",
            "content_addition": "addition_count",
            "content_omission": "omission_count",
            "style_appropriateness": "violation_count",
        }
        return key_map.get(category, "issue_count")

    def _default_prompt(self, category: str) -> str:
        return (
            f"You are an expert translation quality evaluator. "
            f"Evaluate the '{category}' of the model output compared to the golden reference. "
            f"BE EXTREMELY CONSERVATIVE. Only count OBVIOUS errors. WHEN IN DOUBT, DON'T COUNT. "
            f"Return JSON with a single field 'error_count' (integer)."
        )