# backend.py
from fastapi import FastAPI, UploadFile, Form
from pydantic import BaseModel
from typing import List, Literal
import os
import shutil
import torch
import numpy as np
from PIL import Image
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.utils import render_utils, postprocessing_utils

app = FastAPI()

MAX_SEED = np.iinfo(np.int32).max
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
os.makedirs(TMP_DIR, exist_ok=True)

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


@app.on_event("startup")
def preload_model():
    try:
        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))  # Preload rembg
    except:
        pass


@app.post("/image-to-3d")
async def image_to_3d(
    image: UploadFile,
    seed: int = Form(...),
    ss_guidance_strength: float = Form(...),
    ss_sampling_steps: int = Form(...),
    slat_guidance_strength: float = Form(...),
    slat_sampling_steps: int = Form(...),
    req_session: str = Form(...)
):
    user_dir = os.path.join(TMP_DIR, req_session)
    os.makedirs(user_dir, exist_ok=True)
    image_data = Image.open(image.file)

    outputs = pipeline.run(
        image_data,
        seed=seed,
        formats=["gaussian", "mesh"],
        sparse_structure_sampler_params={
            "steps": ss_sampling_steps,
            "cfg_strength": ss_guidance_strength,
        },
        slat_sampler_params={
            "steps": slat_sampling_steps,
            "cfg_strength": slat_guidance_strength,
        },
    )

    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
    video_path = os.path.join(user_dir, 'sample.mp4')
    render_utils.save_video(video, video_path)
    
    torch.cuda.empty_cache()
    return {"video_path": video_path}


@app.post("/extract-glb")
async def extract_glb(
    mesh_simplify: float = Form(...),
    texture_size: int = Form(...),
    req_session: str = Form(...),
):
    user_dir = os.path.join(TMP_DIR, req_session)
    glb_path = os.path.join(user_dir, 'sample.glb')
    postprocessing_utils.export_glb(glb_path, simplify=mesh_simplify, texture_size=texture_size)

    torch.cuda.empty_cache()
    return {"glb_path": glb_path}