"""
clankerDiffusion — Modal L4 training.

Builds a CUDA image, pulls the latest code from the HF code repo at runtime,
keeps data + checkpoints on a persistent Modal Volume (/vol), injects HF_TOKEN
from .env, trains a big hybrid model on an L4 (24 GB), and pushes every
checkpoint to the HF repo:

    coderofpears/clankerDiffusion-checkpoints

Data is generated in-container by prep.py (fast HF egress), scaled by --scale,
so no large file transfer is needed.

Launch (after `modal token new` on this machine):
    modal run modal_train.py::train_on_l4 --hours 20 --size large --scale 3
"""
import os
from modal import App, Image, Volume, Secret

HERE = os.path.dirname(os.path.abspath(__file__))
CODE_REPO = "coderofpears/clankerDiffusion-base"
DATA_REPO = "coderofpears/clankerDiffusion-data"
CKPT_REPO = "coderofpears/clankerDiffusion-checkpoints"
DOTENV = os.path.join(HERE, ".env")

image = (
    Image.debian_slim()
    .pip_install("torch==2.11.0", index_url="https://download.pytorch.org/whl/cu128")
    .pip_install("numpy", "tokenizers", "datasets", "safetensors",
                 "huggingface_hub", "hf-transfer")
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1", "TOKENIZERS_PARALLELISM": "false"})
)

app = App("clanker-diffusion", image=image)
volume = Volume.from_name("clanker-vol", create_if_missing=True)

SIZES = {
    "base":  dict(d_model=768,  n_layers=12, n_heads=12, d_ff=2048, batch=24),
    "large": dict(d_model=2048, n_layers=24, n_heads=16, d_ff=5504, batch=8),
    "xl":    dict(d_model=2560, n_layers=28, n_heads=20, d_ff=6912, batch=4),
    # ~0.8B: fits L4 comfortably, more throughput than the 1.35B "large"
    "08b":   dict(d_model=1536, n_layers=20, n_heads=16, d_ff=4096, batch=12,
                  rope_scale=16.0),
}


@app.function(
    image=image,
    gpu="L4",
    timeout=82800,  # 23h, under Modal's 24h cap
    volumes={"/vol": volume},
    secrets=[Secret.from_dotenv(DOTENV)],
)
def train_on_l4(hours: float = 20.0, ckpt_every: int = 250,
                size: str = "large", batch: int = None, scale: float = 3.0):
    import subprocess
    import sys
    import shutil

    # ---- pull latest code from HF (use huggingface_hub, not the `hf` CLI) ----
    from huggingface_hub import snapshot_download
    os.makedirs("/root/clanker", exist_ok=True)
    snapshot_download(repo_id=CODE_REPO, repo_type="model",
                      local_dir="/root/clanker",
                      token=os.environ.get("HF_TOKEN"))
    os.chdir("/root/clanker")

    data_dir, ckpt_dir = "/vol/data", "/vol/checkpoints"
    os.makedirs(data_dir, exist_ok=True)
    os.makedirs(ckpt_dir, exist_ok=True)

    # reuse the canonical tokenizer / RAG corpus already fetched with the code
    for name in ("tokenizer.json", "meta.json", "rag_corpus.txt"):
        src = os.path.join("/root/clanker/data", name)
        if os.path.exists(src):
            shutil.copy(src, os.path.join(data_dir, name))

    os.makedirs(data_dir, exist_ok=True)
    os.makedirs(ckpt_dir, exist_ok=True)

    # ---- ensure data (generate in-container; fast HF egress) ----
    # Only treat the dataset as complete if meta.json exists (a partial
    # train.bin from a killed run must be regenerated).
    meta_path = os.path.join(data_dir, "meta.json")
    if not os.path.exists(meta_path):
        stale = os.path.join(data_dir, "train.bin")
        if os.path.exists(stale):
            print("[modal] found partial train.bin without meta.json; discarding.")
            os.remove(stale)
        print(f"[modal] generating data locally (scale={scale}) ...")
        subprocess.run([sys.executable, "prep.py", "--out-dir", data_dir,
                        "--scale", str(scale)], check=True)

    spec = dict(SIZES.get(size, SIZES["large"]))
    if batch:
        spec["batch"] = batch

    cmd = [
        sys.executable, "train.py",
        "--hours", str(hours),
        "--batch", str(spec["batch"]),
        "--ckpt-every", str(ckpt_every),
        "--data-dir", data_dir,
        "--ckpt-dir", ckpt_dir,
        "--hf-repo", CKPT_REPO,
        "--d-model", str(spec["d_model"]),
        "--n-layers", str(spec["n_layers"]),
        "--n-heads", str(spec["n_heads"]),
        "--d-ff", str(spec["d_ff"]),
    ]
    if "rope_scale" in spec:
        cmd += ["--rope-scale", str(spec["rope_scale"])]
    print("[modal] launching:", " ".join(cmd), flush=True)
    subprocess.run(cmd, check=True)
    print("[modal] training finished", flush=True)


@app.local_entrypoint()
def main(hours: float = 20.0, ckpt_every: int = 250, size: str = "large",
         batch: int = None, scale: float = 3.0):
    train_on_l4.remote(hours=hours, ckpt_every=ckpt_every, size=size,
                        batch=batch, scale=scale)