import streamlit as st
import subprocess
import shutil
import os
import zipfile
import io

# تابع اجرای دستورات ترمینال
def run_command(command):
    try:
        # استفاده از shell=True برای دستورات ترکیبی
        result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True)
        return True, result.stdout
    except subprocess.CalledProcessError as e:
        return False, e.stderr

# تابع ساخت فایل زیپ در حافظه
def create_zip_buffer(source_dir, flatten=False):
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(source_dir):
            if '.git' in root:
                continue
            for file in files:
                file_path = os.path.join(root, file)
                if flatten:
                    arcname = file
                else:
                    arcname = os.path.relpath(file_path, source_dir)
                zipf.write(file_path, arcname=arcname)
    return zip_buffer.getvalue()

# تنظیمات صفحه
st.set_page_config(page_title="Hugging Face Space Tools", page_icon="🛠️")

st.title("🛠️ ابزارهای مدیریت اسپیس هاگینگ‌فیس")

# ایجاد دو تب
tab1, tab2 = st.tabs(["🔄 انتقال به گیت‌هاب", "📦 دانلود اسپیس (فایل ZIP)"])

# ---------------------------------------------------------
# تب اول: انتقال به گیت‌هاب (پشتیبانی از لینک و فایل زیپ + حل مشکل پوشه اضافه)
# ---------------------------------------------------------
with tab1:
    st.info("منبع خود را برای انتقال به گیت‌هاب انتخاب کنید:")
    
    # انتخاب نوع منبع خارج از فرم برای رندر داینامیک
    source_type = st.radio("نوع منبع:", ["لینک اسپیس هاگینگ‌فیس", "آپلود فایل ZIP"])
    
    with st.form("transfer_form"):
        hf_url = ""
        uploaded_zip = None
        
        if source_type == "لینک اسپیس هاگینگ‌فیس":
            hf_url = st.text_input("لینک صفحه اسپیس هاگینگ فیس (منبع)", 
                                   placeholder="https://huggingface.co/spaces/user/space")
        else:
            uploaded_zip = st.file_uploader("آپلود فایل ZIP (منبع)", type=["zip"])
            
        st.markdown("---")
        gh_url = st.text_input("لینک کامل ریپازیتوری گیت‌هاب (مقصد)", 
                               placeholder="https://github.com/Hamed744/Zirnavis.git")
        
        gh_token = st.text_input("توکن گیت‌هاب (Classic Token)", type="password")
        
        submitted = st.form_submit_button("شروع عملیات انتقال 🚀")

    if submitted:
        if not gh_url or not gh_token:
            st.error("لطفاً لینک گیت‌هاب و توکن را وارد کنید.")
        elif source_type == "لینک اسپیس هاگینگ‌فیس" and not hf_url:
            st.error("لطفاً لینک اسپیس را وارد کنید.")
        elif source_type == "لینک اسپیس هاگینگ‌فیس" and "hf.space" in hf_url:
            st.error("❌ لینک اشتباه است! لطفاً لینک اصلی پروفایل را وارد کنید.")
        elif source_type == "آپلود فایل ZIP" and not uploaded_zip:
            st.error("لطفاً فایل ZIP را آپلود کنید.")
        else:
            status_area = st.empty()
            temp_dir = "temp_repo_gh"
            
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir, ignore_errors=True)
                
            try:
                auth_gh_url = gh_url.replace("https://", f"https://{gh_token}@")
                
                # مرحله 1: آماده‌سازی فایل‌ها (دانلود یا اکسترکت)
                if source_type == "لینک اسپیس هاگینگ‌فیس":
                    status_area.info("📥 در حال دانلود از هاگینگ فیس...")
                    success, msg = run_command(f"git clone {hf_url} {temp_dir}")
                    if not success:
                        st.error(f"خطا در دانلود: {msg}")
                        st.stop()
                else:
                    status_area.info("📥 در حال استخراج فایل ZIP...")
                    os.makedirs(temp_dir, exist_ok=True)
                    with zipfile.ZipFile(uploaded_zip, 'r') as zip_ref:
                        zip_ref.extractall(temp_dir)
                    
                    # بررسی و رفع مشکل پوشه اضافه (انتقال محتویات به روت اصلی)
                    # پوشه‌های سیستمی مک را نادیده می‌گیریم تا اختلال ایجاد نکنند
                    extracted_items = [item for item in os.listdir(temp_dir) if item != '__MACOSX']
                    
                    # اگر فقط یک پوشه در مسیر وجود داشت
                    if len(extracted_items) == 1:
                        single_folder_path = os.path.join(temp_dir, extracted_items[0])
                        if os.path.isdir(single_folder_path):
                            # تمام فایل‌های داخل آن پوشه را یک سطح بالاتر (به روت) منتقل می‌کنیم
                            for item in os.listdir(single_folder_path):
                                shutil.move(os.path.join(single_folder_path, item), temp_dir)
                            # پوشه خالی باقیمانده را حذف می‌کنیم
                            os.rmdir(single_folder_path)
                
                # مرحله 2: عملیات گیت و پوش
                os.chdir(temp_dir)
                status_area.info("⚙️ در حال تنظیمات گیت...")
                
                if os.path.exists(".git"):
                    shutil.rmtree(".git", ignore_errors=True)
                
                run_command("git init")
                run_command("git add .")
                run_command('git config user.email "bot@transfer.com"')
                run_command('git config user.name "Transfer Bot"')
                run_command('git commit -m "Moved to GitHub automatically"')
                run_command("git branch -M main")
                
                status_area.info("📤 در حال آپلود به گیت‌هاب...")
                push_success, push_msg = run_command(f"git push -f {auth_gh_url} main")
                
                if push_success:
                    st.success("✅ تمام شد! فایل‌ها دقیقاً در روت ریپازیتوری منتقل شدند.")
                    st.balloons()
                else:
                    clean_error = push_msg.replace(gh_token, "***")
                    st.error(f"خطا در آپلود: {clean_error}")
                
                os.chdir("..")
                    
            except Exception as e:
                st.error(f"خطای سیستمی: {e}")
            finally:
                try:
                    if os.getcwd().endswith(temp_dir):
                        os.chdir("..")
                except:
                    pass
                if os.path.exists(temp_dir):
                    shutil.rmtree(temp_dir, ignore_errors=True)

