import os
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DDPMScheduler, Trainer, TrainingArguments, DiffusionPipeline
from datasets import load_dataset
from transformers import CLIPTextModel, CLIPTokenizer, PreTrainedModel
from werkzeug.security import generate_password_hash, check_password_hash
import json
import time
from pathlib import Path
from PIL import Image
from diffusers import DreamBoothTrainer, DreamBoothPipeline
from accelerate import Accelerator

# تحميل نموذج 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)

# قاعدة بيانات للمستخدمين والنماذج المحفوظة (سيتم حفظها في ملف JSON)
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 save_users_data():
    with open(user_data_file, "w") as f:
        json.dump(users, f)

# وظيفة تدريب النموذج على الصور المرفوعة باستخدام DreamBooth أو LoRA
def train_model(user_email, images, progress):
    if not user_email:
        return "يجب تسجيل الدخول لحفظ النموذج."
    
    user_model_id = f"user_model_{user_email}"
    
    # إعداد مجلد لحفظ الصور
    user_image_folder = Path(f"user_images/{user_email}")
    user_image_folder.mkdir(parents=True, exist_ok=True)

    # حفظ الصور في المجلد
    for img in images:
        img.save(user_image_folder / img.name)
    
    # إعداد بيانات التدريب (يمكنك تخصيص هذه الطريقة لتناسب النموذج الخاص بك)
    image_paths = [str(user_image_folder / img.name) for img in images]
    dataset = load_dataset('imagefolder', data_dir=str(user_image_folder))

    # تجهيز المعلمات لتدريب DreamBooth
    accelerator = Accelerator()

    # إعداد DreamBooth باستخدام مكتبة diffusers
    model = StableDiffusionPipeline.from_pretrained(MODEL_NAME)
    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")

    # تفعيل التدريب باستخدام DreamBooth
    trainer = DreamBoothTrainer(
        model=model,
        args=TrainingArguments(
            output_dir=f"user_models/{user_email}",
            per_device_train_batch_size=4,
            num_train_epochs=1,
            gradient_accumulation_steps=2,
            logging_dir='./logs',
            logging_steps=10,
            report_to="none"
        ),
        train_dataset=dataset["train"],
        tokenizer=tokenizer
    )
    
    # بداية عملية التدريب
    trainer.train()
    
    # حفظ النموذج المدرب
    model.save_pretrained(f"user_models/{user_email}")
    saved_models[user_email] = user_model_id

    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]
        model = StableDiffusionPipeline.from_pretrained(f"user_models/{user_email}")  # تحميل النموذج المدرب للمستخدم
        model.to(device)
        result = model(prompt).images[0]
    else:
        result = pipe(prompt).images[0]

    return result

# واجهة 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="🔔 نتيجة التدريب")
            progress_bar = gr.Progress()  # بدون label هنا
 
        with gr.Column():
            prompt_input = gr.Textbox(label="✏️ أدخل البرومبت لإنشاء صورة")
            generate_button = gr.Button("🎨 إنشاء صورة")
            output_image = gr.Image(label="📷 الصورة الناتجة")
 
    train_button.click(train_model, inputs=[user_email_input, image_input, progress_bar], 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)