# bora_finalize.py -- runs on the VM after training: pushes everything to the Hub model repo and verifies it.
# Waits for "TRAIN_EXIT=0" in /ephemeral/bora/train.log, uploads, then checks every remote file size against the local one.
# Writes FINALIZE_OK or FINALIZE_FAILED (with the reason) to /ephemeral/bora/finalize.log.
import sys
import time
from pathlib import Path

from huggingface_hub import HfApi, CommitOperationAdd
from huggingface_hub.hf_api import RepoFile

REPO = "Ami1ne/bora-v1"
RUN = Path("/ephemeral/bora_runs/v1")
OVF = Path("/ephemeral/bora_runs/overfit")
CODE = Path("/ephemeral/bora/code")
TRAIN_LOG = Path("/ephemeral/bora/train.log")
LOG = Path("/ephemeral/bora/finalize.log")


def log(msg):
    line = f"[{time.strftime('%m-%d %H:%M:%S')}] {msg}"
    print(line, flush=True)
    with open(LOG, "a") as f:
        f.write(line + "\n")


def main():
    log("waiting for TRAIN_EXIT in train.log")
    while True:
        text = TRAIN_LOG.read_text() if TRAIN_LOG.exists() else ""
        exits = [l for l in text.splitlines() if l.startswith("TRAIN_EXIT=")]
        if exits:
            break
        time.sleep(60)
    if exits[-1] != "TRAIN_EXIT=0" or "done at step 50000" not in text:
        log(f"FINALIZE_FAILED: training did not finish cleanly ({exits[-1]}); nothing pushed")
        sys.exit(1)

    emas = sorted(RUN.glob("ema_step*.safetensors"), key=lambda p: int(p.stem[8:]))
    files = {"ema_latest.safetensors": emas[-1]}
    for p in emas:
        files[f"checkpoints/{p.name}"] = p
    files["checkpoints/ckpt_last.pt"] = RUN / "ckpt_last.pt"              # model + EMA + optimizer, to resume training later
    for n in ("config.json", "log.csv", "val.csv", "train.log"):
        files[f"logs/{n}"] = RUN / n
    for p in sorted((RUN / "samples").glob("*")):
        files[f"samples/{p.name}"] = p
    for p in sorted((OVF / "samples").glob("*.mp4")):
        files[f"overfit/{p.name}"] = p
    for n in ("log.csv", "train.log"):
        files[f"overfit/{n}"] = OVF / n
    for p in sorted(CODE.glob("*.py")):
        files[f"code/{p.name}"] = p
    readme = RUN / "README.md"
    readme.write_text(f"""# Bora V1

Smallest video generation model from the @Codeically series. Flow-matching video DiT (336M parameters, dim 1024, depth 20)
over frozen Wan2.1-VAE latents. 256x256, 8 fps, blocks of 4 latent frames with 1 frame of overlap, 8 sampling steps.
Trained 50,000 steps at batch 32 on one A100-80G, on `Ami1ne/Bora` (Pexels / Pixabay / NASA clips).
Text encoder: google/t5-v1_1-base, 32 tokens.

- `ema_latest.safetensors`: EMA weights ({emas[-1].name}), bf16
- `checkpoints/ckpt_last.pt`: full training state to resume
- `code/`: bora_dit.py (model + sampler), bora_train.py, bora_dataset.py, bora_run.py
- `samples/`: sample videos every 2,000 steps; `logs/`: training and validation curves

Generate: `python code/bora_run.py sample --runs <dir with ema file> --prompts "ocean waves crashing"`.
Decode latents with the Wan2.1 VAE: `vae.decode(z * latents_std + latents_mean)`.
""")
    files["README.md"] = readme
    files = {k: v for k, v in files.items() if v.exists()}
    total = sum(v.stat().st_size for v in files.values())
    log(f"pushing {len(files)} files, {total / 1024**3:.2f} GB to {REPO}")

    api = HfApi()
    log(f"logged in as {api.whoami()['name']}")
    api.create_repo(REPO, repo_type="model", exist_ok=True)
    for attempt in range(1, 5):
        try:
            remote = {it.path: it.size for it in api.list_repo_tree(REPO, repo_type="model", recursive=True) if isinstance(it, RepoFile)}
            ops = [CommitOperationAdd(path_in_repo=k, path_or_fileobj=str(v)) for k, v in files.items() if remote.get(k) != v.stat().st_size]
            if ops:
                api.create_commit(REPO, repo_type="model", operations=ops, commit_message=f"Bora V1 final: {emas[-1].stem}")
            break
        except Exception as e:
            log(f"push attempt {attempt} failed: {type(e).__name__}: {str(e)[:200]}")
            time.sleep(60 * attempt)

    remote = {it.path: it.size for it in api.list_repo_tree(REPO, repo_type="model", recursive=True) if isinstance(it, RepoFile)}
    bad = [k for k, v in files.items() if remote.get(k) != v.stat().st_size]
    if bad:
        log(f"FINALIZE_FAILED: {len(bad)} files missing or wrong size on the Hub: {bad[:10]}")
        sys.exit(1)
    log(f"FINALIZE_OK: {len(files)} files verified on https://huggingface.co/{REPO} ({total / 1024**3:.2f} GB)")


if __name__ == "__main__":
    main()