import torch
import base64
import io
import tempfile
from PIL import Image
# Note: The import might change to 'trellis.2' or stay 'trellis' depending on how they packaged the V2 repo.
# Standard import usually resolves to the installed package name 'trellis' even for v2.
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.utils import postprocessing_utils

class EndpointHandler:
    def __init__(self, path=""):
        print("Loading Trellis 2 Model...")
        # The CORRECT Model ID for V2 (4 Billion parameters)
        self.pipeline = TrellisImageTo3DPipeline.from_pretrained(
            "microsoft/TRELLIS.2-4B",
            torch_dtype=torch.float16,
            use_safetensors=True
        )
        self.pipeline.cuda()
        print("Trellis 2 Loaded!")

    def __call__(self, data):
        # 1. Parse Input
        inputs = data.pop("inputs", data)
        if isinstance(inputs, dict) and "image" in inputs:
            inputs = inputs["image"]
            
        if isinstance(inputs, str):
            image_data = base64.b64decode(inputs)
            image = Image.open(io.BytesIO(image_data)).convert("RGB")
        else:
            image = inputs

        # 2. Inference
        outputs = self.pipeline.run(image, seed=42, formats=["mesh"])
        
        # 3. Export to GLB
        video_content = postprocessing_utils.to_glb(
            outputs['mesh'][0], 
            simplify=0.95,      
            texture_size=1024
        )
        
        # 4. Handle Output
        if isinstance(video_content, str): 
            with open(video_content, "rb") as f: glb_bytes = f.read()
        elif isinstance(video_content, bytes):
            glb_bytes = video_content
        else:
            with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmp:
                outputs['mesh'][0].export(tmp.name, file_type='glb')
                with open(tmp.name, "rb") as f: glb_bytes = f.read()

        return {"glb": base64.b64encode(glb_bytes).decode('utf-8')}