"""Load the actual Small export and reuse the existing channel-wise CFG loop."""
import json
import sys
import threading
import time
from contextlib import nullcontext

import numpy as np
import torch
from huggingface_hub import snapshot_download

import settings

sys.path.append(str(settings.PARENT))
from tts_engine import TTSEngine


def zero_temperature_sampler(sampler):
    """The upstream sampler requires positive T; T=0 means genuine argmax."""
    def sample(*args, **kwargs):
        if float(kwargs.get("temperature", 1)) == 0:
            kwargs["do_sample"] = False
        return sampler(*args, **kwargs)
    return sample


class SmallEngine(TTSEngine):
    def __init__(self):
        started = time.monotonic()
        self.device = settings.DEVICE
        self.lock = threading.Lock()
        self.refs = {}
        self.last_codes = []
        self.last_performance = None

        root = snapshot_download(settings.REPO, revision=settings.REVISION,
            allow_patterns=["code/*.py", "assets/**", f"checkpoints/{settings.STAGE}/*", "README.md", "LICENSE"])
        from pathlib import Path
        root = Path(root)
        sys.path.insert(0, str(root / "code"))
        import moss_small
        from large_talker import build_fresh
        from packing import ScorePacker
        moss_small.SFT3 = str(root / "assets/sft3")
        moss_small.QWEN = str(root / "assets/qwen3")
        moss_small.load_qwen_backbone = lambda model, log=print: None
        schema = json.loads((root / "assets/score_schema.json").read_text())
        print(f"[small] loading {settings.STAGE} ({settings.REVISION[:8]})", flush=True)
        self.model, cfg = build_fresh(schema, log=lambda _: None)
        state = torch.load(root / "checkpoints" / settings.STAGE / "model_bf16.pt", map_location="cpu", weights_only=True)
        self.model.load_state_dict(state, strict=True)
        self.model.tie_weights()
        self.model = self.model.to(self.device, dtype=torch.bfloat16).eval()
        self.model._sample_next_token = zero_temperature_sampler(self.model._sample_next_token)
        del state
        _, _, Processor = moss_small.export_classes()
        self.proc = Processor.from_pretrained(moss_small.SFT3,
            codec_path="OpenMOSS-Team/MOSS-Audio-Tokenizer-v2",
            codec_weight_dtype="bf16" if settings.FAST_INFERENCE else "fp32", codec_compute_dtype="bf16")
        self.proc.audio_tokenizer = self.proc.audio_tokenizer.to(self.device).eval()
        self.packer = ScorePacker(self.proc, cfg, schema)
        self.sr = int(self.proc.model_config.sampling_rate)
        self.parameters = sum(p.numel() for p in self.model.parameters())
        from performance_core import FastGenerator
        self.fast = FastGenerator(self)
        if settings.FAST_INFERENCE and settings.COMPILE_GLOBAL:
            # The global cache grows between frames; dynamic shapes avoid a
            # separate compile for every context length. Local decode instead
            # uses fixed-size CUDA graphs.
            from torch import _dynamo
            _dynamo.config.cache_size_limit = 64
            self.model.transformer.forward = torch.compile(
                self.model.transformer.forward, dynamic=True, mode="default")
        print(f"[small] ready: {self.parameters:,} parameters, {self.sr} Hz, {time.monotonic()-started:.1f}s", flush=True)

    @torch.inference_mode()
    def warm_search(self):
        """Compile continuation/cache-fork paths before advertising ready."""
        from beam_search import ChunkSearch
        from prompting import build
        class WarmScorer:
            device = "none (startup warm-up)"
            last_timing = {"asr_seconds":0., "quality_seconds":0., "total_seconds":0.}
            def score(self, waves, text, final):
                return [{"reward":float(len(w)), "prefix_wer":0.} for w in waves]
        spec = build({"text":"Ready when you are.","caption":"Warm conversation."},"caption","en",seconds=6.6)
        path = settings.HERE / "performance/benchmark-reference.wav"
        for n in [4,8,12,16]:
            ChunkSearch(self,WarmScorer()).generate(spec,n=n,beam_width=1,chunk_frames=10,
                max_seconds=6.6,seed=991,reference_path=path if path.is_file() else None)
        self.last_codes=[]
        self.last_performance=None

    @staticmethod
    def _sample_continue_or_end(mdl, lh, p, stop_bias=None):
        # The old 4.55B model's stop bias is not a measured Small default.
        return mdl._sample_next_assistant_text_token(local_hidden_states=lh,
            do_sample=p["text_temperature"] > 0, temperature=p["text_temperature"], top_k=p["text_top_k"], top_p=p["text_top_p"])

    @torch.inference_mode()
    def reference(self, path):
        key = str(path)
        if key not in self.refs:
            codes = self.proc.encode_audios_from_path(path, n_vq=12)[0].cpu()
            if not 0 < len(codes) <= 37:
                raise ValueError("Referenzaudio muss höchstens 37 Codec-Frames lang sein.")
            self.refs[key] = codes
        return self.refs[key]

    def forget_reference(self, path):
        self.refs.pop(str(path), None)

    def pack(self, prompt, spec, reference, n):
        record = {"prompt": prompt, "text": spec["text"], "frames": spec["frames"], "lang": spec["lang"]}
        row = self.packer.pack_mode(record, [], "reference" if reference is not None else "instruction", reference, generation=True)
        batch = self.packer.collate([row] * n)
        return batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device)

    def pack_many(self, specs, reference, neutral=False):
        from performance_core import combine_branches
        # The upstream collator right-pads, inappropriate for different
        # generation prefixes. Collate one row then explicitly left-pad.
        rows = [self.pack(s["unconditional"] if neutral else s["prompt"], s, reference, 1) for s in specs]
        return combine_branches(rows, self.model.config.pad_token_id, self.model.config.audio_pad_token_id)

    @torch.inference_mode()
    def generate_many(self, specs, guidance=1.5, seed=777, reference_path=None,
                      temperature=1.0, progress=None, streaming=False, batch_size=None, chunk_frames=10,
                      max_frames=None):
        from performance_core import StreamingCodec
        waves, self.last_codes, batches = [], [], []
        with self.lock, torch.cuda.device(self.device), torch.autocast("cuda", dtype=torch.bfloat16):
            reference = self.reference(reference_path) if reference_path else None
            batch_size = batch_size or settings.BATCH_SIZE
            for offset in range(0, len(specs), batch_size):
                group = specs[offset:offset + batch_size]
                count = len(group)
                branches = [self.pack_many(group, reference)]
                if guidance != 1 and any(s["prompt"] != s["unconditional"] for s in group):
                    branches.append(self.pack_many(group, reference, neutral=True))
                torch.manual_seed(seed + offset)
                torch.cuda.manual_seed_all(seed + offset)
                p = {"max_new_tokens": min(settings.MAX_FRAMES + 60, max(s["frames"] for s in group) + 60),
                     "audio_temperature": temperature, "audio_top_p": .95, "audio_top_k": 50,
                     "audio_repetition_penalty": 1., "text_temperature": 0. if temperature == 0 else 1.,
                     "text_top_p": 1., "text_top_k": 50}
                if max_frames is not None:
                    p["max_new_tokens"] = min(p["max_new_tokens"], max_frames)
                torch.cuda.synchronize(self.device)
                start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)]
                capture_before = self.fast.capture_seconds
                wall_start = time.perf_counter()
                start.record()
                decoder = self.proc.audio_tokenizer
                context = StreamingCodec(decoder, count, self.model.config.audio_pad_token_id) if streaming else nullcontext()
                outputs, frames, previous = [], [], 0
                first = None
                with context as incremental:
                    iterator = (self.fast.stream_frames(branches, p, chunk_frames=chunk_frames if streaming else p["max_new_tokens"]+1,
                                guidance=guidance, graph=True) if settings.FAST_INFERENCE else
                                self._stream_frames_cfg(branches, p, chunk_frames=chunk_frames if streaming else p["max_new_tokens"]+1,
                                guidance=guidance))
                    for frames in iterator:
                        if streaming:
                            outputs.append(incremental.decode(torch.stack(frames[previous:], dim=1)))
                            previous = len(frames)
                            if first is None:
                                first = torch.cuda.Event(enable_timing=True)
                                first.record()
                    if not frames:
                        raise RuntimeError("Das Modell hat keinen Audio-Frame erzeugt.")
                    codes = torch.stack(frames, dim=1)
                    valid = codes[:, :, 0].ne(self.model.config.audio_pad_token_id)
                    if not streaming:
                        result = decoder.decode(codes.masked_fill(~valid.unsqueeze(-1), 0).permute(2, 0, 1).contiguous(),
                            padding_mask=valid, num_quantizers=12, return_dict=True, chunk_duration=8)
                        outputs.append((result.audio.float().mean(1), result.audio_lengths))
                end.record()
                end.synchronize()
                seconds = start.elapsed_time(end) / 1000
                durations = valid.sum(-1).cpu().tolist()
                batches.append({"batch_size": count, "device_seconds": seconds,
                    "audio_seconds": [f / 12.5 for f in durations],
                    "model_wall_seconds": time.perf_counter() - wall_start,
                    "first_audio_seconds": start.elapsed_time(first) / 1000 if first else seconds,
                    "graph_capture_seconds": self.fast.capture_seconds - capture_before})
                for i, c in enumerate(codes):
                    c = c[valid[i]]
                    if len(c) < 2:
                        raise RuntimeError("Ein Take enthält weniger als zwei Audio-Frames.")
                    self.last_codes.append(c.cpu())
                    pieces = [pcm[i, :int(lengths[i])].cpu() for pcm, lengths in outputs if int(lengths[i]) > 0]
                    wave = torch.cat(pieces).numpy().reshape(-1)
                    if not np.isfinite(wave).all():
                        raise RuntimeError("Nicht-endliche Audiowerte im Take.")
                    waves.append(wave)
                if progress:
                    progress(len(waves), len(specs))
        total_device = sum(b["device_seconds"] for b in batches)
        total_audio = sum(sum(b["audio_seconds"]) for b in batches)
        self.last_performance = {"device_seconds": total_device, "raw_audio_seconds": total_audio,
            "aggregate_rtf": total_device / total_audio, "audio_seconds_per_second": total_audio / total_device,
            "batches": batches, "streaming": streaming, "chunk_frames": chunk_frames if streaming else None,
            "backend": "fused CFG + local CUDA graphs + compiled global" if settings.FAST_INFERENCE and settings.COMPILE_GLOBAL
                else "fused CFG + local CUDA graphs" if settings.FAST_INFERENCE else "legacy",
            "includes": "TTS prefill, autoregressive decode and PCM codec; GPU event span including launch gaps",
            "excludes": "reference encoding, prompt packing, Luna, alignment, scoring, CPU audio copy, MP3, network",
            "graph_capture_seconds": sum(b["graph_capture_seconds"] for b in batches)}
        return waves

    @torch.inference_mode()
    def warm_up(self):
        from prompting import build
        spec = build({"text": "Ready when you are.", "caption": "Warm, natural conversation."}, "caption", "en", seconds=2.)
        path = settings.HERE / "performance/benchmark-reference.wav"
        references = [None]
        with self.lock, torch.cuda.device(self.device), torch.autocast("cuda", dtype=torch.bfloat16):
            if path.is_file():
                references.append(self.reference(path))
            for n in [1, settings.BATCH_SIZE]:
                for reference in references:
                    branches = [self.pack(spec["prompt"], spec, reference, n),
                                self.pack(spec["unconditional"], spec, reference, n)]
                    params = {"max_new_tokens":20,"audio_temperature":1.,"audio_top_p":.95,"audio_top_k":50,
                        "audio_repetition_penalty":1.,"text_temperature":1.,"text_top_p":1.,"text_top_k":50}
                    frames=[]
                    for frames in self.fast.stream_frames(branches,params,chunk_frames=21,guidance=1.5,min_frames=20,graph=True):
                        pass
                    codes=torch.stack(frames,dim=1).permute(2,0,1).contiguous()
                    self.proc.audio_tokenizer.decode(codes,return_dict=True,chunk_duration=8)
            torch.cuda.synchronize(self.device)
        self.last_codes = []
        self.last_performance = None

    @torch.inference_mode()
    def generate(self, spec, n=1, guidance=1.5, seed=777, reference_path=None, temperature=1.0, progress=None, streaming=False, chunk_frames=10,
                 max_frames=None):
        batch_size, retries = min(n, settings.BATCH_SIZE), []
        while True:
            try:
                waves = self.generate_many([spec] * n, guidance, seed, reference_path, temperature, progress,
                                           streaming=streaming, batch_size=batch_size, chunk_frames=chunk_frames,
                                           max_frames=max_frames)
                self.last_performance["oom_batch_retries"] = retries
                return waves
            except torch.cuda.OutOfMemoryError:
                # The lock has been released by generate_many. Discard this
                # request's partial candidates and retry a smaller batch;
                # never reuse a half-finished codec state.
                self.fast.clear_graphs()
                torch.cuda.empty_cache()
                if batch_size == 1:
                    raise
                retries.append(batch_size)
                batch_size = max(1, batch_size // 2)

    @torch.inference_mode()
    def generate_legacy(self, spec, n=1, guidance=1.5, seed=777, reference_path=None, temperature=1.0, progress=None):
        waves, self.last_codes = [], []
        with self.lock, torch.autocast("cuda", dtype=torch.bfloat16):
            reference = self.reference(reference_path) if reference_path else None
            for offset in range(0, n, 4):
                count = min(4, n - offset)
                branches = [self.pack(spec["prompt"], spec, reference, count)]
                if guidance != 1.0 and spec["prompt"] != spec["unconditional"]:
                    branches.append(self.pack(spec["unconditional"], spec, reference, count))
                torch.manual_seed(seed + offset)
                torch.cuda.manual_seed_all(seed + offset)
                p = {"max_new_tokens": min(settings.MAX_FRAMES + 60, spec["frames"] + 60),
                     "audio_temperature": temperature, "audio_top_p": 0.95,
                     "audio_top_k": 50, "audio_repetition_penalty": 1.0,
                     "text_temperature": 0.0 if temperature == 0 else 1.0, "text_top_p": 1.0, "text_top_k": 50}
                frames = []
                for frames in self._stream_frames_cfg(branches, p, chunk_frames=p["max_new_tokens"] + 1, guidance=guidance):
                    pass
                if not frames:
                    raise RuntimeError("Das Modell hat keinen Audio-Frame erzeugt.")
                codes = torch.stack(frames, dim=1)
                for c in codes:
                    c = c[c[:, 0] != self.model.config.audio_pad_token_id]
                    if len(c) < 2:
                        raise RuntimeError("Ein Take enthält weniger als zwei Audio-Frames.")
                    self.last_codes.append(c.cpu())
                    decoded = self.proc.decode_audio_codes([c], return_stereo=False)[0]
                    wave = decoded.float().cpu().numpy().reshape(-1)
                    if not np.isfinite(wave).all():
                        raise RuntimeError("Nicht-endliche Audiowerte im Take.")
                    waves.append(wave)
                if progress:
                    progress(len(waves), n)
        return waves


class Scorers:
    """Use the same proven Judge/Parakeet/VoiceCLAP ranker, without its LoRAs."""
    def __init__(self):
        from asr_engine import ParakeetASR
        from bestofn import Judge
        from retrieval import Retriever
        self.asr = ParakeetASR(device=settings.SCORE_DEVICE)
        self.retriever = Retriever(device=settings.SCORE_DEVICE)
        self.retriever._load()
        self.judge = Judge(device=settings.SCORE_DEVICE, asr=self.asr)
        if not self.judge.ok or not self.retriever.ok:
            raise RuntimeError("Die Best-of-N-Scorer konnten nicht geladen werden.")
        self.judge.attach_clap(self.retriever)
        if not self.judge._load_heads():
            raise RuntimeError("Die Burst-Detektorköpfe fehlen.")
        # Confirmed numerically by search_validation.py: both packaged heads
        # use exactly this VoiceCLAP tower. Share its weights instead of keeping
        # three resident copies. The old scoring interface remains unchanged.
        self.judge.gen.embedder = self.retriever.model
        self.judge.bl.embedder = self.retriever.model

    def score(self, waves, spec):
        from bestofn import rank
        scores = self.judge.score(waves, 48000, spec["text"], general=spec["general"],
            script=spec["tagged"], tagged=spec["tagged"])
        from metrics import word_errors
        for score in scores:
            score["legacy_word_distance"] = score["wer"]
            score.update(word_errors(spec["text"], score["hyp"]))
        rank(scores)
        return scores