# ---------------------------------------------------------
# تب دوم: دانلود فایل‌ها به صورت زیپ
# ---------------------------------------------------------
with tab2:
    st.info("با استفاده از این بخش می‌توانید کل فایل‌های یک اسپیس عمومی را به صورت ZIP دانلود کنید.")
    
    hf_url_zip = st.text_input("لینک صفحه اسپیس هاگینگ فیس", 
                               placeholder="https://huggingface.co/spaces/elias207/zirnavis52", key="zip_url")
    
    zip_mode = st.radio("ساختار فایل ZIP خروجی را انتخاب کنید:", 
                        ("حفظ ساختار پوشه‌ها (فایل‌ها در پوشه‌های اصلی خودشان)", 
                         "همه فایل‌ها در روت (حذف پوشه‌بندی و انتقال همه فایل‌ها به روت زیپ)"))

    if st.button("آماده‌سازی فایل دانلود 🗜️"):
        if not hf_url_zip:
            st.error("لطفاً لینک اسپیس را وارد کنید.")
        elif "hf.space" in hf_url_zip:
            st.error("❌ لینک اشتباه است! لینک اصلی (huggingface.co/spaces/...) را وارد کنید.")
        else:
            status_zip = st.empty()
            temp_dir_zip = "temp_repo_zip"
            
            if os.path.exists(temp_dir_zip):
                shutil.rmtree(temp_dir_zip, ignore_errors=True)
                
            try:
                status_zip.info("📥 در حال دریافت فایل‌ها از هاگینگ‌فیس...")
                success, msg = run_command(f"git clone {hf_url_zip} {temp_dir_zip}")
                
                if not success:
                    st.error(f"خطا در دریافت فایل‌ها: {msg}")
                else:
                    status_zip.info("⚙️ در حال فشرده‌سازی فایل‌ها...")
                    flatten_files = True if zip_mode == "همه فایل‌ها در روت (حذف پوشه‌بندی و انتقال همه فایل‌ها به روت زیپ)" else False
                    zip_data = create_zip_buffer(temp_dir_zip, flatten=flatten_files)
                    
                    status_zip.success("✅ فایل با موفقیت آماده شد!")
                    
                    repo_name = hf_url_zip.strip("/").split("/")[-1]
                    file_suffix = "flat" if flatten_files else "structured"
                    download_filename = f"{repo_name}_{file_suffix}.zip"
                    
                    st.download_button(
                        label="⬇️ دانلود فایل ZIP",
                        data=zip_data,
                        file_name=download_filename,
                        mime="application/zip",
                        type="primary"
                    )
                    
            except Exception as e:
                st.error(f"خطای سیستمی: {e}")
            finally:
                if os.path.exists(temp_dir_zip):
                    shutil.rmtree(temp_dir_zip, ignore_errors=True)