#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import argparse
from PIL import Image

from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.utils import render_utils, postprocessing_utils

def build_parser():
    p = argparse.ArgumentParser("TRELLIS CLI: single image -> 3D")

    p.add_argument("-i", "--image", required=True, help="Input image path")
    p.add_argument("-o", "--out_dir", default="trellis_out", help="Output directory")

    p.add_argument("--seed", type=int, default=1)
    p.add_argument("--steps_sparse", type=int, default=12)
    p.add_argument("--cfg_sparse", type=float, default=7.5)
    p.add_argument("--steps_slat", type=int, default=12)
    p.add_argument("--cfg_slat", type=float, default=3.0)

    p.add_argument("--simplify", type=float, default=0.95)
    p.add_argument("--texture_size", type=int, default=1024)

    # Export GLB (default True)
    p.add_argument("--export_glb", dest="export_glb", action="store_true", default=True)
    p.add_argument("--no-export_glb", dest="export_glb", action="store_false")

    # Save PLY (default True)
    p.add_argument("--save_ply", dest="save_ply", action="store_true", default=True)
    p.add_argument("--no-save_ply", dest="save_ply", action="store_false")

    # Save videos (default False, plus explicit toggle)
    p.add_argument("--save_video", dest="save_video", action="store_true", default=True)
    p.add_argument("--no-save_video", dest="save_video", action="store_false")

    p.add_argument("--fps", type=int, default=30)
    p.add_argument("--video_gs_name", type=str, default="sample_gs.mp4")
    p.add_argument("--video_rf_name", type=str, default="sample_rf.mp4")
    p.add_argument("--video_mesh_name", type=str, default="sample_mesh.mp4")

    return p

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

    # Optional env
    os.environ.setdefault("SPCONV_ALGO", "native")

    pipeline = TrellisImageTo3DPipeline.from_pretrained("microsoft/TRELLIS-image-large")
    pipeline.cuda()

    image = Image.open(args.image)

    outputs = pipeline.run(
        image,
        seed=args.seed,
        sparse_structure_sampler_params={
            "steps": args.steps_sparse,
            "cfg_strength": args.cfg_sparse,
        },
        slat_sampler_params={
            "steps": args.steps_slat,
            "cfg_strength": args.cfg_slat,
        },
    )

    if args.save_video:
        import imageio

        video = render_utils.render_video(outputs["gaussian"][0])["color"]
        imageio.mimsave(os.path.join(args.out_dir, args.video_gs_name), video, fps=args.fps)

        video = render_utils.render_video(outputs["radiance_field"][0])["color"]
        imageio.mimsave(os.path.join(args.out_dir, args.video_rf_name), video, fps=args.fps)

        video = render_utils.render_video(outputs["mesh"][0])["normal"]
        imageio.mimsave(os.path.join(args.out_dir, args.video_mesh_name), video, fps=args.fps)

    if args.export_glb:
        glb = postprocessing_utils.to_glb(
            outputs["gaussian"][0],
            outputs["mesh"][0],
            simplify=args.simplify,
            texture_size=args.texture_size,
        )
        glb.export(os.path.join(args.out_dir, "sample.glb"))

    if args.save_ply:
        outputs["gaussian"][0].save_ply(os.path.join(args.out_dir, "sample.ply"))

if __name__ == "__main__":
    main()