#!/usr/bin/env python
"""
Orpheus-3B Hebrew TTS — 5 sequential continual-fine-tuning rounds.

Self-contained. Runs on a RunPod A100 started with dockerStartCmd
(no SSH). Writes progress to HuggingFace Hub (checkpoint per round =
round complete) and to a log file that is uploaded to a progress repo
so the driving agent can retrieve it without pod access.

Critical reference: canopyai/Orpheus-TTS (GitHub)
  - prompt format from orpheus_tts_pypi/orpheus_tts/engine_class.py::_format_prompt
  - audio encode/decode layout from orpheus_tts_pypi/orpheus_tts/decoder.py
  - training data format: (input_ids, labels) pre-tokenized, labels=input_ids default

Vocab mapping (verified from unsloth/orpheus-3b-0.1-ft tokenizer.json):
  <custom_token_N>  ->  vocab_id 128256 + N
  128257 = SOA (<custom_token_1>)
  128258 = EOA (<custom_token_2>)
  128259 = SOT (<custom_token_3>)
  128260, 128261 = separators (<custom_token_4>, <custom_token_5>)

Audio token formula (per canopy decoder.turn_token_into_id):
  vocab_id = 128266 + (pos_in_frame * 4096) + snac_code
  where pos_in_frame is 0..6 and the 7-slot frame layout (from decoder) is:
    slot 0: L0[j]         (coarsest, 1 per frame)
    slot 1: L1[2j]        (mid, first of pair)
    slot 2: L2[4j]        (finest, first quad)
    slot 3: L2[4j+1]
    slot 4: L1[2j+1]
    slot 5: L2[4j+2]
    slot 6: L2[4j+3]
"""
from __future__ import annotations

import json
import os
import pathlib
import sys
import time
import traceback
from datetime import datetime

# ---------- CONFIG ----------
BASE_MODEL = "unsloth/orpheus-3b-0.1-ft"
SNAC_MODEL = "hubertsiuzdak/snac_24khz"
DATASET_ID = "imvladikon/hebrew_speech_kan"
WHISPER_MODEL = "openai/whisper-large-v3"

HF_USER = "oridror"
REPO_TEMPLATE = "{user}/orpheus-3b-hebrew-r{round}"
PROGRESS_REPO = f"{HF_USER}/orpheus-hebrew-progress"

ROUND_SIZES = {1: 1000, 2: 1500, 3: 1500, 4: 1500, 5: 1500}
EVAL_SIZE = 10
MAX_AUDIO_SECONDS = 12.0
MIN_AUDIO_SECONDS = 1.0
MAX_SEQ_LEN = 2560

WORKDIR = pathlib.Path("/workspace/orpheus")
DATA_CACHE = WORKDIR / "data"
MODEL_CACHE = WORKDIR / "models"
CKPT_DIR = WORKDIR / "ckpt"
LOG_FILE = WORKDIR / "progress.jsonl"

# Orpheus special token IDs
BOS = 128000              # <|begin_of_text|>
EOT = 128009              # <|eot_id|>
SOT = 128259              # <custom_token_3>  start of "speak this" prompt
SEP_A = 128260
SEP_B = 128261
SOA = 128257              # <custom_token_1>  start of audio
EOA = 128258              # <custom_token_2>  end of audio
PAD = 128263

AUDIO_TOK_BASE = 128266   # code offset base (from canopy decoder formula)
CODEBOOK_SIZE = 4096

# LoRA
LORA_R = 64
LORA_ALPHA = 128
LORA_DROPOUT = 0.05
LR = 5e-5
WARMUP_STEPS = 20
BATCH_SIZE = 1
GRAD_ACCUM = 8
EPOCHS_PER_ROUND = 1


def log(msg: str, **extra) -> None:
    rec = {"ts": datetime.utcnow().isoformat() + "Z", "msg": msg, **extra}
    line = json.dumps(rec, ensure_ascii=False)
    print(line, flush=True)
    WORKDIR.mkdir(parents=True, exist_ok=True)
    with LOG_FILE.open("a", encoding="utf-8") as f:
        f.write(line + "\n")


def setup_env() -> None:
    WORKDIR.mkdir(parents=True, exist_ok=True)
    DATA_CACHE.mkdir(parents=True, exist_ok=True)
    MODEL_CACHE.mkdir(parents=True, exist_ok=True)
    CKPT_DIR.mkdir(parents=True, exist_ok=True)
    os.environ["HF_HOME"] = str(MODEL_CACHE)
    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
    log("env.ready", workdir=str(WORKDIR))


