import argparse
import os
import re

os.environ.setdefault("SPCONV_ALGO", "native")
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import numpy as np

from trellis.pipelines import (
    DVDTextToVoxelPipeline,
    TrellisTextTo3DPipeline,
    export_cubified_voxels,
    run_text_stage2_from_dvd_voxels,
)
from trellis.utils import postprocessing_utils


def slugify(text: str, max_length: int = 48) -> str:
    slug = re.sub(r"[^a-zA-Z0-9]+", "_", text.lower()).strip("_")
    return (slug[:max_length].strip("_") or "text_prompt")


def parse_args():
    parser = argparse.ArgumentParser(description="DVD text-to-voxel generation followed by TRELLIS stage 2.")
    parser.add_argument("--prompt", required=True, help="Input text prompt.")
    parser.add_argument("--name", default=None, help="Output asset name. Defaults to a slug of the prompt.")
    parser.add_argument("--dvd-config", default="ckpts/dvd_text.json", help="DVD text model config JSON.")
    parser.add_argument("--dvd-checkpoint", default="ckpts/dvd_text.safetensors", help="DVD text safetensors checkpoint.")
    parser.add_argument("--trellis-model", default="microsoft/TRELLIS-text-large", help="TRELLIS text pipeline.")
    parser.add_argument("--output-dir", default="example_results", help="Directory for generated assets.")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--dvd-steps", type=int, default=256)
    parser.add_argument("--stage2-steps", type=int, default=25)
    parser.add_argument("--stage2-cfg", type=float, default=5.0)
    parser.add_argument("--skip-stage2", action="store_true", help="Only generate and save DVD voxels.")
    return parser.parse_args()


def main():
    args = parse_args()
    os.makedirs(args.output_dir, exist_ok=True)

    name = args.name or slugify(args.prompt)

    dvd = DVDTextToVoxelPipeline.from_files(
        args.dvd_config,
        args.dvd_checkpoint,
        device=args.device,
    )
    voxels = dvd.sample_voxels(
        args.prompt,
        seed=args.seed,
        steps=args.dvd_steps,
    )

    voxel_coords_path = os.path.join(args.output_dir, f"voxel64_{name}_dis.npy")
    np.save(voxel_coords_path, voxels.coords_without_batch.numpy())
    export_cubified_voxels(voxels, os.path.join(args.output_dir, f"{name}_dvd_voxels.glb"))
    print(f"Saved DVD coords: {voxel_coords_path}")

    if args.skip_stage2:
        return

    trellis = TrellisTextTo3DPipeline.from_pretrained(args.trellis_model)
    trellis.to(args.device)
    outputs = run_text_stage2_from_dvd_voxels(
        trellis,
        args.prompt,
        voxels,
        seed=args.seed,
        slat_sampler_params={
            "steps": args.stage2_steps,
            "cfg_strength": args.stage2_cfg,
        },
    )

    glb = postprocessing_utils.to_glb(outputs["gaussian"][0], outputs["mesh"][0])
    glb_path = os.path.join(args.output_dir, f"{name}.glb")
    glb.export(glb_path)
    print(f"Saved GLB: {glb_path}")


if __name__ == "__main__":
    main()