import gradio as gr
from gradio_client import Client
import os
import logging
import json
from datetime import datetime
import tempfile
import numpy as np
from PIL import Image
import shutil
import httpx
import time
import base64
from gradio_client import Client, handle_file

# 로깅 설정
logging.basicConfig(level=logging.INFO)

# 타임아웃 설정을 30초로 늘림
httpx_client = httpx.Client(timeout=30.0)

max_retries = 3
retry_delay = 5  # 5초 대기

for attempt in range(max_retries):
    try:
        api_client = Client("http://211.233.58.202:7960/")
        api_client.httpx_client = httpx_client  # httpx 클라이언트 설정
        break  # 성공하면 루프 종료
    except httpx.ReadTimeout:
        if attempt < max_retries - 1:  # 마지막 시도가 아니면
            print(f"Connection timed out. Retrying in {retry_delay} seconds...")
            time.sleep(retry_delay)
        else:
            print("Failed to connect after multiple attempts.")
            raise  # 모든 시도 실패 시 예외 발생

            
# 갤러리 저장 디렉토리 설정
GALLERY_DIR = "gallery"
GALLERY_JSON = "gallery.json"

# 갤러리 디렉토리 생성
os.makedirs(GALLERY_DIR, exist_ok=True)

def save_to_gallery(video_path, prompt):
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    new_video_path = os.path.join(GALLERY_DIR, f"{timestamp}.mp4")
    
    # 비디오 파일 복사
    shutil.copy2(video_path, new_video_path)
    
    # 갤러리 정보 저장
    gallery_info = {
        "video": new_video_path,
        "prompt": prompt,
        "timestamp": timestamp
    }
    
    if os.path.exists(GALLERY_JSON):
        with open(GALLERY_JSON, "r") as f:
            gallery = json.load(f)
    else:
        gallery = []
    
    gallery.append(gallery_info)
    
    with open(GALLERY_JSON, "w") as f:
        json.dump(gallery, f, indent=2)
    
    return new_video_path


def load_gallery():
    if os.path.exists(GALLERY_JSON):
        with open(GALLERY_JSON, "r") as f:
            gallery = json.load(f)
        return [(item["video"], item["prompt"]) for item in reversed(gallery)]
    return []



def respond(image, prompt, steps, cfg_scale, eta, fs, seed, video_length):
    logging.info(f"Received prompt: {prompt}, steps: {steps}, cfg_scale: {cfg_scale}, "
                 f"eta: {eta}, fs: {fs}, seed: {seed}, video_length: {video_length}")
    
    try:
        # 이미지 파일 처리
        if image is not None:
            image_file = handle_file(image)
        else:
            image_file = None

        # 비디오 생성 요청
        result = api_client.predict(
            image=image_file,
            prompt=prompt,
            steps=steps,
            cfg_scale=cfg_scale,
            eta=eta,
            fs=fs,
            seed=seed,
            video_length=video_length,
            api_name="/infer"
        )
        logging.info("API response received: %s", result)
        
        # 결과 확인 및 처리
        if isinstance(result, dict) and 'video' in result:
            saved_video_path = save_to_gallery(result['video'], prompt)
            return saved_video_path
        else:
            raise ValueError("Unexpected API response format")
    except Exception as e:
        logging.error("Error during API request: %s", str(e))
        return "Failed to generate video due to an error."

css = """
footer {
    visibility: hidden;
}
"""


# 이미지 생성을 위한 예제 프롬프트
examples = [
    ["A glamorous young woman with long, wavy blonde hair and smokey eye makeup, posing in a luxury hotel room. She’s wearing a sparkly gold cocktail dress and holding up a white card with 'openfree.ai' written on it in elegant calligraphy. Soft, warm lighting creates a luxurious atmosphere. ", "q1.webp"],
    ["A fantasy map of a fictional world, with detailed terrain and cities.", "q19.webp"]
]

def use_prompt(prompt):
    return prompt

with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
    with gr.Tab("Generate"):
        with gr.Row():
            input_image = gr.Image(label="Upload an image", type="filepath")
            input_text = gr.Textbox(label="Enter your prompt for video generation")
            output_video = gr.Video(label="Generated Video")
        
        with gr.Row():
            steps = gr.Slider(minimum=1, maximum=100, step=1, label="Steps", value=30)
            cfg_scale = gr.Slider(minimum=1, maximum=15, step=0.1, label="CFG Scale", value=3.5)
            eta = gr.Slider(minimum=0, maximum=1, step=0.1, label="ETA", value=1)
            fs = gr.Slider(minimum=1, maximum=30, step=1, label="FPS", value=8)
            seed = gr.Slider(minimum=0, maximum=1000000, step=1, label="Seed", value=123)
            video_length = gr.Slider(minimum=1, maximum=10, step=1, label="Video Length (seconds)", value=2)
        
        with gr.Row():
            for prompt, image_file in examples:
                with gr.Column():
                    gr.Image(image_file, label=prompt[:50] + "...")  # 프롬프트의 처음 50자만 표시
                    gr.Button("Use this prompt").click(
                        fn=use_prompt,
                        inputs=[],
                        outputs=input_text,
                        api_name=False
                    ).then(
                        lambda x=prompt: x,
                        inputs=[],
                        outputs=input_text
                    )
    
    with gr.Tab("Gallery"):
        gallery = gr.Gallery(
            label="Generated Videos",
            show_label=False,
            elem_id="gallery",
            columns=[5],
            rows=[3],
            object_fit="contain",
            height="auto"
        )
        refresh_btn = gr.Button("Refresh Gallery")

    def update_gallery():
        return load_gallery()

    refresh_btn.click(fn=update_gallery, inputs=None, outputs=gallery)
    demo.load(fn=update_gallery, inputs=None, outputs=gallery)

    input_text.submit(
        fn=respond, 
        inputs=[input_image, input_text, steps, cfg_scale, eta, fs, seed, video_length], 
        outputs=output_video
    ).then(
        fn=update_gallery,
        inputs=None,
        outputs=gallery
    )

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