# ==============================================================================
# VIU-MINI-DENSE-360M: RTX 5090 / 4090 CLOUD NOTEBOOK MASTER RUNNER
# Works out-of-the-box on: RunPod, Vast.ai, Colab Pro, Lambda, Paperspace, Local PC
# ==============================================================================
import os, sys, shutil, subprocess, glob

# 1. Environment & Auth Setup
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["PYTHONUNBUFFERED"] = "1"
os.environ["HF_TOKEN"] = os.environ.get("HF_TOKEN", "")

print("=" * 70)
print("  🚀 ViuMini-Dense-360M: RTX 5090 High-Throughput Trainer")
print("=" * 70)

# 2. Hardware Audit
import torch
if not torch.cuda.is_available():
    raise RuntimeError("CUDA is not available. Please ensure GPU runtime is enabled!")

gpu_name = torch.cuda.get_device_name(0)
props = torch.cuda.get_device_properties(0)
vram_gb = props.total_memory / (1024**3)
num_gpus = torch.cuda.device_count()
print(f"[GPU] Detected: {gpu_name} ({vram_gb:.1f} GB VRAM x {num_gpus} GPUs)")

# 3. Dependencies
print("\n[1/4] Ensuring required dependencies...")
subprocess.run([
    sys.executable, "-m", "pip", "install", "-q", "--upgrade",
    "transformers", "datasets", "tokenizers", "bitsandbytes>=0.43.0",
    "pyarrow", "pyyaml", "tqdm", "huggingface_hub", "pandas"
], check=True)

# 4. Clone or Pull Latest Codebase
if os.path.exists("/workspace"):
    repo_dir = "/workspace/ViuMini-Dense-360M"
elif os.path.exists("/content"):
    repo_dir = "/content/ViuMini-Dense-360M"
elif os.path.exists("/kaggle/working"):
    repo_dir = "/kaggle/working/ViuMini-Dense-360M"
else:
    repo_dir = os.path.abspath("ViuMini-Dense-360M")

print(f"\n[2/4] Setting up repository at {repo_dir}...")
if not os.path.exists(os.path.join(repo_dir, ".git")):
    subprocess.run(["git", "clone", "https://huggingface.co/ViuAI/ViuMini-Dense-360M", repo_dir], check=True)
else:
    subprocess.run(["git", "-C", repo_dir, "pull"], check=False)

os.chdir(repo_dir)

# 5. Check Tokenizer
tok_path = os.path.join(repo_dir, "tokenizer", "outputs", "tokenizer.json")
if not os.path.exists(tok_path) or os.path.getsize(tok_path) < 100000:
    from huggingface_hub import hf_hub_download
    print("[Tok] Downloading 4.25 MB tokenizer.json from Hub...")
    hf_hub_download(
        repo_id="ViuAI/ViuMini-Dense-360M",
        filename="tokenizer/outputs/tokenizer.json",
        local_dir=repo_dir,
        token=os.environ["HF_TOKEN"]
    )
print(f"[Tok] Tokenizer verified ({os.path.getsize(tok_path):,} bytes)")

# 6. Check Latest Checkpoint on Hugging Face to Auto-Resume
print("\n[3/4] Checking checkpoint history on Hugging Face Hub...")
from huggingface_hub import HfApi
api = HfApi(token=os.environ["HF_TOKEN"])
remote_ckpts = []
try:
    files = api.list_repo_files(repo_id="ViuAI/ViuMini-Dense-360M", repo_type="model")
    for f in files:
        if f.startswith("model/checkpoints/step_") and f.endswith(".pt"):
            st = int(f.split("step_")[1].split(".pt")[0])
            remote_ckpts.append((st, f))
    remote_ckpts.sort(key=lambda x: x[0])
except Exception as e:
    print(f"[HF] Note: {e}")

if remote_ckpts:
    latest_step, latest_file = remote_ckpts[-1]
    print(f"[Resume] Found latest checkpoint on Hub: Step {latest_step} ({latest_file})")
    print(f"[Resume] Model will seamlessly continue from Step {latest_step} onwards!")
else:
    print("[Resume] No previous checkpoint found. Starting from Step 0.")

# 7. Select Hardware Config
cfg_file = "model/configs/train_rtx5090.yaml"
if not os.path.exists(cfg_file):
    cfg_file = "model/configs/train_rtx4090.yaml"

print(f"\n[4/4] Launching High-Throughput Training with {cfg_file}...")
cmd = [
    sys.executable, "model/scripts/train.py",
    "--config", cfg_file,
    "--resume", "auto"
]

# Launch with unbuffered real-time stdout streaming (step 10, 20, 30...)
proc = subprocess.Popen(
    cmd,
    stdout=subprocess.PIPE,
    stderr=subprocess.STDOUT,
    text=True,
    bufsize=1,
    universal_newlines=True
)
for line in iter(proc.stdout.readline, ''):
    print(line, end='', flush=True)
proc.stdout.close()
proc.wait()