import threading
import time
import asyncio
import os
import json
import gradio as gr
from PIL import Image
from icrawler.builtin import BingImageCrawler
from telegram import Update
from telegram.ext import ApplicationBuilder, ContextTypes, CommandHandler, MessageHandler, filters
from huggingface_hub import InferenceClient

# --- الإعدادات والمتغيرات العالمية ---
LOGS = ["🤖 بدء تشغيل نظام التدريب الذكي مع ميزة الذاكرة الدائمة..."]
DATASET_DIR = "agent_dataset"
MEMORY_FILE = "bot_memory.json"

# متغيرات العقل والذاكرة الافتراضية (سيتم تحديثها تلقائياً إذا وجد ملف الحفظ)
CURRENT_KEYWORD = "لا يوجد"
TRAINING_STEP = 0
MAX_TRAINING_STEPS = 10
IS_TRAINING_COMPLETE = False
LEARNED_FEATURES = "ولاشيء بعد"
CURRENT_STATUS = "خامل - بانتظار تحديد الستايل وبدء التدريب"

if not os.path.exists(DATASET_DIR):
    os.makedirs(DATASET_DIR)

# --- 💾 دالات إدارة الذاكرة الدائمة (حفظ واسترجاع) ---
def save_memory_to_space():
    """حفظ حالة البوت الحالية داخل ملف JSON في السبيس"""
    try:
        memory_data = {
            "current_keyword": CURRENT_KEYWORD,
            "training_step": TRAINING_STEP,
            "is_training_complete": IS_TRAINING_COMPLETE,
            "learned_features": LEARNED_FEATURES
        }
        with open(MEMORY_FILE, "w", encoding="utf-8") as f:
            json.dump(memory_data, f, ensure_ascii=False, indent=4)
    except Exception as e:
        LOGS.append(f"⚠️ فشل حفظ الذاكرة في الملف: {str(e)}")

def load_memory_from_space():
    """استرجاع حالة البوت وعقله من ملف الـ JSON عند الإقلاع"""
    global CURRENT_KEYWORD, TRAINING_STEP, IS_TRAINING_COMPLETE, LEARNED_FEATURES, CURRENT_STATUS
    if os.path.exists(MEMORY_FILE):
        try:
            with open(MEMORY_FILE, "r", encoding="utf-8") as f:
                memory_data = json.load(f)
                CURRENT_KEYWORD = memory_data.get("current_keyword", "لا يوجد")
                TRAINING_STEP = memory_data.get("training_step", 0)
                IS_TRAINING_COMPLETE = memory_data.get("is_training_complete", False)
                LEARNED_FEATURES = memory_data.get("learned_features", "ولاشيء بعد")
                
                if IS_TRAINING_COMPLETE:
                    CURRENT_STATUS = f"✅ مسترجع من الذاكرة: اكتمل التدريب على {CURRENT_KEYWORD}"
                elif CURRENT_KEYWORD != "لا يوجد":
                    CURRENT_STATUS = f"🔄 مسترجع من الذاكرة: جاهز لمتابعة تدريب {CURRENT_KEYWORD} ({TRAINING_STEP}/{MAX_TRAINING_STEPS})"
                else:
                    CURRENT_STATUS = "خامل - بانتظار تحديد الستايل وبدء التدريب"
                    
                LOGS.append(f"🧠 [الذاكرة الدائمة]: تم استرجاع عقلك بنجاح! الستايل المحفوظ هو: ({CURRENT_KEYWORD})")
        except Exception as e:
            LOGS.append(f"⚠️ خطأ أثناء قراءة ملف الذاكرة المحفوظة: {str(e)}")
    else:
        LOGS.append("ℹ️ لم يتم العثور على ملف ذاكرة سابق، البوت يبدأ الآن بعقل نظيف.")

# استدعاء دالة التحميل فوراً عند تشغيل السيرفر لاستعادة البيانات السابقة
load_memory_from_space()

# استخدام السيرفرات السحابية للرسم السريع (موديل DreamShaper المخصص للأنمي والفن)
MODEL_ID = "Lykon/dreamshaper-8"
client = InferenceClient(model=MODEL_ID)

pause_event = threading.Event()
pause_event.set()

