# bora_run.py -- command-line runner for a rented GPU box (same steps as bora_train.ipynb, no Jupyter needed).
#
#   python bora_run.py download --data /ephemeral/bora_data
#   python bora_run.py overfit  --data /ephemeral/bora_data --runs /ephemeral/bora_runs
#   python bora_run.py train    --data /ephemeral/bora_data --runs /ephemeral/bora_runs --batch 32 --steps 50000
#   python bora_run.py sample   --data ... --runs ... --prompts "ocean waves crashing" "rocket launch"
#   python bora_run.py push     --runs ... --repo Ami1ne/bora-v1          (token from HF_TOKEN or ~/.cache/huggingface/token)
#
# Run train inside tmux so it survives disconnects. It resumes from <runs>/<name>/ckpt_last.pt.
import argparse
import json
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path

os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60")
os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "60")
sys.path.insert(0, str(Path(__file__).resolve().parent))

REPO_ID, REPO_TYPE = "Ami1ne/Bora", "dataset"
LAT_DIR, LAT_META_DIR, LAT_MERGED = "latents_256", "latents_256_manifests", "latents_256_manifest.csv"
LABELS_DIR = "text_t5v11base_32"
LABEL_FILES = ["latents_256_manifest_labeled.csv", "captions.json", "text_t5v11base_32.safetensors"]
VAE_REPO, VAE_SUBFOLDER = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "vae"
SAMPLE_CAPTIONS = ("ocean waves crashing", "fog rolling through valley", "motorcycle overtaking car", "rocket launch")


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


def with_retries(fn, what, log=print, tries=6, base=15, cap=600):
    for i in range(1, tries + 1):
        try:
            return fn()
        except KeyboardInterrupt:
            raise
        except Exception as e:
            st = getattr(getattr(e, "response", None), "status_code", None)
            if st in (401, 403, 404) or i == tries:
                raise
            wait = min(cap, base * 2 ** (i - 1)) if st != 429 else max(300, base * 2 ** (i - 1))
            log(f"  retry {i}/{tries - 1} for {what}: {type(e).__name__}: {str(e)[:160]} -> sleep {wait}s")
            time.sleep(wait)


# ----------------------------------------------------------------------------- download

def cmd_download(a):
    from huggingface_hub import HfApi, snapshot_download, hf_hub_download
    from huggingface_hub.hf_api import RepoFile
    data = Path(a.data); data.mkdir(parents=True, exist_ok=True)
    log = log_to(data / "download.log")
    api = HfApi()
    remote = {}
    for folder in (LAT_DIR, LABELS_DIR):
        remote.update({it.path: it.size for it in with_retries(lambda: list(api.list_repo_tree(REPO_ID, path_in_repo=folder, repo_type=REPO_TYPE)), f"list {folder}", log)
                       if isinstance(it, RepoFile)})
    remote = {p: s for p, s in remote.items() if p.endswith((".safetensors", ".csv", ".json"))}
    missing = lambda: [p for p, s in remote.items() if not (data / p).exists() or (data / p).stat().st_size != s]
    GB = 1024 ** 3
    todo = missing()
    log(f"{len(remote)} files on the Hub, {sum(remote.values()) / GB:.1f} GB | to download: {len(todo)} files, {sum(remote[p] for p in todo) / GB:.1f} GB")
    t0 = time.time()
    if todo:
        with_retries(lambda: snapshot_download(REPO_ID, repo_type=REPO_TYPE, local_dir=str(data), max_workers=16,
                                               allow_patterns=[f"{LAT_DIR}/*", f"{LAT_META_DIR}/*", LAT_MERGED, f"{LABELS_DIR}/*"]),
                     "snapshot_download", log, tries=5, base=20)
    for p in missing():
        log(f"size mismatch or missing: {p} -> downloading again")
        with_retries(lambda: hf_hub_download(REPO_ID, p, repo_type=REPO_TYPE, local_dir=str(data), force_download=True), f"download {p}", log, tries=4, base=20)
    assert not missing(), f"still wrong: {missing()[:5]}"
    for f in LABEL_FILES:
        shutil.copy(data / LABELS_DIR / f, data / f)
    log(f"ok in {time.time() - t0:.0f}s: {sum(1 for p in remote if p.startswith(LAT_DIR))} shards + label files in {data}")


# ----------------------------------------------------------------------------- VAE decode -> mp4

