#!/usr/bin/env python3
"""Smoke Faz3 TRELLIS forward LoRA then optional Paint LoRA. Uses HF Preprocessed-4K."""
from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--paint", action="store_true", help="Ayrica faz3_paint smoke")
    parser.add_argument("--limit", type=int, default=4)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--lora-rank", type=int, default=8)
    args = parser.parse_args()

    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if not token:
        print("HF_TOKEN gerekli")
        return 1

    os.environ.setdefault("ATTN_BACKEND", "xformers")
    os.environ.setdefault("SPARSE_ATTN_BACKEND", "xformers")
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

    # Re-enter train_pipeline CLI
    cmd_geo = [
        sys.executable,
        str(ROOT / "train_pipeline.py"),
        "--mode",
        "faz3",
        "--epochs",
        str(args.epochs),
        "--limit",
        str(args.limit),
        "--lora-rank",
        str(args.lora_rank),
        "--checkpoint-every",
        "20",
        "--validation-every",
        "50",
    ]
    print("RUN", " ".join(cmd_geo))
    import subprocess

    env = os.environ.copy()
    env["HF_TOKEN"] = token
    r = subprocess.call(cmd_geo, cwd=str(ROOT), env=env)
    if r != 0:
        return r
    ckpt = ROOT / "checkpoints" / "latest_faz3_lora.pt"
    if not ckpt.is_file():
        print("FAIL: latest_faz3_lora.pt yok")
        return 1
    print("OK geo ckpt", ckpt, ckpt.stat().st_size)

    if args.paint:
        cmd_p = [
            sys.executable,
            str(ROOT / "train_pipeline.py"),
            "--mode",
            "faz3_paint",
            "--epochs",
            str(args.epochs),
            "--limit",
            str(min(2, args.limit)),
            "--lora-rank",
            str(args.lora_rank),
            "--checkpoint-every",
            "5",
        ]
        print("RUN", " ".join(cmd_p))
        r2 = subprocess.call(cmd_p, cwd=str(ROOT), env=env)
        if r2 != 0:
            return r2
        ckpt2 = ROOT / "checkpoints" / "latest_faz3_paint_lora.pt"
        print("OK paint ckpt", ckpt2, ckpt2.stat().st_size if ckpt2.is_file() else None)
    print("SMOKE OK")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())