# ---------- SNAC ENCODE / DECODE (canopy-compatible) ----------

def snac_encode_one(wav_24k, snac_model, device):
    """Encode one 24kHz mono numpy array to a flat list of per-position vocab IDs.

    Returns a list of length 7 * num_frames of int32 vocab IDs.
    """
    import numpy as np
    import torch

    w = torch.tensor(wav_24k, dtype=torch.float32, device=device).unsqueeze(0).unsqueeze(0)
    with torch.no_grad():
        codes = snac_model.encode(w)
    c0 = codes[0][0].cpu().tolist()   # (T,)
    c1 = codes[1][0].cpu().tolist()   # (2T,)
    c2 = codes[2][0].cpu().tolist()   # (4T,)
    T = len(c0)

    flat = []
    for j in range(T):
        flat.append(AUDIO_TOK_BASE + 0 * CODEBOOK_SIZE + c0[j])
        flat.append(AUDIO_TOK_BASE + 1 * CODEBOOK_SIZE + c1[2 * j])
        flat.append(AUDIO_TOK_BASE + 2 * CODEBOOK_SIZE + c2[4 * j])
        flat.append(AUDIO_TOK_BASE + 3 * CODEBOOK_SIZE + c2[4 * j + 1])
        flat.append(AUDIO_TOK_BASE + 4 * CODEBOOK_SIZE + c1[2 * j + 1])
        flat.append(AUDIO_TOK_BASE + 5 * CODEBOOK_SIZE + c2[4 * j + 2])
        flat.append(AUDIO_TOK_BASE + 6 * CODEBOOK_SIZE + c2[4 * j + 3])
    return flat


def snac_decode(flat_vocab_ids, snac_model, device):
    """Inverse of snac_encode_one. Takes flat list of vocab IDs, returns 24kHz wav np array."""
    import torch

    # Convert to per-position codes
    codes_by_pos = [[] for _ in range(7)]
    for idx, vid in enumerate(flat_vocab_ids):
        pos = idx % 7
        code = vid - AUDIO_TOK_BASE - pos * CODEBOOK_SIZE
        if code < 0 or code >= CODEBOOK_SIZE:
            return None
        codes_by_pos[pos].append(code)

    T = len(codes_by_pos[0])
    if T == 0:
        return None

    c0_list = codes_by_pos[0]
    c1_list = []
    c2_list = []
    for j in range(T):
        c1_list.append(codes_by_pos[1][j])
        c2_list.append(codes_by_pos[2][j])
        c2_list.append(codes_by_pos[3][j])
        c1_list.append(codes_by_pos[4][j])
        c2_list.append(codes_by_pos[5][j])
        c2_list.append(codes_by_pos[6][j])

    c0 = torch.tensor(c0_list, dtype=torch.int32, device=device).unsqueeze(0)
    c1 = torch.tensor(c1_list, dtype=torch.int32, device=device).unsqueeze(0)
    c2 = torch.tensor(c2_list, dtype=torch.int32, device=device).unsqueeze(0)
    with torch.inference_mode():
        wav = snac_model.decode([c0, c1, c2])
    return wav.squeeze().cpu().float().numpy()


# ---------- TRAIN SAMPLE BUILD ----------

def build_sample(text, audio_vocab_ids, tokenizer, voice_prefix="hebrew: "):
    """Build input_ids per canopy format:
       BOS, SOT, <text>, EOT, SEP_A, SEP_B, SOA, <audio>, EOA
    """
    text_with_voice = voice_prefix + text
    text_ids = tokenizer.encode(text_with_voice, add_special_tokens=False)
    input_ids = (
        [BOS, SOT]
        + text_ids
        + [EOT, SEP_A, SEP_B, SOA]
        + audio_vocab_ids
        + [EOA]
    )
    # Loss: mask prompt, learn audio+EOA
    prefix_len = 1 + 1 + len(text_ids) + 4  # BOS, SOT, text, EOT+SEP_A+SEP_B+SOA
    labels = [-100] * prefix_len + audio_vocab_ids + [EOA]
    return input_ids, labels


# ---------- DATA STREAM ----------

