import os
import uuid
from flask import jsonify, send_file, request
from main import *
import torch
import io
from skimage import img_as_ubyte
import imageio

def text_to_video_func(prompt, output_path="output_video.mp4"):
    if text_to_video_model is None:
        return "Text-to-Video model not initialized."
    video_frames_list = text_to_video_model(prompt)
    if video_frames_list and hasattr(video_frames_list, 'frames'):
        video_frames = video_frames_list.frames
        export_to_video_pure(video_frames, output_video=output_path)
        return output_path
    return "Video generation failed."

def export_to_video_pure(video_frames, output_video="output_video.mp4", fps=25):
    writer = imageio.get_writer(output_video, fps=fps)
    for frame in video_frames:
        writer.append_data(img_as_ubyte(frame))
    writer.close()

def text_to_video_api():
    data = request.get_json()
    prompt = data.get('prompt')
    if not prompt:
        return jsonify({"error": "Prompt is required"}), 400
    output_file = text_to_video_func(prompt)
    if output_file == "Text-to-Video model not initialized." or output_file == "Video generation failed.":
        return jsonify({"error": "Text to video failed"}), 500
    with open(output_file, 'rb') as f:
        video_content = f.read()
    return send_file(io.BytesIO(video_content), mimetype='video/mp4', as_attachment=True, download_name="output_video.mp4")