class Decoder:
    def __init__(self, data):
        import torch
        from safetensors import safe_open
        with safe_open(str(next((Path(data) / LAT_DIR).glob("shard_*.safetensors"))), framework="pt") as f:
            meta = f.metadata()
        self.mean = torch.tensor(json.loads(meta["latents_mean"])).view(1, -1, 1, 1, 1)
        self.std = torch.tensor(json.loads(meta["latents_std"])).view(1, -1, 1, 1, 1)
        self.vae = None

    def decode(self, lat):
        import numpy as np
        import torch
        if self.vae is None:
            from diffusers import AutoencoderKLWan
            self.vae = AutoencoderKLWan.from_pretrained(VAE_REPO, subfolder=VAE_SUBFOLDER, torch_dtype=torch.float32).eval()
        vae = self.vae.to("cuda")
        out = []
        with torch.inference_mode():
            for z in lat:
                x = vae.decode(z[None].float().cuda() * self.std.cuda() + self.mean.cuda()).sample.float()
                out.append(((x[0] + 1) * 127.5).round().clamp(0, 255).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy())
        vae.to("cpu"); torch.cuda.empty_cache()
        return np.stack(out)

    @staticmethod
    def write_mp4(frames, path, fps=8):
        import numpy as np
        both = np.concatenate(list(frames), axis=2)
        h, w = both.shape[1:3]
        subprocess.run(["ffmpeg", "-nostdin", "-v", "error", "-y", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{w}x{h}", "-r", str(fps), "-i", "-",
                        "-c:v", "libx264", "-crf", "16", "-pix_fmt", "yuv420p", str(path)], input=both.tobytes(), check=True, timeout=300)
        return path

    def decode_fn(self, out_dir, log=print):
        def fn(lat, prompts, step):
            try:
                p = self.write_mp4(self.decode(lat), Path(out_dir) / "samples" / f"step{step}.mp4")
                log(f"  samples at step {step}: {p.name} | {[x[:40] for x in prompts]}")
            except Exception as e:                       # a failed preview must never kill the run
                log(f"  sample decode failed at step {step}: {type(e).__name__}: {e}")
        return fn


# ----------------------------------------------------------------------------- overfit / train

def cmd_overfit(a):
    import csv
    from safetensors.torch import load_file
    from bora_train import TrainConfig, train
    out = Path(a.runs) / "overfit"
    shutil.rmtree(out, ignore_errors=True); out.mkdir(parents=True)
    log = log_to(out / "train.log")
    dec = Decoder(a.data)
    cfg = TrainConfig(data_root=a.data, out_dir=str(out), preset=a.preset, model=json.loads(a.model), overfit_batch=a.batch, max_steps=a.steps,
                      warmup=50, lr=3e-4, min_lr_frac=1.0, p_uncond=0.0, p_cond=0.5, num_workers=0, log_every=25, val_every=0,
                      sample_every=a.steps, sample_steps=8, sample_cfg=1.0, save_every=0, ema=0.99, compile=a.compile)
    train(cfg, dec.decode_fn(out, log), log=log)
    truth = load_file(str(out / "overfit_batch.safetensors"))["latents"][:4]
    dec.write_mp4(dec.decode(truth), out / "samples" / "truth.mp4")
    rows = list(csv.DictReader(open(out / "log.csv")))
    first, last = float(rows[0]["loss"]), float(rows[-1]["loss"])
    log(f"OVERFIT RESULT: loss {first:.3f} at step {rows[0]['step']} -> {last:.3f} at step {rows[-1]['step']} | "
        f"{float(rows[-1]['samples_per_s']):.1f} samples/s | videos: {out / 'samples'}/step{a.steps}.mp4 and truth.mp4")


def cmd_train(a):
    from bora_train import TrainConfig, train
    out = Path(a.runs) / a.name
    out.mkdir(parents=True, exist_ok=True)
    log = log_to(out / "train.log")
    dec = Decoder(a.data)
    cfg = TrainConfig(data_root=a.data, out_dir=str(out), preset=a.preset, model=json.loads(a.model), batch_size=a.batch, max_steps=a.steps,
                      lr=a.lr, warmup=a.warmup, num_workers=a.workers, grad_checkpoint=a.grad_checkpoint, compile=a.compile,
                      log_every=50, val_every=1000, sample_every=2000, save_every=1000, sample_captions=SAMPLE_CAPTIONS)
    train(cfg, dec.decode_fn(out, log), log=log)


# ----------------------------------------------------------------------------- sample / push