def stream_hebrew_clips(n_samples, seed, eval_n=0):
    import numpy as np
    import librosa
    from datasets import load_dataset

    ds = load_dataset(DATASET_ID, split="train", streaming=True)
    ds = ds.shuffle(seed=seed, buffer_size=5000)

    target = n_samples + eval_n
    collected = 0
    for sample in ds:
        audio = sample.get("audio")
        text = sample.get("sentence") or sample.get("text") or sample.get("transcription")
        if not audio or not text:
            continue
        arr = audio.get("array")
        sr = audio.get("sampling_rate", 16000)
        if arr is None:
            continue
        dur = len(arr) / sr
        if dur < MIN_AUDIO_SECONDS or dur > MAX_AUDIO_SECONDS:
            continue
        arr = np.asarray(arr, dtype=np.float32)
        if sr != 24000:
            arr = librosa.resample(arr, orig_sr=sr, target_sr=24000)
        yield arr, text.strip()
        collected += 1
        if collected >= target:
            return


def build_round_dataset(n_samples, eval_n, seed, snac_model, tokenizer, device):
    train_records = []
    eval_pairs = []  # (gt_text, gt_wav_24k)

    log("data.stream.start", target=n_samples + eval_n, seed=seed)
    count = 0
    for wav, text in stream_hebrew_clips(n_samples, seed, eval_n=eval_n):
        if count < eval_n:
            eval_pairs.append((text, wav))
            count += 1
            continue
        try:
            audio_ids = snac_encode_one(wav, snac_model, device)
        except Exception as e:
            log("data.encode_fail", error=str(e))
            continue
        input_ids, labels = build_sample(text, audio_ids, tokenizer)
        if len(input_ids) > MAX_SEQ_LEN:
            continue
        train_records.append({"input_ids": input_ids, "labels": labels})
        count += 1
        if len(train_records) % 100 == 0:
            log("data.progress", processed=count, kept=len(train_records))
    log("data.stream.done", train=len(train_records), eval=len(eval_pairs))
    return train_records, eval_pairs


# ---------- SANITY CHECK ----------

