"""The playground's task presets: prompts, the JSON schema each can be constrained to, and
whether the schema is on by default.

  Emotion   the seven-label emotion prompt, no system prompt of our own; free output by default
  Combined  the Inter-2 Audio system prompt + one instruction covering engagement, social
            signals, emotion and delivery; constrained by default, since free output can skip
            sections

When constrained, every label is limited to the values listed in the prompt and every
rationale is capped at 300 characters.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

PROMPTS = Path(__file__).parent / "prompts"
EMOTIONS = ["angry", "disgusted", "fearful", "happy", "neutral", "sad", "surprised"]
CATALOG_EMOTIONS = ["Anger", "Disgust", "Fear", "Happiness", "Neutral", "Sadness", "Surprise"]
SIGNALS = ["Hesitation", "Agreement", "Uncertainty", "Frustration", "Interest", "Confidence"]
ENGAGEMENT = ["engaged", "neutral", "disengaged"]
DELIVERY = ["Normal", "Crying", "Laughing", "Mumbling", "Shouting", "Singing", "Whispering"]
LEVELS = ["high", "medium", "low"]
RATIONALE_MAX_CHARS = 300  # keeps the longest constrained answer well inside 768 new tokens


def _prompt(name: str) -> str:
    return (PROMPTS / name).read_text(encoding="utf-8").strip()


def emotion_schema(confidence_steps: int = 4, evidence_max_chars: int = 0) -> dict:
    """{primary_emotion, confidence, evidence}: the label limited to the 7 emotions.

    confidence is limited to i/confidence_steps (a power of 2, so xgrammar keeps
    the short float form) and evidence is capped at evidence_max_chars (0 forces
    it empty): the label is what matters for this task, and an unbounded number
    can loop on digits. Properties are generated in this order,
    so the label comes first and neither cap can change it under greedy decoding.
    """
    if confidence_steps < 1 or confidence_steps & (confidence_steps - 1):
        raise ValueError(f"confidence_steps must be a power of 2, got {confidence_steps}")
    return {
        "type": "object",
        "properties": {
            "primary_emotion": {"type": "string", "enum": EMOTIONS},
            "confidence": {"enum": [i / confidence_steps for i in range(confidence_steps + 1)]},
            "evidence": {"type": "string", "maxLength": evidence_max_chars},
        },
        "required": ["primary_emotion", "confidence", "evidence"],
    }


def _labelled_list(names: list[str], min_items: int, max_items: int, rationale: bool = True) -> dict:
    fields = {
        "name": {"type": "string", "enum": names},
        "probability": {"type": "string", "enum": LEVELS},
    }
    if rationale:
        fields["rationale"] = {"type": "string", "maxLength": RATIONALE_MAX_CHARS}
    return {
        "type": "array",
        "minItems": min_items,
        "maxItems": max_items,
        "items": {
            "type": "object",
            "properties": fields,
            "required": list(fields),
        },
    }


def combined_schema() -> dict:
    """{engagement_status, signals, emotions, delivery}, in that order (the instruction's order)."""
    return {
        "type": "object",
        "properties": {
            "engagement_status": {"type": "string", "enum": ENGAGEMENT},
            "signals": _labelled_list(SIGNALS, 0, len(SIGNALS)),
            "emotions": _labelled_list(CATALOG_EMOTIONS, 1, 1),
            "delivery": {"type": "string", "enum": DELIVERY},
        },
        "required": ["engagement_status", "signals", "emotions", "delivery"],
    }


@dataclass(frozen=True)
class Task:
    system_prompt: str
    user_prompt: str
    schema: dict
    constrain: bool  # the schema checkbox's default for this task


TASKS: dict[str, Task] = {
    "Emotion": Task("", _prompt("emotion_user.txt"), emotion_schema(), constrain=False),
    "Combined": Task(_prompt("shared_system.txt"), _prompt("combined_user.txt"), combined_schema(), constrain=True),
}
DEFAULT_TASK = "Emotion"