# --- 🎨 دالة الرسم الذكية (تدمج الستايل المدرب تلقائياً مع طلبك) ---
def generate_image_logic(prompt):
    global CURRENT_STATUS
    pause_event.clear() # إيقاف مؤقت للتدريب أثناء الرسم لحماية الموارد
    
    old_status = CURRENT_STATUS
    CURRENT_STATUS = "🎨 جاري توليد الصورة بالستايل الخاص بك..."
    
    # دمج الستايل الذي تدرب عليه البوت والمخزن بالذاكرة مع وصفك
    if CURRENT_KEYWORD != "لا يوجد":
        final_prompt = f"{prompt}, in {CURRENT_KEYWORD} style, highly detailed, perfect masterpiece"
    else:
        final_prompt = f"{prompt}, highly detailed, perfect masterpiece"
        
    LOGS.append(f"🎨 أمر رسم نشط للوصف المدمج: ({final_prompt})")
    
    try:
        image = client.text_to_image(final_prompt)
        output_path = "bot_output.png"
        image.save(output_path)
        LOGS.append("✅ تم إنتاج الصورة بنجاح بالستايل المطلوب!")
        return output_path
    except Exception as e:
        LOGS.append(f"❌ خطأ بالرسم: {str(e)}")
        return None
    finally:
        CURRENT_STATUS = old_status
        pause_event.set()

# --- 🧠 حلقة التدريب والتحليل التفصيلي للستايل ---
def continuous_learning_loop():
    global CURRENT_STATUS, TRAINING_STEP, IS_TRAINING_COMPLETE, LEARNED_FEATURES
    
    feature_steps = [
        "ملامح الوجه وتناسق العينين (Face Anatomy)",
        "الخطوط الأساسية والـ Outlines ورسم الهياكل",
        "توزيع الألوان ودرجات الظل والنور (Color Palette)",
        "المؤثرات البصرية والخلفيات العمق الضوئي",
        "ستايل ضربات الفرشاة والتكستشر (Textures)",
        "تعابير الوجه وحركة الجسد الكرتونية",
        "صقل الإضاءة المحيطية وتأثير النيون (Lighting)",
        "تفاصيل اليدين والملابس والإكسسوارات",
        "دمج العناصر الأمامية مع عزل الخلفية",
        "صقل الصورة النهائي ورفع الدقة الفنية (Final Rendering)"
    ]
    
    while True:
        pause_event.wait()
        
        if IS_TRAINING_COMPLETE or CURRENT_KEYWORD == "لا يوجد":
            time.sleep(2)
            continue
            
        CURRENT_STATUS = f"🏋️ جاري التدريب المكثف على ستايل: {CURRENT_KEYWORD} (خطوة {TRAINING_STEP + 1}/{MAX_TRAINING_STEPS})"
        
        try:
            # 1. سحب صورة لتحليلها
            crawler = BingImageCrawler(storage={'root_dir': DATASET_DIR})
            crawler.crawl(keyword=CURRENT_KEYWORD, max_num=1, overwrite=False)
            
            # 2. استخراج الميزة الحالية التي يتعلمها البوت
            current_feature = feature_steps[TRAINING_STEP % len(feature_steps)]
            LEARNED_FEATURES = f"🔄 يتعلم الآن: {current_feature}"
            LOGS.append(f"🧠 [تدريب]: تم هضم ميزة جديدة -> {current_feature}")
            
            TRAINING_STEP += 1
            
            # 3. حفظ التقدم والخطوة الحالية فوراً في ملف الذاكرة لضمان عدم ضياعها
            save_memory_to_space()
            
            # تحقق من الوصول للحد النهائي للتدريب
            if TRAINING_STEP >= MAX_TRAINING_STEPS:
                IS_TRAINING_COMPLETE = True
                CURRENT_STATUS = f"✅ اكتمل التدريب بنجاح 100% على ستايل: {CURRENT_KEYWORD}"
                LEARNED_FEATURES = "🏆 تم تعلم الستايل بالكامل وحفظه بذاكرة السبيس الدائمة!"
                LOGS.append(f"🎉 [مبارك]: البوت أصبح خبيراً الآن وتم تجميد الذاكرة بنجاح.")
                save_memory_to_space() # حفظ نهائي بعد الاكتمال
                
        except Exception as e:
            LOGS.append(f"⚠️ تنبيه بالتدريب: {str(e)}")
            
        time.sleep(10)

