import gradio as gr
import requests
import base64
import json
import os
from time import sleep

def file_to_base64(filepath):
    with open(filepath, "rb") as file:
        file_data = file.read()
        base64_encoded_data = base64.b64encode(file_data)
        base64_message = base64_encoded_data.decode('utf-8')
    return base64_message

def submit_job(audio_file, preset, beat_sensitivity, fps):
    api_key = os.getenv('DEFORUM_API_KEY')  
    if not api_key:
        return "Error: API key is not set in environment variables", None

    server = "https://deforum.studio"
    api_url = f'{server}/api/public/v1/audiovis1'
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }

    audio_base64 = file_to_base64(audio_file.name)

    payload = {
        "audioData": audio_base64,
        "presetName": preset,
        "beatSensitivity": beat_sensitivity,
        "fps": fps
    }

    response = requests.post(api_url, headers=headers, json=payload)
    if response.status_code != 201:
        return f"Error submitting job: {response.status_code} - {response.text}", None

    data = response.json()
    tracking_url = f"{server}{data['links']['audiovis1']}"

    while True:
        response = requests.get(tracking_url, headers=headers)
        if response.status_code != 200:
            return f"Error getting job status: {response.text}", None

        tracking_data = response.json()
        if tracking_data['status'] in ['canceled', 'failed', 'succeeded']:
            break
        sleep(10)

    if tracking_data['status'] != 'succeeded':
        return f"Job ended with status: {tracking_data['status']}", None
    else:
        return tracking_data['links']['outputUrls'][0]

description1 = """
# Audio Visualizer Playground
A way for you to easily create audio-synced animation masks for AnimateDiff or Deforum Animations.
"""
description2 = """
### Please provide feedback for us about which masks you like, or would like to see in our animation channel in our Discord: [discord.gg/deforum](https://discord.gg/deforum)
"""

custom_css = """
@import url('//fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700&display=swap');
@import url('//fonts.googleapis.com/css?family=Aldrich&display=swap');
@import-normalize;

html, body {
    height: 100%;
    margin: 0;
    font-family: 'Open Sans', sans-serif;
    background-color: #121212;
    color: #e0e0e0;
    text-align: center;
}

.gradio-container {
    max-width: 900px;
    margin: 0 auto;
    padding: 20px;
    border-radius: 10px;
    box-shadow: 0 0 20px rgba(0, 0, 0, 0.5);
    background-color: #1e1e1e;
}

h1, h2, h3, h4, h5, h6 {
    font-family: 'Aldrich', sans-serif;
    color: #ffffff;
    text-align: center;
}

label {
    font-weight: bold;
    color: #bbbbbb;
    text-align: center;
}

.gr-button {
    background-color: #ff5722 !important;
    color: #fff !important;
    border: none !important;
    padding: 10px 20px !important;
    border-radius: 5px !important;
    cursor: pointer !important;
    transition: background-color 0.3s ease !important;
}

.gr-button:hover {
    background-color: #e64a19 !important;
}

.gr-box, .gr-box-labeled, .gr-file, input[type=text], textarea, .gr-input, .gr-textbox input {
    background-color: #444444 !important;
    color: #e0e0e0 !important;
    border: 1px solid #333333 !important;
    border-radius: 5px !important;
    padding: 10px !important;
}

.gr-file {
    display: flex;
    align-items: center;
    justify-content: center;
    height: 60px;  /* Adjusted height */
    border: 2px dashed #ff5722 !important;
    transition: background-color 0.3s ease !important;
}

.gr-file:hover {
    background-color: #555555 !important;
}

video, .gr-video {
    border-radius: 5px !important;
}

.markdown-text {
    color: #e0e0e0 !important;
}
"""

def main(audio_file, preset, beat_sensitivity, fps):
    video_url = submit_job(audio_file, preset, beat_sensitivity, fps)
    return video_url

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown(description1, elem_id="markdown-text")
    gr.Markdown(description2, elem_id="markdown-text")
    audio_file = gr.File(label="Upload MP3 File", file_types=['audio'])
    preset = gr.Dropdown(label="Select Visualization Preset", choices=[
        "pulsing-circles", "pulsing-hexagons", "pulsing-squares", "pulsing-triangles", 
        "rewbs-01", "rewbs-02", "rotating-shapes"
    ])
    beat_sensitivity = gr.Slider(label="Beat Sensitivity", minimum=1, maximum=5, step=0.1, value=3)
    fps = gr.Slider(label="FPS", minimum=12, maximum=24, step=1)
    submit_button = gr.Button("Submit")
    output_video = gr.Video(label="Output MP4")

    def update_output(video_url):
        output_video.update(value=video_url, visible=True)

    submit_button.click(main, inputs=[audio_file, preset, beat_sensitivity, fps], outputs=[output_video])

if __name__ == "__main__":
    demo.launch()