from __future__ import annotations
import argparse, json, os, sys, time, gc
from pathlib import Path

os.environ.setdefault("ATTN_BACKEND", "sdpa")

import torch

MODEL_ID = "VAST-AI/AniGen"


class SSDecoderExport(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, z, z_skl):
        occ, occ_skl = self.model(z, z_skl)
        return occ, occ_skl


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=os.environ.get("CF_ONNX_OUT", "/tmp/cf-ss-decoder"))
    ap.add_argument("--model-root", default=os.environ.get("ANIGEN_MODEL_ROOT", "/tmp/anigen-model"))
    ap.add_argument("--app-root", default=os.environ.get("ANIGEN_APP_ROOT", "/home/user/app"))
    args = ap.parse_args()

    out = Path(args.out)
    model_root = Path(args.model_root)
    app_root = Path(args.app_root)
    out_dir = out / "onnx/anigen/ss-decoder"
    out_dir.mkdir(parents=True, exist_ok=True)
    sys.path.insert(0, str(app_root))

    from huggingface_hub import snapshot_download
    snapshot_download(
        MODEL_ID,
        token=os.environ.get("HF_TOKEN"),
        local_dir=model_root,
        allow_patterns=[
            "ckpts/anigen/ss_dae/config.json",
            "ckpts/anigen/ss_dae/ckpts/decoder_final.pt",
        ],
    )
    os.chdir(model_root)

    from anigen.utils.model_utils import load_decoder
    model = load_decoder(str(model_root / "ckpts/anigen/ss_dae"), "final", "cuda").eval()
    wrapper = SSDecoderExport(model).eval()

    z = torch.zeros((1, 8, 16, 16, 16), device="cuda", dtype=torch.float32)
    z_skl = torch.zeros((1, 4, 16, 16, 16), device="cuda", dtype=torch.float32)
    path = out_dir / "model.onnx"

    started = time.time()
    with torch.inference_mode():
        torch.onnx.export(
            wrapper,
            (z, z_skl),
            str(path),
            input_names=["z", "z_skl"],
            output_names=["occupancy", "occupancy_skl"],
            opset_version=23,
            dynamo=True,
            external_data=True,
        )
    export_s = time.time() - started

    import onnx
    onnx.checker.check_model(str(path))

    meta = {
        "component": "anigen-ss-decoder",
        "source": MODEL_ID,
        "checkpoint": "ckpts/anigen/ss_dae/ckpts/decoder_final.pt",
        "opset": 23,
        "static_profile": {
            "z": [1, 8, 16, 16, 16],
            "z_skl": [1, 4, 16, 16, 16],
            "occupancy": [1, 1, 64, 64, 64],
            "occupancy_skl": [1, 1, 64, 64, 64],
        },
        "export_seconds": round(export_s, 3),
        "torch": torch.__version__,
        "cuda": torch.version.cuda,
        "gpu": torch.cuda.get_device_name(0),
    }
    (out_dir / "export_meta.json").write_text(json.dumps(meta, indent=2))
    print("SS_DECODER_EXPORTED", json.dumps(meta), flush=True)

    del wrapper, model, z, z_skl
    gc.collect(); torch.cuda.empty_cache()


if __name__ == "__main__":
    main()