learning_thread = threading.Thread(target=continuous_learning_loop, daemon=True)
learning_thread.start()

# --- 🕹️ دالة بدء تدريب جديد من واجهة Gradio ---
def start_new_training(style_name):
    global CURRENT_KEYWORD, TRAINING_STEP, IS_TRAINING_COMPLETE, LEARNED_FEATURES, CURRENT_STATUS
    if not style_name.strip():
        return "❌ يرجى كتابة اسم الستايل أولاً", "لا شيء"
    
    CURRENT_KEYWORD = style_name
    TRAINING_STEP = 0
    IS_TRAINING_COMPLETE = False
    LEARNED_FEATURES = "⏳ جاري بدء الدورة الأولى وحفظ حجر الأساس في الذاكرة..."
    CURRENT_STATUS = f"🏋️ جاري بدء تدريب الستايل: {style_name}"
    LOGS.append(f"🚀 [إشارة تحكم]: أطلقت تدريباً لستايل جديد: {style_name}")
    
    # حفظ الحالة الجديدة فوراً بالملف
    save_memory_to_space()
    
    return CURRENT_STATUS, LEARNED_FEATURES

# --- 🤖 إعدادات بوت التليجرام ---
async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text(
        f"مرحباً بك في بوت الرسم الذكي ذو الذاكرة الدائمة! 🧠💾\n\n"
        f"• الستايل المسترجع حالياً: **{CURRENT_KEYWORD}**\n"
        f"• حالة التدريب: {CURRENT_STATUS}\n\n"
        f"اكتب لي أي وصف بالإنجليزية، وسأرسمه لك مدموجاً بالستايل المحفوظ فوراً!"
    )

async def draw_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
    prompt = update.message.text
    await update.message.reply_text("🎨 أمرك مطاع! تم تجميد التدريب مؤقتاً.. جاري الرسم واستدعاء الذاكرة الحالية...")
    
    image_path = generate_image_logic(prompt)
    
    if image_path and os.path.exists(image_path):
        await context.bot.send_photo(chat_id=update.effective_chat.id, photo=open(image_path, 'rb'))
        await update.message.reply_text("✅ تفضل النتيجة! تم استئناف وضع التدريب/الحفظ تلقائياً.")
    else:
        await update.message.reply_text("❌ واجهت السيرفرات ضغطاً، يرجى إعادة المحاولة.")

def run_telegram_bot():
    token = os.getenv("TELEGRAM_TOKEN")
    if not token:
        return
    application = ApplicationBuilder().token(token).build()
    application.add_handler(CommandHandler('start', start))
    application.add_handler(MessageHandler(filters.TEXT & (~filters.COMMAND), draw_command))
    application.run_polling(stop_signals=None)

telegram_thread = threading.Thread(target=lambda: asyncio.run(run_telegram_bot()), daemon=True)
telegram_thread.start()

# --- 🖥️ واجهة التحكم الرسومية (Gradio) ---
def get_ui_updates():
    return CURRENT_STATUS, LEARNED_FEATURES, "\n".join(LOGS[-12:])

with gr.Blocks() as demo:
    gr.Markdown("# 🚀 مركز الإدارة والتدريب الذكي لـ Catbot (نسخة الذاكرة الدائمة 💾)")
    
    with gr.Row():
        style_input = gr.Textbox(label="✍️ أدخل الستايل الجديد الذي تريد تدريب البوت عليه (بالإنجليزية)", placeholder="مثال: neon anime cyberpunk, pixel art, oily classic portrait")
        start_btn = gr.Button("🚀 ابدأ تدريب هذا الستايل فوراً")
        
    with gr.Row():
        status_box = gr.Textbox(label="📊 حالة السيرفر والتدريب المسترجعة")
        feature_box = gr.Textbox(label="🧠 تفاصيل الجزء الذي يتعلمه الذكاء الاصطناعي الآن")
        
    logs_box = gr.Code(label="📜 سجل الأحداث والعمليات المباشر (Logs)", language="markdown")
    
    start_btn.click(fn=start_new_training, inputs=style_input, outputs=[status_box, feature_box])
    
    demo.load(fn=get_ui_updates, outputs=[status_box, feature_box, logs_box])
    timer = gr.Timer(3)
    timer.tick(fn=get_ui_updates, outputs=[status_box, feature_box, logs_box])

demo.launch()