import os
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from transformers import CLIPTextModel, CLIPTokenizer, Trainer, TrainingArguments
from datasets import load_dataset
import json
import time

# تحميل نموذج Stable Diffusion الأساسي
MODEL_NAME = "runwayml/stable-diffusion-v1-5"
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = StableDiffusionPipeline.from_pretrained(MODEL_NAME)
pipe.to(device)

# قاعدة بيانات للمستخدمين والنماذج المحفوظة
user_data_file = "user_data.json"

# تحميل البيانات من ملف JSON إذا كان موجودًا
if os.path.exists(user_data_file):
    with open(user_data_file, "r") as f:
        users = json.load(f)
else:
    users = {}

saved_models = {}

# وظيفة لتدريب النموذج على الصور المرفوعة
def train_model(user_email, images):
    if not user_email:
        return "يجب تسجيل الدخول لحفظ النموذج."
    
    user_model_id = f"user_model_{user_email}"
    saved_models[user_email] = user_model_id

    # تأكد من أن الصور مرفوعة
    if len(images) == 0:
        return "⚠️ يجب رفع صور للتدريب."

    # تجهيز البيانات الخاصة بالتدريب
    dataset = load_dataset('imagefolder', data_dir=images)

    # إعداد الـ Trainer للتدريب
    training_args = TrainingArguments(
        output_dir="./output",            # مكان حفظ النموذج المدرب
        per_device_train_batch_size=1,    # حجم الدفعة
        num_train_epochs=3,               # عدد العصور (epochs)
        logging_dir='./logs',             # مكان تسجيل السجلات
        logging_steps=10,
    )

    # تحميل Tokenizer و Model من CLIP لتوليد النصوص
    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch16")
    text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch16")

    # إعداد الـ Trainer باستخدام Trainer من Hugging Face
    trainer = Trainer(
        model=pipe.unet, 
        args=training_args, 
        train_dataset=dataset["train"],
        tokenizer=tokenizer,
        data_collator=None,
    )

    # بدء عملية التدريب
    trainer.train()

    return f"✅ تم حفظ النموذج بنجاح: {user_model_id}"

# وظيفة لإنشاء الصور باستخدام البرومبتات
def generate_image(prompt, user_email=None):
    if user_email and user_email in saved_models:
        model_id = saved_models[user_email]
        result = pipe(prompt).images[0]
    else:
        result = pipe(prompt).images[0]

    return result

# وظيفة لتسجيل المستخدمين
def register(email, password):
    if email in users:
        return "⚠️ البريد الإلكتروني مسجل بالفعل!"
    users[email] = {"password": generate_password_hash(password)}
    save_users_data()
    return "✅ تم التسجيل بنجاح!"

# وظيفة لتسجيل الدخول
def login(email, password):
    if email not in users or not check_password_hash(users[email]["password"], password):
        return "❌ البريد الإلكتروني أو كلمة المرور غير صحيحة!"
    return f"✅ تم تسجيل الدخول: {email}"

# حفظ بيانات المستخدمين إلى ملف JSON
def save_users_data():
    with open(user_data_file, "w") as f:
        json.dump(users, f)

# واجهة Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🖼️ إنشاء صور مخصصة باستخدام الذكاء الاصطناعي")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Files(label="📤 رفع صورك للتدريب")  # رفع الصور
            user_email_input = gr.Textbox(label="📧 بريدك الإلكتروني (اختياري لحفظ النموذج)")
            train_button = gr.Button("🔧 تدريب النموذج")
            train_output = gr.Textbox(label="🔔 نتيجة التدريب")
 
        with gr.Column():
            prompt_input = gr.Textbox(label="✏️ أدخل البرومبت لإنشاء صورة")
            generate_button = gr.Button("🎨 إنشاء صورة")
            output_image = gr.Image(label="📷 الصورة الناتجة")
 
    # تم إزالة progress_bar من هنا
    train_button.click(train_model, inputs=[user_email_input, image_input], outputs=train_output)
    generate_button.click(generate_image, inputs=[prompt_input, user_email_input], outputs=output_image)
 
    gr.Markdown("### 🔑 تسجيل الدخول / التسجيل")
    
    with gr.Row():
        with gr.Column():
            login_email = gr.Textbox(label="📧 البريد الإلكتروني")
            login_password = gr.Textbox(label="🔑 كلمة المرور", type="password")
            login_button = gr.Button("🚀 تسجيل الدخول")
            login_output = gr.Textbox(label="🔔 حالة تسجيل الدخول")
 
        with gr.Column():
            register_email = gr.Textbox(label="📧 البريد الإلكتروني للتسجيل")
            register_password = gr.Textbox(label="🔑 كلمة المرور للتسجيل", type="password")
            register_button = gr.Button("📝 تسجيل حساب جديد")
            register_output = gr.Textbox(label="🔔 حالة التسجيل")
 
    login_button.click(login, inputs=[login_email, login_password], outputs=login_output)
    register_button.click(register, inputs=[register_email, register_password], outputs=register_output)

# تشغيل التطبيق مع رابط عام
demo.launch(share=True)