# /// script
# dependencies = ["huggingface_hub", "datasets==2.18.0", "soundfile", "tqdm"]
# ///
"""
Fish Audio S2 Pro Hebrew LoRA fine-tuning on HF Jobs.

Usage: hf jobs uv run --flavor a100-large --timeout 6h --secrets HF_TOKEN https://hf.co/sm3222/fish-audio-training-scripts/resolve/main/train_s2_pro_hebrew.py
"""

import os
import shutil
import subprocess
import sys
import json
from datetime import datetime
from pathlib import Path

HF_TOKEN = os.environ["HF_TOKEN"]
OUTPUT_REPO = os.environ.get("OUTPUT_REPO", "sm3222/fish-audio-s2-pro-hebrew-lora")
BASE_MODEL = "fishaudio/s2-pro"
DATASET_ID = "shunyalabs/hebrew-speech-dataset"
WORKDIR = Path("/workspace")
DATA_DIR = WORKDIR / "data"
CKPT_DIR = WORKDIR / "checkpoints" / "s2-pro"
RESULTS_DIR = WORKDIR / "results"
FISH_DIR = WORKDIR / "fish-speech"
LOGS_DIR = WORKDIR / "logs"

for d in [DATA_DIR, CKPT_DIR, RESULTS_DIR, LOGS_DIR]:
    d.mkdir(parents=True, exist_ok=True)


def log(msg):
    print(f"[{datetime.now():%H:%M:%S}] {msg}", flush=True)


def run(cmd, **kwargs):
    log(f"$ {' '.join(str(p) for p in cmd)}")
    subprocess.run(cmd, check=True, text=True,
                   cwd=kwargs.pop("cwd", WORKDIR), **kwargs)


# ── 0. Clone fish-speech ────────────────────────────────────────────
log("=== 0/6 Clone fish-speech ===")
if not FISH_DIR.exists():
    run(["git", "clone", "https://github.com/fishaudio/fish-speech.git", str(FISH_DIR)])
run(["uv", "sync", "--python", "3.12", "--extra", "cu129"], cwd=FISH_DIR)

# ── 1. Download base model ──────────────────────────────────────────
log("=== 1/6 Download S2 Pro base model ===")
run(["hf", "download", "--token", HF_TOKEN, "--local-dir", str(CKPT_DIR), BASE_MODEL],
    cwd=FISH_DIR)

# ── 2. Download & prepare dataset ───────────────────────────────────
log("=== 2/6 Prepare Hebrew dataset ===")
run([sys.executable, "-c", f"""
from datasets import load_dataset
import soundfile as sf
import numpy as np
from pathlib import Path

ds = load_dataset("{DATASET_ID}", split="train", token="{HF_TOKEN}")
out = Path("{DATA_DIR}") / "speaker_0"
out.mkdir(parents=True, exist_ok=True)

for i, row in enumerate(ds):
    audio = row["audio"]["array"]
    sr = row["audio"]["sampling_rate"]
    text = row["transcript"]
    path = out / f"{{i:05d}}"
    sf.write(str(path) + ".wav", audio, sr)
    path.with_suffix(".lab").write_text(text, encoding="utf-8")
    if (i + 1) % 1000 == 0:
        print(f"  {{i+1}}/{{len(ds)}}")

print(f"Done: {{len(ds)}} pairs")
"""], cwd=FISH_DIR)

# ── 3. Extract VQGAN tokens ─────────────────────────────────────────
log("=== 3/6 Extract VQGAN semantic tokens ===")
run([
    "uv", "run", "python", "tools/vqgan/extract_vq.py",
    str(DATA_DIR), "--num-workers", "4", "--batch-size", "16",
    "--checkpoint-path", str(CKPT_DIR / "codec.safetensors"),
], cwd=FISH_DIR)

# ── 4. Pack protobuf ────────────────────────────────────────────────
log("=== 4/6 Pack protobuf ===")
run([
    "uv", "run", "python", "tools/llama/build_dataset.py",
    "--input", str(DATA_DIR),
    "--output", str(DATA_DIR / "protos"),
    "--text-extension", ".lab", "--num-workers", "4",
], cwd=FISH_DIR)

# ── 5. LoRA fine-tune ──────────────────────────────────────────────
log("=== 5/6 LoRA fine-tuning ===")
os.environ["HYDRA_FULL_ERROR"] = "1"
run([
    "uv", "run", "python", "fish_speech/train.py",
    f"--config-name=text2semantic_finetune",
    f"project=hebrew_s2_pro_lora",
    f"pretrained_ckpt_path={CKPT_DIR}",
    f"train_dataset.proto_files=[{DATA_DIR / 'protos'}]",
    f"val_dataset.proto_files=[{DATA_DIR / 'protos'}]",
    "+lora@model.model.lora_config=r_8_alpha_16",
    "trainer.max_steps=5000",
    "trainer.accumulate_grad_batches=4",
    "data.batch_size=2",
    "data.num_workers=4",
], cwd=FISH_DIR)

# ── 6. Merge LoRA + push to Hub ─────────────────────────────────────
log("=== 6/6 Merge LoRA + push ===")
ckpt_dir = WORKDIR / "hebrew_s2_pro_lora" / "checkpoints"
checkpoints = sorted(ckpt_dir.glob("step_*.ckpt"))
if not checkpoints:
    raise FileNotFoundError(f"No checkpoints in {ckpt_dir}")

best = checkpoints[-2] if len(checkpoints) > 1 else checkpoints[0]
log(f"Merging: {best.name}")

merged = RESULTS_DIR / "merged"
run([
    "uv", "run", "python", "tools/llama/merge_lora.py",
    "--lora-config", "r_8_alpha_16",
    "--base-weight", str(CKPT_DIR),
    "--lora-weight", str(best),
    "--output", str(merged),
], cwd=FISH_DIR)

from huggingface_hub import HfApi
api = HfApi(token=HF_TOKEN)
api.create_repo(repo_id=OUTPUT_REPO, repo_type="model", exist_ok=True)
api.upload_folder(folder_path=str(merged), repo_id=OUTPUT_REPO,
                  repo_type="model", ignore_patterns=[".*"])
log(f"✓ {OUTPUT_REPO}")