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}

def get_client_info(request: gr.Request):
    clientHost = request.client.host
    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, transfer_duration, request: gr.Request):
    client_host = get_client_info(request)
    flag = check_ip_limit(client_host)
    if not flag:
        raise gr.Error(f"Exceed the request limit of 1 requests per five minute.")

    print(f"Received video: {in_video}, aspect_ratio: {aspect_ratio}, style: {style}, transfer_duration: {transfer_duration}")

    if in_video is None:
        raise Exception("input_video is None")

    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)
        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)
            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)

        with gr.Row():
            
            aspect_ratio = gr.Radio(["16:9", "9:16", "1:1"], label="3. Aspect Ratio(Width:Height)", value="16:9", interactive=True)
            transfer_duration = gr.Slider(
                minimum=0,
                maximum=3,
                step=1,
                value=3,
                label="Transfer Duration (seconds), for more than 3 seconds, please join our Discord(https://discord.gg/cN8geBsBmu) or website",
                interactive=True
            )
        submit_button = gr.Button("Generate", interactive=True)

        examples = gr.Examples(
            [
                ["test32.mp4", "ClayStyle1", "9:16" ],
                ["test32.mp4", "3D P1", "9:16" ],
            ],
            inputs=[in_video, style, aspect_ratio, transfer_duration],
            outputs=[out_video],
            fn=fn,
        )

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


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