import argparse
import json
import os
from pathlib import Path

import numpy as np
import torch
from PIL import Image

os.environ.setdefault("SPCONV_ALGO", "native")
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("DVD_MODEL_REPO", "Zhengrui/dvd")


def cfg_schedule(mode: str, constant: float, early: float, late: float, split: float):
    if mode == "Constant":
        return float(constant)
    if mode == "Two-stage":
        split = float(split)
        early = float(early)
        late = float(late)
        return lambda t: early if t < split else late
    return None


def run_generate(config_path: str):
    with open(config_path, "r") as f:
        cfg = json.load(f)

    from dvd import DVDImageToVoxelPipeline, export_cubified_voxels

    repo = os.environ.get("DVD_MODEL_REPO", "Zhengrui/dvd")
    subfolder = os.environ.get("DVD_MODEL_SUBFOLDER") or None
    revision = os.environ.get("DVD_MODEL_REVISION") or None
    token = os.environ.get("DVD_MODEL_TOKEN") or os.environ.get("HF_TOKEN") or None

    print(f"[DVD Worker] loading DVD image pipeline from {repo}", flush=True)
    pipeline = DVDImageToVoxelPipeline.from_pretrained(
        repo,
        variant="base",
        device="cuda",
        subfolder=subfolder,
        revision=revision,
        token=token,
    )
    print("[DVD Worker] pipeline ready", flush=True)

    image = Image.open(cfg["image_path"]).convert("RGBA")
    sampler_kwargs = {"steps": int(cfg["dvd_steps"])}
    schedule = cfg_schedule(
        cfg["dvd_cfg_mode"],
        float(cfg["dvd_cfg_constant"]),
        float(cfg["dvd_cfg_early"]),
        float(cfg["dvd_cfg_late"]),
        float(cfg["dvd_cfg_split"]),
    )
    if schedule is not None:
        sampler_kwargs["cfg_strength"] = schedule

    print(f"[DVD Worker] sampling seed={cfg['seed']} steps={cfg['dvd_steps']}", flush=True)
    voxels = pipeline.sample_voxels(
        image,
        seed=int(cfg["seed"]),
        preprocess_image=bool(cfg["preprocess_image"]),
        **sampler_kwargs,
    )

    mesh_path = cfg["mesh_path"]
    npy_path = cfg["npy_path"]
    export_cubified_voxels(voxels, mesh_path)
    np.save(npy_path, voxels.coords_without_batch.detach().cpu().numpy().astype(np.int32))
    torch.cuda.empty_cache()
    print(f"[DVD Worker] done mesh={mesh_path} npy={npy_path}", flush=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=["generate"])
    parser.add_argument("config")
    args = parser.parse_args()
    if args.command == "generate":
        run_generate(args.config)


if __name__ == "__main__":
    main()