"""
Colab runner for Inflect-Nano smoothness fine-tuning.

Single script that handles the full pipeline with graceful checkpointing —
safe to re-run after interruptions.  Each step is skipped if its output
already exists.

Usage in Colab:
    !wget https://your-server/colab_run.py
    !python colab_run.py
"""
from __future__ import annotations

import os
import subprocess
import sys
import time
from pathlib import Path

# ---------------------------------------------------------------------------
# Config — change these
# ---------------------------------------------------------------------------
REPO_URL = "https://git.quietidiot.com/michael-treadgold/speech-nano.git"
REPO_DIR = Path("/content/speech-nano")
DATASET = "MikhailT/cmu-arctic"
SPEAKER = "rms"                # "rms", "bdl", "jmk", "awb", "ksp" (CMU ARCTIC male speakers)
TRAIN_STEPS = 5000
BATCH_SIZE = 8

# ---------------------------------------------------------------------------
# Step helpers
# ---------------------------------------------------------------------------
def step(name: str) -> None:
    print(f"\n{'='*60}\n  {name}\n{'='*60}")

def run(cmd: str, **kwargs) -> bool:
    print(f"  $ {cmd}")
    result = subprocess.run(cmd, shell=True, **kwargs)
    return result.returncode == 0

def already(path: Path | str) -> bool:
    p = Path(path)
    ok = p.exists()
    if ok:
        print(f"  ⏭  Already exists: {p}")
    return ok

def pip_installed(pkg: str) -> bool:
    result = subprocess.run(
        [sys.executable, "-c", f"import {pkg}"],
        capture_output=True, text=True, timeout=30,
    )
    return result.returncode == 0


def main() -> None:
    start = time.time()

    # ---- Step 1: Clone or pull repo ----
    step("1/6  Clone / pull repo")
    if REPO_DIR.exists():
        print(f"  Repo exists, pulling latest...")
        subprocess.run(["git", "-C", str(REPO_DIR), "pull"], check=False)
        subprocess.run(["git", "-C", str(REPO_DIR), "lfs", "pull"], check=False)
    else:
        run(f"git clone {REPO_URL} {REPO_DIR}")
    sys.path.insert(0, str(REPO_DIR))
    sys.path.insert(0, str(REPO_DIR / "third_party" / "tiny_tts_frontend"))
    os.chdir(REPO_DIR)
    print(f"  Working dir: {REPO_DIR}")

    # ---- Step 2: Install dependencies ----
    step("2/6  Install dependencies")
    pkgs = "torch torchaudio soundfile numpy g2p_en transformers gradio numba scipy datasets"
    if not pip_installed("g2p_en"):
        run(f"{sys.executable} -m pip install -q {pkgs}")
    else:
        print("  ⏭  Already installed (g2p_en found)")

    import nltk
    nltk.download("averaged_perceptron_tagger_eng", quiet=True)
    nltk.download("cmudict", quiet=True)
    print("  NLTK data OK")

    # ---- Step 3: Download model weights ----
    step("3/6  Download model weights")
    WEIGHTS_DIR = REPO_DIR / "weights"
    AC_WEIGHTS = WEIGHTS_DIR / "inflect_nano_v1_acoustic.pt"
    VO_WEIGHTS = WEIGHTS_DIR / "inflect_nano_v1_vocoder.pt"

    if AC_WEIGHTS.exists() and AC_WEIGHTS.stat().st_size > 1000:
        print(f"  ⏭  Acoustic weights OK ({AC_WEIGHTS.stat().st_size / 1e6:.1f} MB)")
    else:
        run(f"{sys.executable} download_weights.py")

    # ---- Step 4: Preprocess dataset ----
    step("4/6  Preprocess dataset")
    JSONL = Path(f"/content/durations_{SPEAKER}.jsonl")

    if JSONL.exists() and JSONL.stat().st_size > 100:
        print(f"  ⏭  JSONL exists ({JSONL.stat().st_size / 1e6:.2f} MB)")
    else:
        ok = run(
            f"{sys.executable} preprocess_dataset.py hf "
            f"--dataset {DATASET} --split {SPEAKER} "
            f"--out {JSONL} --voice-id mark"
        )
        if not ok:
            print("  ❌  Preprocessing failed. Check error above.")
            sys.exit(1)

    # ---- Step 5: Train ----
    step("5/6  Train smoothness model")
    OUT_DIR = Path("/content/checkpoints/smooth-v1")
    LATEST_CKPT = OUT_DIR / "inflect-smooth-latest.pt"

    import torch
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"  Device: {device}  ({torch.cuda.get_device_name(0) if device == 'cuda' else 'CPU'})")

    if LATEST_CKPT.exists():
        print(f"  ⏭  Checkpoint exists — adding {TRAIN_STEPS} more steps")
        resume_flag = "--resume"
    else:
        print(f"  Starting fresh — {TRAIN_STEPS} steps")
        resume_flag = ""

    ok = run(
        f"{sys.executable} -m inflect_nano.train_smooth "
        f"--durations-jsonl {JSONL} "
        f"--out-dir {OUT_DIR} "
        f"--init-checkpoint {AC_WEIGHTS} "
        f"--vocoder-checkpoint {VO_WEIGHTS} "
        f"--steps {TRAIN_STEPS} "
        f"--batch-size {BATCH_SIZE} "
        f"--device {device} "
        f"{resume_flag}"
    )
    if not ok:
        print("  ❌  Training failed. Check error above.")
        sys.exit(1)

    # ---- Step 6: Generate sample ----
    step("6/6  Generate sample")
    from inference import load_acoustic, load_vocoder, synthesize
    import soundfile as sf
    import numpy as np

    ac_path = LATEST_CKPT if LATEST_CKPT.exists() else AC_WEIGHTS
    acoustic, speakers, ap = load_acoustic(ac_path, torch.device(device))
    vocoder, vp = load_vocoder(VO_WEIGHTS, torch.device(device))
    print(f"  Acoustic: {ap:,} params  Vocoder: {vp:,} params  Total: {ap+vp:,}")

    text = "Every man is destined to die, but his work echoes through the ages."
    audio = synthesize(
        text, acoustic, vocoder, speakers, torch.device(device),
        smooth_prosody=True, mel_smooth_sigma=0.8,
    )
    out = Path("/content/sample_smooth.wav")
    sf.write(str(out), audio, 24000, subtype="PCM_16")

    elapsed = (time.time() - start) / 60
    print(f"\n{'='*60}")
    print(f"  Done in {elapsed:.0f}m.  Sample: {out}  ({audio.size/24000:.1f}s)")
    print(f"  Checkpoints: {OUT_DIR}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()