import os
import shutil
import gradio as gr
from huggingface_hub import snapshot_download, upload_folder, create_repo

HF_TOKEN = os.environ.get("HF_TOKEN")

# قائمة بأشهر المستودعات المفتوحة اللي فيها النماذج بصيغة ONNX
FALLBACK_REPOS = [
    "InferencePrincess/UVR-MDX-ONNX",
    "jarredou/UVR-Models",
    "seanghay/uvr_models",
    "DakeQQ/mdx-net-onnx"
]

def migrate_vocal_separator():
    if not HF_TOKEN:
        yield "❌ خطأ: لم يتم العثور على التوكن السري HF_TOKEN!"
        return

    TARGET_REPO = "M-hv1/titan-vocal-separator"
    LOCAL_DIR = "./titan_separator_clean"
    
    # 1. إنشاء مستودعك الخاص
    try:
        yield f"⏳ [1/3] جاري إنشاء مستودع الفلتر الصوتي {TARGET_REPO}..."
        create_repo(repo_id=TARGET_REPO, repo_type="model", token=HF_TOKEN, exist_ok=True)
    except Exception as e:
        yield f"❌ خطأ في إنشاء المستودع: {str(e)}"
        return

    # 2. محاولة السحب من المستودعات (الصيد)
    download_success = False
    yield f"⏳ [2/3] جاري البحث عن ملفات ONNX في المستودعات البديلة..."
    
    for repo in FALLBACK_REPOS:
        try:
            yield f"🔍 جاري محاولة السحب من: {repo}..."
            snapshot_download(
                repo_id=repo,
                local_dir=LOCAL_DIR,
                token=HF_TOKEN,
                allow_patterns=["*.onnx", "*.json"], # هنسحب الـ ONNX بس
                ignore_patterns=["*.ckpt", "*.pth", "*.bin"] # تجاهل الأوزان الثقيلة
            )
            download_success = True
            yield f"✅ تم العثور على الملفات وسحبها بنجاح من: {repo}"
            break # لو نجح يوقف اللف
        except Exception as e:
            yield f"⚠️ فشل السحب من {repo} (ممكن يكون اتحذف)، جاري تجربة اللي بعده..."

    if not download_success:
        yield "❌ للأسف، كل المستودعات البديلة مغلقة أو غير متاحة. هنحتاج نرفع الملفات يدوياً."
        return

    # 3. الرفع لمستودع تيتان
    try:
        yield f"🚀 [3/3] جاري رفع الفلتر إلى مستودعات تيتان..."
        upload_folder(
            folder_path=LOCAL_DIR,
            repo_id=TARGET_REPO,
            repo_type="model",
            token=HF_TOKEN,
            commit_message="⚡ Titan Vocal Separator: Fallback ONNX version downloaded"
        )

        if os.path.exists(LOCAL_DIR):
            shutil.rmtree(LOCAL_DIR)

        yield f"🎯 تمت المهمة بنجاح يا جوكر! الفلتر الصوتي جاهز دلوقتي في مستودعك."

    except Exception as e:
        yield f"❌ حدث خطأ أثناء الرفع لمستودعك:\n{str(e)}"

# --- واجهة الـ UI ---
with gr.Blocks(title="Titan Audio Separator Setup") as demo:
    gr.Markdown("# 🎛️ Titan Vocal & Noise Separator Migrator")
    gr.Markdown("نسخة متطورة: تبحث أوتوماتيكياً في 4 مستودعات مختلفة لسحب الفلتر الصوتي.")
    
    btn = gr.Button("تجهيز ورفع الفلتر الصوتي 🚀", variant="primary")
    output_log = gr.Textbox(label="Status Log", lines=10)
    
    btn.click(fn=migrate_vocal_separator, outputs=output_log)

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