import gradio as gr
import requests
import io
from PIL import Image as img
from moviepy.editor import ImageSequenceClip
import numpy as np

google_html = '''

<!-- Google tag (gtag.js) -->
<script async src="https://www.googletagmanager.com/gtag/js?id=G-CDY5VGM0FL"></script>
<script>
  window.dataLayer = window.dataLayer || [];
  function gtag(){dataLayer.push(arguments);}
  gtag('js', new Date());

  gtag('config', 'G-CDY5VGM0FL');
</script>

'''


# Setting up the API URL and headers
API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"

def query(payload, api_key):
    headers = {"Authorization": f"Bearer {api_key}"}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content

def generate_video(video_description, api_key):
    # Setting up the list to store images
    image_list = []

    # Fetching the images based on the description
    image_bytes = query({
        "inputs": "Create 9 sequential images with 1:1 aspect ratio for every one of images of a " + video_description,
    }, api_key)

    # Opening the image using PIL
    image = img.open(io.BytesIO(image_bytes))

    # List of coordinates to crop the image into 9 parts
    coordinates = [
        (0, 340, 0, 340),  
        (340, 680, 0, 340), 
        (680, 1020, 0, 340),
        (0, 340, 340, 680),  
        (340, 680, 340, 680), 
        (680, 1020, 340, 680),
        (0, 340, 680, 1020),  
        (340, 680, 680, 1020), 
        (680, 1020, 680, 1020),
    ]

    # Cropping the parts and converting them to numpy arrays and adding them to image_list
    for i, (x1, x2, y1, y2) in enumerate(coordinates):
        cropped_image = image.crop((x1, y1, x2, y2))
        image_array = np.array(cropped_image)  # Converting the image to a numpy array
        image_list.append(image_array)

    # Repeating the images to get more frames for the video
    image_list_extended = image_list * 10  # Repeating 10 times to get a longer video

    # Creating the video from the list using moviepy
    clip = ImageSequenceClip(image_list_extended, fps=2)
    
    # Saving the video to a file
    output_path = "/tmp/output_video.mp4"
    clip.write_videofile(output_path, codec="libx264")

    return output_path

def main(video_description, api_key):
    if not api_key:
        return "Please Enter your Hugging Face API key."
    
    # Creating the video based on the input description
    video_path = generate_video(video_description, api_key)

    # Returning the video download link to the user
    return "The video has been created, you can download it now:", gr.File(video_path)

# Gradio interface
with gr.Interface(fn=main, inputs=[gr.HTML(google_html),gr.Textbox(label="Enter video description"), gr.Textbox(label="Enter your Hugging Face API key")], outputs=[gr.Textbox(), gr.File()]) as demo:
    demo.launch()