import torch
import base64
import io
import tempfile
from PIL import Image
# CORRECTED IMPORT: Package is 'trellis2', Class is 'Trellis2...'
from trellis2.pipelines import Trellis2ImageTo3DPipeline

class EndpointHandler:
    def __init__(self, path=""):
        # Load the official Microsoft Trellis 2 Model
        print("Loading Trellis 2 (4B) Model...")
        self.pipeline = Trellis2ImageTo3DPipeline.from_pretrained(
            "microsoft/TRELLIS.2-4B",
            torch_dtype=torch.float16,
            use_safetensors=True
        )
        self.pipeline.cuda()
        print("Trellis 2 Loaded!")

    def __call__(self, data):
        """
        Input: {"inputs": "base64_string"}
        Output: {"glb": "base64_string"}
        """
        # 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
        # Trellis 2 returns a list of Mesh objects directly
        outputs = self.pipeline.run(image, seed=42)
        mesh_result = outputs[0]
        
        # 3. Export to GLB
        # The V2 mesh object has a direct export method
        with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmp:
            # export() takes a filepath string
            mesh_result.export(tmp.name)
            
            # Read back the bytes
            with open(tmp.name, "rb") as f:
                glb_bytes = f.read()

        # 4. Return Base64
        out_b64 = base64.b64encode(glb_bytes).decode('utf-8')
        return {"glb": out_b64}