import gradio as gr
import json
from moviepy.video.io.VideoFileClip import VideoFileClip
import threading
import time
global_lock = threading.Lock()
ip_records = {}

def get_video_duration(video_path):
    clip = VideoFileClip(video_path)
    duration = clip.duration
    clip.close()
    return duration

from utils import processing

with open("sorted.json", "r") as f:
    style_dicts = json.load(f)

styles = list(style_dicts.keys())
style_images = {k: v['thumbnailUrl'] for k, v in style_dicts.items() if v is not None}

# results_map = {
#     "test32.mp4ClayStyle19:16": "test32.mp4",
#     "test32.mp43D P19:16": "test32.mp4",
#     }
def get_client_info():
    request = gr.Request()
    clientHost = request
    print("IP address:", clientHost)
    return clientHost

def check_ip_limit(client_ip):
    current_time = time.time()
    if client_ip not in ip_records:
        ip_records[client_ip] = [current_time]
        return True
    else:
        times = ip_records[client_ip]
        times = [t for t in times if current_time - t < 5 * 60]
        if len(times) < 2:
            times.append(current_time)
            ip_records[client_ip] = times
            return True
        else:
            return False

def fn(in_video, style, aspect_ratio, prompt):

    print(f"Received video: {in_video}, aspect_ratio: {aspect_ratio}, style: {style}")
    
    if in_video is None:
        raise Exception("input_video is None")
    # input_str = "".join([str(item) for item in [in_video, style, aspect_ratio]])
    # if input_str in results_map:
    #     return results_map[input_str]

    if aspect_ratio not in ["16:9", "9:16", "1:1"]:
        raise Exception('aspect_ratio not in ["16:9", "9:16", "1:1"]')
    if aspect_ratio == "16:9":
        width, height = 1280, 720
    elif aspect_ratio == "9:16":
        width, height = 720, 1280
    elif aspect_ratio == "1:1":
        width, height = 960, 960

    style_id = style_dicts[style]['id']
    if style not in styles:
        raise Exception("style not in styles")

    video_duration = get_video_duration(in_video)
    if video_duration < 3:
        raise Exception("video_duration < 3")

    transfer_duration = 3

    res = None
    try:
        global_lock.acquire()
        res = processing(in_video, width, height, style_id, transfer_duration, prompt)
        if not res:
            raise Exception("Processing failed")
    finally:
        global_lock.release()

    return res

def load_description(fp):
    with open(fp, 'r', encoding='utf-8') as f:
        content = f.read()
    return content

with gr.Blocks(theme=gr.themes.Ocean()) as demo:
    # Transform your videos with different Anime styles and have fun!
    gr.HTML(load_description("title.md"))
    with gr.Tab("Video Processing"):
        with gr.Row():
            in_video = gr.Video(label="1. Input Video (Please provide a clear video clip)", interactive=True)

            aspect_ratio = gr.Radio(["16:9", "9:16", "1:1"], label="3. Aspect Ratio(Width:Height)", value="16:9", interactive=True)

            prompt = gr.Textbox(label="4. Prompt", value="a character with cute face")


            out_video = gr.Video(label="Final Output Video")


        with gr.Row():
            style = gr.Radio(
                choices=styles,
                label="2. Style, for more styles, please join our [Discord](https://discord.gg/cN8geBsBmu)",
                value="ClayStyle1",
                interactive=True
            )
            preview = gr.Image(
                value=style_images["ClayStyle1"],
                label="Style Preview",
                show_label=True,
                interactive=False
            )

            def update_preview(choice):
                return style_images[choice]

            style.change(fn=update_preview, inputs=style, outputs=preview)

        submit_button = gr.Button("Generate", interactive=True, variant="primary")

        examples = gr.Examples(
            [
                ["test32.mp4", "ClayStyle1", "9:16", "a girl with cute face is walking"],
                ["test32.mp4", "3D P1", "9:16", "a girl with cute face is walking" ],
                ["test32.mp4", "FlatanimeV4-3", "9:16", "a girl with cute face is walking" ],
                ["test32.mp4", "FlatanimeV4 Delicated", "9:16", "a girl with cute face is walking" ],
                ["test32.mp4", "illustrationV4-1", "9:16", "a girl with cute face is walking" ],
                ["test32.mp4", "Child Crayon Drawing", "9:16", "a girl with cute face is walking" ],
                ["test32.mp4", "Pixels", "9:16", "a girl with cute face is walking" ],
                ["test32.mp4", "Color Ink Painting", "9:16", "a girl with cute face is walking" ],
                ["test32.mp4", "cutefaceV4 style1", "9:16", "a girl with cute face is walking" ],
                ["test32.mp4", "flatanime-gibli", "9:16", "a girl with cute face is walking" ],
            ],
            inputs=[in_video, style, aspect_ratio, prompt],
            outputs=[out_video],
            fn=fn,
        )

        submit_button.click(
            fn,
            inputs=[in_video, style, aspect_ratio, prompt],
            outputs=[out_video],
        )
    gr.HTML(load_description("end.md"))


if __name__ == "__main__":
    demo.launch(show_error=True, share=False)