from typing import Dict, List, Any
import torch
import base64
from io import BytesIO
from diffusers import HunyuanVideoPipeline
from diffusers.utils import export_to_video
import tempfile
import os

# Set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if device.type != 'cuda':
    raise ValueError("HunyuanVideo requires GPU")

# Set mixed precision dtype
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16

class EndpointHandler():
    # Upstream model that actually holds the weights. The HF Inference Toolkit
    # passes the local endpoint repo dir (/repository) as `path`, but that repo
    # only contains handler.py + requirements.txt — no model_index.json. So we
    # always load from the upstream model ID, ignoring the passed-in path.
    #
    # NOTE: use the diffusers-format community repo, NOT tencent/HunyuanVideo —
    # the official repo ships raw .pt checkpoints with no model_index.json and
    # is not loadable via HunyuanVideoPipeline.from_pretrained.
    MODEL_ID = "hunyuanvideo-community/HunyuanVideo"

    def __init__(self, path=""):
        """
        Initialize HunyuanVideo pipeline for video generation.
        Args:
            path: Local endpoint dir from the toolkit (ignored — has no weights).
        """
        # Load HunyuanVideo pipeline from the upstream model ID.
        self.pipe = HunyuanVideoPipeline.from_pretrained(
            self.MODEL_ID,
            torch_dtype=dtype
        )
        self.pipe.to(device)

        # Enable memory optimizations
        self.pipe.enable_model_cpu_offload()
        self.pipe.vae.enable_slicing()
        self.pipe.vae.enable_tiling()

    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
        """
        Generate video from text prompt.
        Args:
            data: Dictionary containing:
                - inputs (str): Text prompt for video generation
                - num_inference_steps (int, optional): Number of denoising steps (default: 50)
                - guidance_scale (float, optional): Guidance scale for generation (default: 6.0)
                - num_frames (int, optional): Number of frames (default: 129)
                - height (int, optional): Video height (default: 544)
                - width (int, optional): Video width (default: 960)
        Returns:
            Dictionary with base64 encoded video
        """
        inputs = data.pop("inputs", data)

        # Hyperparameters
        num_inference_steps = data.pop("num_inference_steps", 50)
        guidance_scale = data.pop("guidance_scale", 6.0)
        num_frames = data.pop("num_frames", 129)  # ~5 seconds at 25 FPS
        height = data.pop("height", 544)
        width = data.pop("width", 960)

        # Generate video
        output = self.pipe(
            prompt=inputs,
            height=height,
            width=width,
            num_frames=num_frames,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            generator=torch.Generator(device="cpu").manual_seed(42),
        )

        # Get video frames
        video = output.frames[0]

        # Export to video file (temporary)
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_file:
            tmp_path = tmp_file.name
            export_to_video(video, tmp_path, fps=25)

        # Read video and encode as base64
        with open(tmp_path, "rb") as video_file:
            video_bytes = video_file.read()
            video_base64 = base64.b64encode(video_bytes).decode('utf-8')

        # Clean up temporary file
        os.unlink(tmp_path)

        return {
            "video": video_base64,
            "format": "mp4",
            "fps": 25,
            "frames": num_frames,
            "resolution": f"{width}x{height}"
        }