def cmd_sample(a):
    import torch
    from transformers import AutoTokenizer, T5EncoderModel
    from bora_train import load_ema_model, load_null_text
    from bora_dit import sample_video
    out = Path(a.runs) / a.name
    ema = Path(a.ema) if a.ema else sorted(out.glob("ema_step*.safetensors"), key=lambda p: int(p.stem[8:]))[-1]
    model = load_ema_model(ema, "cuda")
    tok = AutoTokenizer.from_pretrained("google/t5-v1_1-base")
    enc = T5EncoderModel.from_pretrained("google/t5-v1_1-base", torch_dtype=torch.float32).cuda().eval()
    with torch.inference_mode():
        t = tok(a.prompts, max_length=32, padding="max_length", truncation=True, return_tensors="pt").to("cuda")
        text = enc(input_ids=t.input_ids, attention_mask=t.attention_mask).last_hidden_state * t.attention_mask[..., None]
    del enc; torch.cuda.empty_cache()
    nt, nm = load_null_text(Path(a.data) / "text_t5v11base_32.safetensors")
    B = len(a.prompts)
    g = torch.Generator(device="cuda").manual_seed(a.seed)
    lat = sample_video(model, text, t.attention_mask.bool(), nt[None].expand(B, -1, -1).cuda(), nm[None].expand(B, -1).cuda(),
                       n_blocks=3, steps=a.sample_steps, cfg_scale=a.cfg, shift=model.cfg.time_shift, generator=g).float().cpu()
    dec = Decoder(a.data)
    p = dec.write_mp4(dec.decode(lat), out / "samples" / f"prompts_{ema.stem}_s{a.seed}.mp4")
    print(p)


def cmd_push(a):
    from huggingface_hub import HfApi, CommitOperationAdd
    out = Path(a.runs) / a.name
    api = HfApi()
    print("logged in as", api.whoami()["name"])
    api.create_repo(a.repo, repo_type="model", exist_ok=True)
    ema = sorted(out.glob("ema_step*.safetensors"), key=lambda p: int(p.stem[8:]))[-1]
    here = Path(__file__).resolve().parent
    files = {ema.name: ema, "ema_latest.safetensors": ema, "config.json": out / "config.json", "log.csv": out / "log.csv",
             "val.csv": out / "val.csv", "train.log": out / "train.log"}
    for n in ("bora_dit.py", "bora_dataset.py", "bora_train.py", "bora_run.py"):
        files[n] = here / n
    for p in sorted((out / "samples").glob("*.mp4"), key=lambda p: p.stat().st_mtime)[-4:]:
        files[f"samples/{p.name}"] = p
    ops = [CommitOperationAdd(path_in_repo=k, path_or_fileobj=str(v)) for k, v in files.items() if v.exists()]
    api.create_commit(a.repo, repo_type="model", operations=ops, commit_message=f"{a.name}: {ema.stem}")
    print(f"ok: https://huggingface.co/{a.repo}")


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    def common(p, runs=True):
        p.add_argument("--data", default="/ephemeral/bora_data")
        if runs:
            p.add_argument("--runs", default="/ephemeral/bora_runs")
            p.add_argument("--name", default="v1")
    p = sub.add_parser("download"); common(p, runs=False)
    p = sub.add_parser("overfit"); common(p)
    p.add_argument("--preset", default="v1"); p.add_argument("--model", default="{}")
    p.add_argument("--batch", type=int, default=8); p.add_argument("--steps", type=int, default=600); p.add_argument("--compile", action="store_true")
    p = sub.add_parser("train"); common(p)
    p.add_argument("--preset", default="v1"); p.add_argument("--model", default="{}")
    p.add_argument("--batch", type=int, default=32); p.add_argument("--steps", type=int, default=50000)
    p.add_argument("--lr", type=float, default=1e-4); p.add_argument("--warmup", type=int, default=1000)
    p.add_argument("--workers", type=int, default=8); p.add_argument("--grad-checkpoint", action="store_true"); p.add_argument("--compile", action="store_true")
    p = sub.add_parser("sample"); common(p)
    p.add_argument("--prompts", nargs="+", required=True); p.add_argument("--ema", default=None)
    p.add_argument("--sample-steps", type=int, default=8); p.add_argument("--cfg", type=float, default=4.0); p.add_argument("--seed", type=int, default=0)
    p = sub.add_parser("push"); common(p); p.add_argument("--repo", default="Ami1ne/bora-v1")
    a = ap.parse_args()
    {"download": cmd_download, "overfit": cmd_overfit, "train": cmd_train, "sample": cmd_sample, "push": cmd_push}[a.cmd](a)


if __name__ == "__main__":
    main()