def sanity_check(base_path, snac_model, tokenizer, device):
    """Round-trip encode→decode a synthetic sinewave to verify SNAC pipeline,
    and test that the base model generates valid audio tokens on an English prompt.
    If either fails, training is pointless and we should abort."""
    import numpy as np
    import torch
    import soundfile as sf
    from transformers import AutoModelForCausalLM

    log("sanity.start")

    # 1. Encode → Decode round trip on a short synthetic clip
    t = np.linspace(0, 2.0, 48000, endpoint=False, dtype=np.float32)
    wav = 0.1 * np.sin(2 * np.pi * 440 * t).astype(np.float32)
    try:
        vocab_ids = snac_encode_one(wav, snac_model, device)
        log("sanity.encode_ok", n_tokens=len(vocab_ids), n_frames=len(vocab_ids) // 7)
        recon = snac_decode(vocab_ids, snac_model, device)
        if recon is None or len(recon) < 1000:
            log("sanity.decode_fail")
            return False
        sf.write(str(WORKDIR / "sanity_recon.wav"), recon, 24000)
        log("sanity.roundtrip_ok", recon_len=len(recon))
    except Exception as e:
        log("sanity.roundtrip_exception", error=str(e), tb=traceback.format_exc())
        return False

    # 2. Base model speaks English prompt
    try:
        model = AutoModelForCausalLM.from_pretrained(
            base_path, torch_dtype=torch.bfloat16, device_map=device
        )
        model.eval()
        prompt_ids = [BOS, SOT] + tokenizer.encode("tara: hello world", add_special_tokens=False) + [EOT, SEP_A, SEP_B, SOA]
        inp = torch.tensor([prompt_ids], dtype=torch.long, device=device)
        with torch.no_grad():
            out = model.generate(
                inp,
                max_new_tokens=700,
                do_sample=True,
                temperature=0.6,
                top_p=0.9,
                eos_token_id=EOA,
                pad_token_id=PAD,
            )
        gen = out[0, inp.shape[1]:].tolist()
        if EOA in gen:
            gen = gen[:gen.index(EOA)]
        n_audio = sum(1 for t in gen if AUDIO_TOK_BASE <= t < AUDIO_TOK_BASE + 7 * CODEBOOK_SIZE)
        log("sanity.base_gen", n_generated=len(gen), n_audio_tokens=n_audio)
        if n_audio < 14:
            log("sanity.base_gen_too_short")
            del model
            torch.cuda.empty_cache()
            return False
        audio_ids = [t for t in gen if AUDIO_TOK_BASE <= t < AUDIO_TOK_BASE + 7 * CODEBOOK_SIZE]
        # trim to multiple of 7
        audio_ids = audio_ids[: (len(audio_ids) // 7) * 7]
        recon = snac_decode(audio_ids, snac_model, device)
        if recon is not None:
            sf.write(str(WORKDIR / "sanity_base_english.wav"), recon, 24000)
            log("sanity.base_decode_ok", wav_len=len(recon))
        del model
        torch.cuda.empty_cache()
        return True
    except Exception as e:
        log("sanity.base_gen_exception", error=str(e), tb=traceback.format_exc())
        return False


# ---------- TRAIN ROUND ----------

def train_round(round_num, prev_model_path, train_records, save_path):
    import torch
    from datasets import Dataset
    from peft import LoraConfig, get_peft_model
    from transformers import (
        AutoModelForCausalLM,
        AutoTokenizer,
        Trainer,
        TrainingArguments,
    )

    log("train.load_base", path=prev_model_path, round=round_num)
    tokenizer = AutoTokenizer.from_pretrained(prev_model_path)
    model = AutoModelForCausalLM.from_pretrained(
        prev_model_path,
        torch_dtype=torch.bfloat16,
        device_map="cuda",
    )
    model.gradient_checkpointing_enable()
    model.enable_input_require_grads()

    lora_cfg = LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
        use_rslora=True,
    )
    model = get_peft_model(model, lora_cfg)
    model.print_trainable_parameters()

    ds_train = Dataset.from_list(train_records)

    def collate(batch):
        import torch as _t
        max_len = max(len(b["input_ids"]) for b in batch)
        input_ids, labels, attn = [], [], []
        for b in batch:
            pad = max_len - len(b["input_ids"])
            input_ids.append(b["input_ids"] + [PAD] * pad)
            labels.append(b["labels"] + [-100] * pad)
            attn.append([1] * len(b["input_ids"]) + [0] * pad)
        return {
            "input_ids": _t.tensor(input_ids, dtype=_t.long),
            "labels": _t.tensor(labels, dtype=_t.long),
            "attention_mask": _t.tensor(attn, dtype=_t.long),
        }

    out_adapter = CKPT_DIR / f"r{round_num}_adapter"
    args = TrainingArguments(
        output_dir=str(out_adapter),
        per_device_train_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUM,
        num_train_epochs=EPOCHS_PER_ROUND,
        learning_rate=LR,
        warmup_steps=WARMUP_STEPS,
        logging_steps=10,
        save_strategy="no",
        bf16=True,
        report_to="none",
        remove_unused_columns=False,
        optim="adamw_torch_fused",
        gradient_checkpointing=True,
    )
    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=ds_train,
        data_collator=collate,
    )

    t0 = time.time()
    trainer.train()
    train_secs = time.time() - t0
    log("train.done", round=round_num, seconds=int(train_secs), n=len(train_records))

    log("train.merge_save", round=round_num, path=str(save_path))
    merged = model.merge_and_unload()
    merged.save_pretrained(str(save_path), safe_serialization=True)
    tokenizer.save_pretrained(str(save_path))

    del trainer, model, merged
    torch.cuda.empty_cache()
    return train_secs


# ---------- EVAL ----------

def eval_round(round_num, model_path, eval_pairs, snac_model, device):
    import numpy as np
    import torch
    import librosa
    import soundfile as sf
    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

    log("eval.start", round=round_num, n=len(eval_pairs))
    if not eval_pairs:
        return {"wer": None, "n": 0}

    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(
        model_path, torch_dtype=torch.bfloat16, device_map=device
    )
    model.eval()

    asr = pipeline(
        "automatic-speech-recognition",
        model=WHISPER_MODEL,
        device=device,
        torch_dtype=torch.float16,
        generate_kwargs={"language": "he", "task": "transcribe"},
    )

    from jiwer import wer as wer_fn

    refs, hyps = [], []
    samples = []
    for idx, (gt_text, _) in enumerate(eval_pairs):
        prompt_ids = [BOS, SOT] + tokenizer.encode("hebrew: " + gt_text, add_special_tokens=False) + [EOT, SEP_A, SEP_B, SOA]
        inp = torch.tensor([prompt_ids], dtype=torch.long, device=device)
        with torch.no_grad():
            out = model.generate(
                inp,
                max_new_tokens=1400,
                do_sample=True,
                temperature=0.6,
                top_p=0.9,
                eos_token_id=EOA,
                pad_token_id=PAD,
                repetition_penalty=1.1,
            )
        gen = out[0, inp.shape[1]:].tolist()
        if EOA in gen:
            gen = gen[:gen.index(EOA)]
        audio_ids = [t for t in gen if AUDIO_TOK_BASE <= t < AUDIO_TOK_BASE + 7 * CODEBOOK_SIZE]
        audio_ids = audio_ids[: (len(audio_ids) // 7) * 7]
        if len(audio_ids) < 14:
            refs.append(gt_text)
            hyps.append("")
            continue
        wav_24k = snac_decode(audio_ids, snac_model, device)
        if wav_24k is None:
            refs.append(gt_text)
            hyps.append("")
            continue
        wav_16k = librosa.resample(wav_24k, orig_sr=24000, target_sr=16000)
        try:
            asr_out = asr({"array": wav_16k.astype(np.float32), "sampling_rate": 16000})
            hyp_text = asr_out["text"].strip()
        except Exception as e:
            log("eval.asr_fail", error=str(e))
            hyp_text = ""
        refs.append(gt_text)
        hyps.append(hyp_text)
        # Save first 3 wavs for manual inspection
        if idx < 3:
            wav_path = WORKDIR / f"r{round_num}_eval_{idx}.wav"
            sf.write(str(wav_path), wav_24k, 24000)
            samples.append({"ref": gt_text, "hyp": hyp_text, "wav": wav_path.name})

    non_empty = [(r, h) for r, h in zip(refs, hyps) if r]
    w = None
    if non_empty:
        rs, hs = zip(*non_empty)
        try:
            w = float(wer_fn(list(rs), list(hs)))
        except Exception:
            w = None

    del model, asr
    torch.cuda.empty_cache()

    log("eval.done", round=round_num, wer=w, n=len(refs),
        sample_ref=refs[0] if refs else "", sample_hyp=hyps[0] if hyps else "")
    return {"wer": w, "n": len(refs), "samples": samples,
            "sample_ref": refs[0] if refs else "", "sample_hyp": hyps[0] if hyps else ""}


# ---------- HF UPLOAD ----------

def upload_to_hf(round_num, model_path, eval_result, train_secs):
    from huggingface_hub import HfApi, create_repo
    repo_id = REPO_TEMPLATE.format(user=HF_USER, round=round_num)
    token = os.environ["HUGGINGFACE_HUB_TOKEN"]
    api = HfApi(token=token)

    log("hf.upload.start", repo=repo_id)
    try:
        create_repo(repo_id, token=token, repo_type="model", exist_ok=True, private=False)
    except Exception as e:
        log("hf.create_repo.warn", error=str(e))

    readme = f"""---
license: apache-2.0
language: he
base_model: {BASE_MODEL}
tags:
  - text-to-speech
  - hebrew
  - orpheus
  - continual-finetune
---

# Orpheus-3B Hebrew TTS — Round {round_num}

Continual LoRA fine-tune of `{BASE_MODEL}` for Hebrew TTS.

- Round: {round_num} of 5
- Training clips: {ROUND_SIZES[round_num]}
- Dataset: `{DATASET_ID}` (streamed, seed={round_num * 97})
- Audio codec: `{SNAC_MODEL}` (24kHz, 7-token frames)
- LoRA: r={LORA_R}, alpha={LORA_ALPHA}, use_rslora=True
- Training time: {int(train_secs)}s
- Eval WER (Whisper-large-v3 back-inference, n={eval_result.get('n')}): **{eval_result.get('wer')}**

## Eval sample

- Ref: `{eval_result.get('sample_ref','')}`
- Hyp: `{eval_result.get('sample_hyp','')}`

Trained autonomously from MYD monorepo on RunPod A100.
"""
    (pathlib.Path(model_path) / "README.md").write_text(readme, encoding="utf-8")

    # Optionally include the eval wavs if they exist
    for i in range(3):
        src = WORKDIR / f"r{round_num}_eval_{i}.wav"
        if src.exists():
            import shutil
            shutil.copy(src, pathlib.Path(model_path) / f"eval_{i}.wav")

    api.upload_folder(
        folder_path=str(model_path),
        repo_id=repo_id,
        repo_type="model",
        token=token,
        commit_message=f"Round {round_num} — WER={eval_result.get('wer')}",
    )
    log("hf.upload.done", repo=repo_id)
    return repo_id


def push_progress():
    """Upload current progress.jsonl to the progress repo for external monitoring."""
    try:
        from huggingface_hub import HfApi, create_repo
        token = os.environ.get("HUGGINGFACE_HUB_TOKEN")
        if not token:
            return
        create_repo(PROGRESS_REPO, token=token, repo_type="model", exist_ok=True, private=False)
        api = HfApi(token=token)
        api.upload_file(
            path_or_fileobj=str(LOG_FILE),
            path_in_repo="progress.jsonl",
            repo_id=PROGRESS_REPO,
            repo_type="model",
            token=token,
        )
        # Also push sanity wavs if present
        for name in ["sanity_recon.wav", "sanity_base_english.wav", "final.json"]:
            p = WORKDIR / name
            if p.exists():
                try:
                    api.upload_file(path_or_fileobj=str(p), path_in_repo=name,
                                    repo_id=PROGRESS_REPO, repo_type="model", token=token)
                except Exception:
                    pass
    except Exception as e:
        log("progress.push.warn", error=str(e))


# ---------- MAIN ----------

def main():
    sanity_only = "--sanity-only" in sys.argv
    skip_sanity = "--skip-sanity" in sys.argv

    try:
        setup_env()
        from huggingface_hub import snapshot_download

        log("base.download.start", model=BASE_MODEL)
        base_path = snapshot_download(
            BASE_MODEL,
            cache_dir=str(MODEL_CACHE),
            token=os.environ["HUGGINGFACE_HUB_TOKEN"],
        )
        log("base.download.done", path=base_path)

        from snac import SNAC
        import torch
        from transformers import AutoTokenizer

        device = "cuda"
        snac_model = SNAC.from_pretrained(SNAC_MODEL).to(device)
        snac_model.eval()
        log("snac.ready")

        tokenizer = AutoTokenizer.from_pretrained(base_path)

        # Sanity check
        if not skip_sanity:
            ok = sanity_check(base_path, snac_model, tokenizer, device)
            push_progress()
            if not ok:
                log("sanity.failed_abort")
                return
            if sanity_only:
                log("sanity.only.done")
                return

        current_model_path = base_path
        summary = []

        for r in range(1, 6):
            round_start = time.time()
            log("round.begin", round=r)
            save_path = CKPT_DIR / f"r{r}_merged"
            n = ROUND_SIZES[r]
            seed = r * 97

            train_records, eval_pairs = build_round_dataset(
                n, EVAL_SIZE, seed, snac_model, tokenizer, device
            )
            push_progress()
            if len(train_records) < 50:
                log("round.abort.insufficient_data", round=r, got=len(train_records))
                break

            train_secs = train_round(r, current_model_path, train_records, save_path)
            push_progress()

            eval_result = eval_round(r, save_path, eval_pairs, snac_model, device)
            push_progress()

            repo_id = upload_to_hf(r, save_path, eval_result, train_secs)

            summary.append({
                "round": r,
                "train_clips": len(train_records),
                "train_seconds": int(train_secs),
                "eval": {k: v for k, v in eval_result.items() if k != "samples"},
                "hf_repo": repo_id,
                "wall_seconds": int(time.time() - round_start),
            })
            push_progress()

            current_model_path = str(save_path)
            prev_prev = CKPT_DIR / f"r{r-2}_merged"
            if prev_prev.exists() and prev_prev != save_path:
                import shutil
                shutil.rmtree(prev_prev, ignore_errors=True)
                log("round.cleanup", removed=str(prev_prev))

            # Stop early if eval shows zero progress after round 1
            if r == 1:
                wer = eval_result.get("wer")
                if wer is not None and wer > 0.95:
                    log("round.abort.hopeless_wer", wer=wer)
                    break

        final = {
            "ended": datetime.utcnow().isoformat() + "Z",
            "base_model": BASE_MODEL,
            "dataset": DATASET_ID,
            "rounds": summary,
            "hf_user": HF_USER,
        }
        (WORKDIR / "final.json").write_text(json.dumps(final, indent=2), encoding="utf-8")
        log("run.complete", rounds_done=len(summary))
        push_progress()

    except Exception as e:
        log("fatal", error=str(e), tb=traceback.format_exc())
        push_progress()
        sys.exit(1)


if __name__ == "__main__":
    main()