import uuid
from dataclasses import dataclass


@dataclass(frozen=True)
class RatingCriterion:
    key: str
    label: str
    description: str


RATING_CRITERIA: list[RatingCriterion] = [
    RatingCriterion(
        key="prompt_adherence",
        label="Prompt Adherence",
        description="Subject, attributes, and composition match the prompt exactly",
    ),
    RatingCriterion(
        key="aesthetics",
        label="Aesthetics",
        description="Visual beauty, color harmony, and overall appeal",
    ),
    RatingCriterion(
        key="technical_quality",
        label="Technical Quality",
        description="Sharpness, clean render, and lack of artifacts",
    ),
    RatingCriterion(
        key="lighting",
        label="Lighting & Atmosphere",
        description="Light quality, shadow detail, highlights, and mood",
    ),
    RatingCriterion(
        key="detail_richness",
        label="Detail & Texture",
        description="Fine-grained surface textures and element detail",
    ),
    RatingCriterion(
        key="coherence",
        label="Coherence",
        description="Correct anatomy, physics, and spatial relationships",
    ),
]

RATING_OPTIONS = [
    "1 - Poor",
    "2 - Weak",
    "3 - Okay",
    "4 - Good",
    "5 - Excellent",
]

_RATING_VALUE_MAP = {
    option: index
    for index, option in enumerate(RATING_OPTIONS, start=1)
}


class PBRating:
    @classmethod
    def INPUT_TYPES(cls):
        required = {
            "notes": (
                "STRING",
                {
                    "multiline": True,
                    "default": "",
                    "tooltip": "Optional notes about why this image was rated this way.",
                },
            )
        }

        for criterion in RATING_CRITERIA:
            required[criterion.key] = (
                "COMBO",
                {
                    "options": RATING_OPTIONS,
                    "default": RATING_OPTIONS[2],
                    "tooltip": criterion.description,
                },
            )

        return {"required": required}

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return str(uuid.uuid4())

    RETURN_TYPES = ("RATING",)
    RETURN_NAMES = ("rating",)
    FUNCTION = "rate_prompt"
    CATEGORY = "Prompt Bench"

    def rate_prompt(self, notes: str = "", **ratings: str):
        rating_id = str(uuid.uuid4())

        criterion_scores: dict[str, int] = {}

        for criterion in RATING_CRITERIA:
            selected_option = ratings.get(criterion.key, None)
            if selected_option:
                criterion_scores[criterion.key] = _RATING_VALUE_MAP[selected_option]

        average_score = sum(criterion_scores.values()) / len(RATING_CRITERIA)
        payload = {
            "rating_id": rating_id,
            "notes": notes,
            "criteria": criterion_scores,
            "average_score": average_score,
        }

        return (payload,)