"""Versioned, checkpoint-native prompting and generation defaults."""
from __future__ import annotations
from dataclasses import dataclass, field, asdict
import math
import re

EOD = 151643
ABC_START, ABC_END = 151847, 151848
MUSIC_START, MUSIC_END = 151851, 151852
CODEC_OFFSET, CODEC_SIZE = 151853, 32768
LATENT_START, LATENT_END, LATENT_PAD = 184621, 184622, 184623
VOCAB_SIZE, CONTEXT = 184704, 24576
PROTOCOL_VERSION = "yue2-native-v1"
INSTRUCTIONS = {
    "off": "Generate music with codec tokens from the given conditions.",
    "melody": "Generate a melody-only ABC transcription without chord symbols, then generate music with codec tokens from the given conditions.",
    "full": "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions.",
}


@dataclass(frozen=True)
class Sampling:
    temperature: float = 1.0
    top_p: float = 0.95
    top_k: int = 100
    repetition_penalty: float = 1.2
    penalty_window: int = 50
    min_tokens: int = 200
    max_tokens: int = 9000

    def __post_init__(self):
        if any(type(x) is not int for x in (self.top_k, self.penalty_window, self.min_tokens, self.max_tokens)):
            raise ValueError("Sampling counts must be integers")
        if not all(math.isfinite(x) for x in (self.temperature, self.top_p, self.repetition_penalty)):
            raise ValueError("Sampling numbers must be finite")
        if not 0 <= self.temperature <= 5 or not 0 < self.top_p <= 1 or self.top_k < 1:
            raise ValueError("Invalid sampling temperature/top_p/top_k")
        if self.repetition_penalty <= 0 or not 1 <= self.penalty_window <= 100:
            raise ValueError("Invalid repetition penalty/window")
        if not 0 <= self.min_tokens <= self.max_tokens or self.max_tokens < 1:
            raise ValueError("Require 0 <= min_tokens <= max_tokens")


@dataclass(frozen=True)
class GenerationConfig:
    abc: Sampling = field(default_factory=lambda: Sampling(.7, .9, 30, 1.005, 100, 32, 4096))
    semantic: Sampling = field(default_factory=Sampling)
    ode_steps: int = 32
    ode_method: str = "midpoint"
    context: int = CONTEXT
    version: str = PROTOCOL_VERSION

    def __post_init__(self):
        if self.context != CONTEXT or self.ode_method != "midpoint" or type(self.ode_steps) is not int or self.ode_steps < 1:
            raise ValueError("Require context=24576 and midpoint with positive integer steps")

    def to_dict(self):
        return asdict(self)

    @classmethod
    def from_dict(cls, data):
        value = dict(data)
        defaults = cls()
        for key in ("abc", "semantic"):
            if key in value and isinstance(value[key], dict):
                value[key] = Sampling(**{**asdict(getattr(defaults, key)), **value[key]})
        return cls(**value)


def resolve_sampling(value, default):
    if value is None:
        return default
    if isinstance(value, Sampling):
        return value
    if isinstance(value, dict):
        return Sampling(**{**asdict(default), **value})
    raise TypeError("Sampling must be a Sampling object or a dictionary of overrides")


@dataclass(frozen=True)
class SongRequest:
    style: str
    lyrics: str
    cot: str = "full"
    seed: int = 831001
    abc: str | None = None
    cfg_scale: float | None = None
    id: str = "song"

    def __post_init__(self):
        if self.cot not in INSTRUCTIONS:
            raise ValueError("cot must be off, melody or full")
        if not isinstance(self.style, str) or not isinstance(self.lyrics, str):
            raise TypeError("style and lyrics must be strings")
        if type(self.seed) is not int or not 0 <= self.seed < 2**63:
            raise ValueError("seed must be an integer in [0, 2**63)")
        if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,179}", self.id) or self.id in {".", ".."}:
            raise ValueError("id must be a filename-safe identifier")
        if self.abc is not None and (self.cot == "off" or not isinstance(self.abc, str) or not self.abc.strip()):
            raise ValueError("External ABC requires nonempty text and cot=melody/full")
        if self.cfg_scale is not None and (not math.isfinite(self.cfg_scale) or not 0 <= self.cfg_scale <= 20):
            raise ValueError("cfg_scale must be finite and in [0,20]")

    @property
    def guidance(self):
        return (1.01 if self.cot == "off" else 1.0) if self.cfg_scale is None else self.cfg_scale

    def text(self):
        return f"{INSTRUCTIONS[self.cot]}\n[Tags]\n{self.style}\n[Lyrics]\n{self.lyrics}\n"

    def to_dict(self):
        return asdict(self)


def token_prefixes(request, tokenizer, abc_ids=None):
    base = [EOD] + tokenizer.encode(request.text())
    if request.cot == "off":
        return base + [ABC_START, ABC_END, MUSIC_START]
    if abc_ids is None:
        if request.abc is None:
            return base + [ABC_START]
        abc_ids = tokenizer.encode(request.abc)
    abc_ids = list(abc_ids)
    if any(type(token) is not int or not 0 <= token < EOD for token in abc_ids):
        raise ValueError("ABC IDs must remain inside the ordinary text vocabulary")
    return base + [ABC_START] + abc_ids + [ABC_END, MUSIC_START]


def negative_prefix(request, tokenizer, abc_ids=None):
    base = [EOD] + tokenizer.encode(INSTRUCTIONS[request.cot])
    if request.cot == "off":
        return base + [MUSIC_START]
    if abc_ids is None:
        raise ValueError("Symbolic CFG must retain the exact positive-branch ABC IDs")
    abc_ids = list(abc_ids)
    if any(type(token) is not int or not 0 <= token < EOD for token in abc_ids):
        raise ValueError("ABC IDs must remain inside the ordinary text vocabulary")
    return base + [ABC_START] + abc_ids + [ABC_END, MUSIC_START]


def chunk_ranges(frames, prefix_tokens, context=CONTEXT):
    size = min((context - prefix_tokens - 3) // 2, CONTEXT)
    if frames < 1 or size < 1:
        raise ValueError("Empty codec or prefix leaves no acoustic context")
    return [(a, min(a + size, frames)) for a in range(0, frames, size)]