import tempfile

import numpy as np
import PIL.Image
import torch
import trimesh
from diffusers import ShapEImg2ImgPipeline, ShapEPipeline
from diffusers.utils import export_to_ply, export_to_gif

from typing import Dict, List, Any

class EndpointHandler():
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.pipe = ShapEPipeline.from_pretrained("openai/shap-e", torch_dtype=torch.float16)
        self.pipe.to(self.device)

        self.pipe_img = ShapEImg2ImgPipeline.from_pretrained("openai/shap-e-img2img", torch_dtype=torch.float16)
        self.pipe_img.to(self.device)

    def to_glb(self, ply_path: str) -> str:
        mesh = trimesh.load(ply_path)
        rot = trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])
        mesh = mesh.apply_transform(rot)
        rot = trimesh.transformations.rotation_matrix(np.pi, [0, 1, 0])
        mesh = mesh.apply_transform(rot)
        mesh_path = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
        mesh.export(mesh_path.name, file_type="glb")
        return mesh_path.name
        
    def run_text(self, prompt: str, seed: int = 0, guidance_scale: float = 15.0, num_steps: int = 64, output_type: str = "pil") -> str:
        generator = torch.Generator(device=self.device).manual_seed(seed)
        
        if output_type=="pil":
            images = self.pipe(
            prompt,
            num_images_per_prompt=4,
            generator=generator,
            guidance_scale=guidance_scale,
            num_inference_steps=num_steps,
            frame_size=64,
        ).images
            gif_path = tempfile.NamedTemporaryFile(suffix=".gif", delete=False, mode="w+b")
            gif_path2 = tempfile.NamedTemporaryFile(suffix=".gif", delete=False, mode="w+b")
            gif_path3 = tempfile.NamedTemporaryFile(suffix=".gif", delete=False, mode="w+b")
            gif_path4 = tempfile.NamedTemporaryFile(suffix=".gif", delete=False, mode="w+b")
            export_to_gif(images[0], gif_path.name)
            export_to_gif(images[1], gif_path2.name)
            export_to_gif(images[2], gif_path3.name)
            export_to_gif(images[3], gif_path4.name)
            return gif_path.name, gif_path2.name, gif_path3.name, gif_path4.name
        else:
            images = self.pipe(
            prompt,
            num_images_per_prompt=1,
            generator=generator,
            guidance_scale=guidance_scale,
            num_inference_steps=num_steps,
            frame_size=64,
            output_type=output_type
        ).images
            ply_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False, mode="w+b")
            export_to_ply(images[0], ply_path.name)
            return self.to_glb(ply_path.name)

    def __call__(self, prompt: str, seed: int = 0, guidance_scale: float = 15.0, num_steps: int = 64) -> str:
        generator = torch.Generator(device=self.device).manual_seed(seed)
        
        images = self.pipe(
        prompt,
        num_images_per_prompt=1,
        generator=generator,
        guidance_scale=guidance_scale,
        num_inference_steps=num_steps,
        frame_size=64,
        output_type="mesh"
        ).images
        ply_path = tempfile.NamedTemporaryFile(suffix=".ply", delete=False, mode="w+b")
        export_to_ply(images[0], ply_path.name)
        return self.to_glb(ply_path.name)