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

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

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, width, height):
    # Check file size
    file_size_mb = os.path.getsize(audio_file.name) / (1024 * 1024)
    if file_size_mb > 3:
        raise gr.Error(f"File size is {file_size_mb:.2f}MB, which exceeds the 3MB limit. Please use a file that is 3MB or smaller.")

    api_key = os.getenv('DEFORUM_API_KEY')  # Ensure your environment variable is set
    if not api_key:
        raise gr.Error("API key is not set in environment variables")

    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,
        "width": width,
        "height": height
    }

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

    data = response.json()
    tracking_url = f"{server}{data['links']['audiovis1']}"
    logging.info(f"Job submitted successfully. Tracking URL: {tracking_url}")

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

        tracking_data = response.json()
        logging.info(f"Job status: {tracking_data['status']}")
        if tracking_data['status'] in ['canceled', 'failed', 'succeeded']:
            break
        sleep(10)

    if tracking_data['status'] != 'succeeded':
        raise gr.Error(f"Job ended with status: {tracking_data['status']}")
    else:
        output_url = tracking_data['links']['outputUrls'][0]
        logging.info(f"Job succeeded. Output URL: {output_url}")
        return output_url

description1 = """
# Audio Visualizer Playground
### Easily create audio-synced animation masks for AnimateDiff or Deforum Animations.
"""
description2 = """
#### Please provide feedback 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: #000000; /* Change label color to black */
    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, width, height):
    result = submit_job(audio_file, preset, beat_sensitivity, fps, width, height)
    return result

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown(description1, elem_id="markdown-text")
    gr.Markdown(description2, elem_id="markdown-text")
    with gr.Group():
        with gr.Column():
            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", "rotating-shapes"
            ])
        with gr.Row():
            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, value=24)
        with gr.Row():
            width = gr.Slider(label="Width", minimum=512, maximum=1024, step=1, value=512)
            height = gr.Slider(label="Height", minimum=512, maximum=1024, step=1, value=512)
        submit_button = gr.Button("Submit")
    output_video = gr.Video(label="Output MP4")
    output_error = gr.Textbox(label="Error", visible=False)

    def update_output(video_url):
        if isinstance(video_url, str) and video_url.startswith("Error"):
            output_error.update(value=video_url, visible=True)
            output_video.update(value=None, visible=False)
        else:
            output_error.update(value="", visible=False)
            output_video.update(value=video_url, visible=True)

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

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