import base64
import io
import torch
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.utils import export_to_video

class EndpointHandler:
    def __init__(self, path=""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.bfloat16

        # Load VAE and pipeline from the diffusers-format repo
        vae = AutoencoderKLWan.from_pretrained(
            "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
            subfolder="vae",
            torch_dtype=torch.float32,
        )
        self.pipe = WanPipeline.from_pretrained(
            "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
            vae=vae,
            torch_dtype=self.dtype,
        )
        self.pipe.to(self.device)

    def __call__(self, data):
        prompt = data.pop("inputs", "")
        params = data.pop("parameters", {})

        num_frames = params.get("num_frames", 49)
        height = params.get("height", 480)
        width = params.get("width", 832)
        num_inference_steps = params.get("num_inference_steps", 30)
        fps = params.get("fps", 24)
        guidance_scale = params.get("guidance_scale", 5.0)
        negative_prompt = params.get("negative_prompt", "low quality, blurry, distorted")

        frames = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_frames=num_frames,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
        ).frames[0]

        # Export to video in memory
        buf = io.BytesIO()
        export_to_video(frames, buf, fps=fps)
        video_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

        duration = len(frames) / fps
        return {
            "video": video_b64,
            "format": "mp4",
            "frames": len(frames),
            "duration": duration,
        }