# app.py
import os
import oss2
import sys
import uuid
import shutil
import time
import gradio as gr
import requests

import dashscope
from dashscope.utils.oss_utils import check_and_upload_local

DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")

# چک کردن وجود API Key
if not DASHSCOPE_API_KEY:
    print("⚠️ هشدار: DASHSCOPE_API_KEY تنظیم نشده است!")
    print("برای استفاده از برنامه، لطفاً API Key را تنظیم کنید.")
    DASHSCOPE_API_KEY = None
else:
    dashscope.api_key = DASHSCOPE_API_KEY


class WanAnimateApp:
    def __init__(self, url, get_url):
        self.url = url
        self.get_url = get_url

    def predict(
        self, 
        ref_img,
        video,
        model_id,
        model,
    ):
        # چک کردن API Key قبل از شروع
        if not DASHSCOPE_API_KEY:
            return None, "خطا: DASHSCOPE_API_KEY تنظیم نشده است. لطفاً متغیر محیطی را تنظیم کنید."
        
        try:
            # آپلود فایل‌ها به OSS
            _, image_url = check_and_upload_local(model_id, ref_img, DASHSCOPE_API_KEY)
            _, video_url = check_and_upload_local(model_id, video, DASHSCOPE_API_KEY)

            # آماده‌سازی درخواست
            payload = {
                "model": model_id,
                "input": {
                    "image_url": image_url,
                    "video_url": video_url
                },
                "parameters": {
                    "check_image": True,
                    "mode": model,
                }
            }
            
            # تنظیم هدرها
            headers = {
                "X-DashScope-Async": "enable",
                "X-DashScope-OssResourceResolve": "enable",
                "Authorization": f"Bearer {DASHSCOPE_API_KEY}",
                "Content-Type": "application/json"
            }
            
            # ارسال درخواست اولیه
            url = self.url
            response = requests.post(url, json=payload, headers=headers)
            
            # چک کردن موفقیت درخواست
            if response.status_code != 200:
                return None, f"خطا در درخواست: {response.status_code}: {response.text}"
            
            # دریافت task ID
            result = response.json()
            task_id = result.get("output", {}).get("task_id")
            if not task_id:
                return None, "خطا: دریافت task ID انجام نشد"
            
            # پولینگ برای نتایج
            get_url = f"{self.get_url}/{task_id}"
            headers = {
                "Authorization": f"Bearer {DASHSCOPE_API_KEY}",
                "Content-Type": "application/json"
            }
            
            while True:
                response = requests.get(get_url, headers=headers)
                if response.status_code != 200:
                    return None, f"خطا در دریافت وضعیت: {response.status_code}: {response.text}"
                
                result = response.json()
                print(result)
                task_status = result.get("output", {}).get("task_status")
                
                if task_status == "SUCCEEDED":
                    video_url = result["output"]["results"]["video_url"]
                    return video_url, "موفق ✓"
                elif task_status == "FAILED":
                    error_msg = result.get("output", {}).get("message", "خطای ناشناخته")
                    code_msg = result.get("output", {}).get("code", "کد ناشناخته")
                    return None, f"تسک ناموفق: {error_msg} کد: {code_msg} TaskId: {task_id}"
                else:
                    # تسک در حال اجراست
                    time.sleep(5)
        
        except Exception as e:
            return None, f"خطا: {str(e)}"

