import os
import sys
import subprocess
import threading
import time
import glob
import shutil
import zipfile
import json
import dataclasses
import asyncio
import nest_asyncio
import gradio as gr
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import ApplicationBuilder, MessageHandler, CommandHandler, CallbackQueryHandler, filters, ContextTypes

# تفعيل nest_asyncio للتعامل مع الحلقات التكرارية في HF
nest_asyncio.apply()

# --- إعدادات البرنامج الأساسية ---
BASE_DIR = "/data" if os.path.exists("/data") else "."
APP_DIR = os.path.join(BASE_DIR, "BallonsTranslator")
WORK_DIR = os.path.join(BASE_DIR, "bot_work")
TOKEN = os.getenv("7196327387:AAH5VtQM6M_ZoF8Ux6nznapikrgcfNT2CZo")

# تحميل المشروع إذا لم يكن موجوداً
if not os.path.exists(APP_DIR):
    print("Downloading BallonsTranslator...")
    subprocess.run(["git", "clone", "https://github.com/dmMaze/BallonsTranslator.git", APP_DIR])

sys.path.insert(0, APP_DIR)
os.chdir(APP_DIR)

from utils.config import ProgramConfig

MODELS = {
    'lama_large_512px': 'LAMA Large (Default)',
    'lama': 'LAMA Light',
    'aot': 'AOT (Fast)',
    'manga_inpaint': 'Manga Inpaint'
}
current_model = "lama_large_512px"

# --- وظائف المساعدة ---
def save_config(model_name):
    cfg = ProgramConfig()
    cfg.module.inpainter = model_name
    cfg.module.enable_translate = False
    cfg.module.enable_ocr = True
    cfg.module.enable_inpaint = True
    cfg_path = os.path.join(APP_DIR, 'config/config.json')
    os.makedirs(os.path.dirname(cfg_path), exist_ok=True)
    with open(cfg_path, 'w', encoding='utf-8') as f:
        json.dump(dataclasses.asdict(cfg), f, indent=2, ensure_ascii=False)

def run_inpaint(folder):
    save_config(current_model)
    # تشغيل المعالجة بدون واجهة رسومية (Headless)
    subprocess.run(["python", "launch.py", "--headless", "--exec_dirs", folder, "--ldpi", "96"], cwd=APP_DIR)
    return os.path.join(folder, 'inpainted')

# --- منطق البوت ---
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text(f"مرحباً! أرسل صورة أو ملف ZIP لتبييضه.\nالموديل الحالي: {current_model}")

async def on_photo(update: Update, context: ContextTypes.DEFAULT_TYPE):
    chat_id = update.effective_chat.id
    user_folder = os.path.join(WORK_DIR, f"p_{chat_id}")
    if os.path.exists(user_folder): shutil.rmtree(user_folder)
    os.makedirs(user_folder)

    msg = await update.message.reply_text("⏳ جاري المعالجة...")
    try:
        photo_file = await update.message.photo[-1].get_file()
        img_path = os.path.join(user_folder, "page.jpg")
        await photo_file.download_to_drive(img_path)
        
        result_dir = run_inpaint(user_folder)
        results = glob.glob(os.path.join(result_dir, "*.*"))
        
        if results:
            await update.message.reply_document(document=open(results[0], 'rb'), caption="✅ تم التبييض")
        else:
            await update.message.reply_text("❌ فشلت المعالجة.")
    except Exception as e:
        await update.message.reply_text(f"خطأ: {str(e)}")
    finally:
        await msg.delete()
        shutil.rmtree(user_folder, ignore_errors=True)

# --- تشغيل البوت ---
def run_bot():
    if not TOKEN:
        print("Error: BOT_TOKEN not found in environment variables!")
        return

    os.makedirs(WORK_DIR, exist_ok=True)
    app = ApplicationBuilder().token(TOKEN).build()
    app.add_handler(CommandHandler("start", start))
    app.add_handler(MessageHandler(filters.PHOTO, on_photo))
    
    print("Bot is starting...")
    app.run_polling(drop_pending_updates=True)

# --- واجهة Gradio (لإبقاء الهجين شغال) ---
def dummy_fn():
    return "Bot is running..."

if __name__ == "__main__":
    # تشغيل البوت في خيط (Thread) منفصل
    t = threading.Thread(target=run_bot)
    t.daemon = True
    t.start()

    # تشغيل واجهة بسيطة لـ Hugging Face
    gr.Interface(fn=dummy_fn, inputs=[], outputs="text", title="Telegram Inpaint Bot").launch(server_name="0.0.0.0", server_port=7860)