import subprocess
import os
import sys
import shutil
from pathlib import Path
import argparse
import gradio as gr

# ✅ التنظيف أولاً: فقط للمجلدات المؤقتة
folders_to_delete = ["./output", "./__pycache__", "./.cache", "./temp"]
for folder in folders_to_delete:
    if os.path.exists(folder):
        print(f"🗑️ حذف {folder}")
        shutil.rmtree(folder)

# ✅ طباعة حالة الذاكرة
import psutil
mem = psutil.virtual_memory()
print(f"🔍 RAM المستخدمة: {mem.used / 1e9:.2f} GB / {mem.total / 1e9:.2f} GB")

# ✅ تحميل الموديلات إذا ما كانت موجودة
if not os.path.exists("./models/fantasytalking_model.ckpt"):
    print("🛠️ جاري تحميل النماذج عبر download_models.py ...")
    subprocess.run(["python", "download_models.py"])

# ✅ إعداد المسارات
sys.path.append(os.path.abspath("."))

# ✅ استيراد المكونات
from STT.sst import speech_to_text
from LLM.llm import generate_reply
from TTS_X.tts import generate_voice
from FantasyTalking.infer import load_models, main

# ✅ ثابتات النموذج
args_template = argparse.Namespace(
    fantasytalking_model_path="./models/fantasytalking_model.ckpt",
    wav2vec_model_dir="./models/wav2vec2-base-960h",
    wan_model_dir="./models/Wan2.1-I2V-14B-720P",
    image_path="",
    audio_path="",
    prompt="",
    output_dir="./output",
    image_size=512,
    audio_scale=1.0,
    prompt_cfg_scale=5.0,
    audio_cfg_scale=5.0,
    max_num_frames=81,
    inference_steps=20,
    fps=23,
    num_persistent_param_in_dit=None,
    seed=1111
)

# ✅ تحميل النماذج
print("🚀 جاري تحميل FantasyTalking و Wav2Vec...")
pipe, fantasytalking, wav2vec_processor, wav2vec = load_models(args_template)
print("✅ تم التحميل!")

# ✅ توليد فيديو
def generate_video(image_path, audio_path, prompt, output_dir="./output"):
    # انسخي args_template إلى dict عشان نعدل عليه بسهولة
    args_dict = vars(args_template).copy()

    # نحدث فقط اللي نحتاجه
    args_dict.update({
        "image_path": image_path,
        "audio_path": audio_path,
        "prompt": prompt,
        "output_dir": output_dir
    })

    # نحول من dict إلى argparse.Namespace
    args = argparse.Namespace(**args_dict)

    return main(args, pipe, fantasytalking, wav2vec_processor, wav2vec)

# ✅ خط الأنابيب الكامل
def full_pipeline(user_audio, user_image):
    print("🎤 تحويل الصوت إلى نص...")
    user_text = speech_to_text(user_audio)

    print("💬 توليد الرد...")
    reply = generate_reply(user_text)

    print("🔊 تحويل الرد إلى صوت...")
    reply_audio_path = generate_voice(reply)

    print("📽️ توليد الفيديو...")
    Path("./output").mkdir(parents=True, exist_ok=True)
    video_path = generate_video(
        image_path=user_image,
        audio_path=reply_audio_path,
        prompt=reply
    )

    return user_text, reply, reply_audio_path, video_path

# ✅ واجهة Gradio
with gr.Blocks(title="🧠 صوتك يحرك صورة!") as demo:
    gr.Markdown("## 🎤➡️💬➡️🔊➡️📽️ من صوتك إلى فيديو متكلم!")

    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(label="🎙️ ارفع صوتك", type="filepath")
            image_input = gr.Image(label="🖼️ صورة المتحدث", type="filepath")
            btn = gr.Button("🎬 شغل")

        with gr.Column():
            user_text = gr.Textbox(label="📝 النص المسموع")
            reply_text = gr.Textbox(label="🤖 رد المساعد")
            reply_audio = gr.Audio(label="🔊 الرد المنطوق")
            video_output = gr.Video(label="📽️ الفيديو الناتج")

    btn.click(fn=full_pipeline,
              inputs=[audio_input, image_input],
              outputs=[user_text, reply_text, reply_audio, video_output])

demo.launch(inbrowser=True, share=True)