def start_app():
    import argparse
    parser = argparse.ArgumentParser(description="Wan2.2-Animate ابزار تولید ویدیو")
    args = parser.parse_args()
    
    url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2video/video-synthesis/"
    get_url = f"https://dashscope.aliyuncs.com/api/v1/tasks/"
    app = WanAnimateApp(url=url, get_url=get_url)

    with gr.Blocks(title="Wan2.2-Animate تولید ویدیو") as demo:
        gr.HTML("""
            <div style="padding: 2rem; text-align: center; max-width: 1200px; margin: 0 auto; font-family: Arial, sans-serif;">
                <h1 style="font-size: 2.5rem; font-weight: bold; margin-bottom: 0.5rem; color: #333;">
                    Wan2.2-Animate: Unified Character Animation and Replacement
                </h1>
                <h3 style="font-size: 2rem; font-weight: bold; margin-bottom: 0.5rem; color: #333;">
                    Wan2.2-Animate: مدل یکپارچه انیمیشن کاراکتر و جایگزینی شخصیت ویدیو
                </h3>
                <div style="font-size: 1.25rem; margin-bottom: 1.5rem; color: #555;">
                    Tongyi Lab, Alibaba
                </div>
            </div>
        """)
        
        gr.HTML("""
            <details>
                <summary>‼️ راهنمای استفاده</summary>
                
                <h4>دو حالت پشتیبانی می‌شود:</h4>
                <ul>
                    <li><strong>Move Mode</strong>: حرکت دادن کاراکتر تصویر ورودی با حرکات از ویدیو ورودی</li>
                    <li><strong>Mix Mode</strong>: جایگزینی کاراکتر ویدیو ورودی با کاراکتر از تصویر ورودی</li>
                </ul>
                
                <h4>محدودیت‌های ورودی:</h4>
                <ul>
                    <li>حجم فایل ویدیو: کمتر از 200MB</li>
                    <li>رزولوشن ویدیو: حداقل ضلع بیشتر از 200، حداکثر ضلع کمتر از 2048</li>
                    <li>مدت ویدیو: 2 تا 30 ثانیه</li>
                    <li>نسبت تصویر: 1:3 تا 3:1</li>
                    <li>فرمت ویدیو: mp4, avi, mov</li>
                    <li>حجم تصویر: کمتر از 5MB</li>
                    <li>رزولوشن تصویر: حداقل ضلع بیشتر از 200، حداکثر ضلع کمتر از 4096</li>
                    <li>فرمت تصویر: jpg, png, jpeg, webp, bmp</li>
                </ul>
                
                <h4>کیفیت استنتاج:</h4>
                <ul>
                    <li><strong>wan-pro</strong>: 25fps, 720p</li>
                    <li><strong>wan-std</strong>: 15fps, 720p</li>
                </ul>
            </details>
        """)

        with gr.Row():
            with gr.Column():    
                ref_img = gr.Image(
                    label="تصویر مرجع (Reference Image)",
                    type="filepath",
                    sources=["upload"],
                )
                
                video = gr.Video(
                    label="ویدیوی الگو (Template Video)",
                    sources=["upload"],
                )
                
                with gr.Row():
                    model_id = gr.Dropdown(
                        label="حالت (Mode)",
                        choices=["wan2.2-animate-move", "wan2.2-animate-mix"],
                        value="wan2.2-animate-move",
                    )

                    model = gr.Dropdown(
                        label="کیفیت استنتاج (Inference Quality)",
                        choices=["wan-pro", "wan-std"],
                        value="wan-pro",
                    )

                run_button = gr.Button("تولید ویدیو (Generate Video)")

            with gr.Column():
                output_video = gr.Video(label="ویدیوی خروجی (Output Video)")
                output_status = gr.Textbox(label="وضعیت (Status)")
        
        run_button.click(
            fn=app.predict,
            inputs=[ref_img, video, model_id, model],
            outputs=[output_video, output_status],
        )

        # مثال‌ها بدون کش
        example_data = [
            ['./examples/mov/1/1.jpeg', './examples/mov/1/1.mp4', 'wan2.2-animate-move', 'wan-pro'],
            ['./examples/mov/2/2.jpeg', './examples/mov/2/2.mp4', 'wan2.2-animate-move', 'wan-pro'],
            ['./examples/mix/1/1.jpeg', './examples/mix/1/1.mp4', 'wan2.2-animate-mix', 'wan-pro'],
            ['./examples/mix/2/2.jpeg', './examples/mix/2/2.mp4', 'wan2.2-animate-mix', 'wan-pro']
        ]

        if example_data and os.path.exists('./examples'):
            gr.Examples(
                examples=example_data,
                inputs=[ref_img, video, model_id, model],
                # outputs و fn و cache_examples را حذف کنید
            )
    
    demo.queue(default_concurrency_limit=100)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860
    )


if __name__ == "__main__":
    start_app()