app.py

import os
import gradio as gr
import numpy as np
import spaces
import torch
import random
import sqlite3
import json
from datetime import datetime
from PIL import Image, ImageFilter
from typing import Iterable
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes
from deep_translator import GoogleTranslator
from transformers import pipeline

# --- تنظیمات دیتابیس برای مدیریت کردیت کاربران ---
DB_NAME = "user_limits.db"

def init_db():
    """ایجاد جدول برای ذخیره تعداد استفاده کاربران"""
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    c.execute('''
        CREATE TABLE IF NOT EXISTS usage_logs (
            user_id TEXT,
            date TEXT,
            count INTEGER,
            PRIMARY KEY (user_id, date)
        )
    ''')
    conn.commit()
    conn.close()

init_db()

def check_and_update_quota(user_id, is_premium):
    """
    بررسی مجاز بودن کاربر برای تولید تصویر
    کاربران پریمیوم: همیشه مجاز
    کاربران رایگان: حداکثر 5 تصویر در روز
    """
    if not user_id:
        return True, "Guest" # اگر یوزر آیدی نبود، سختگیری نمی‌کنیم (یا می‌توان محدود کرد)

    if is_premium:
        return True, "نامحدود"

    today = datetime.now().strftime("%Y-%m-%d")
    
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    
    c.execute("SELECT count FROM usage_logs WHERE user_id = ? AND date = ?", (user_id, today))
    result = c.fetchone()
    
    current_count = result[0] if result else 0
    limit = 5
    
    conn.close()
    
    if current_count >= limit:
        return False, f"0 از {limit}"
    
    return True, f"{limit - current_count} از {limit}"

def increment_usage(user_id, is_premium):
    """افزایش شمارنده استفاده کاربر پس از تولید موفق"""
    if not user_id or is_premium:
        return

    today = datetime.now().strftime("%Y-%m-%d")
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    
    c.execute("INSERT OR IGNORE INTO usage_logs (user_id, date, count) VALUES (?, ?, 0)", (user_id, today))
    c.execute("UPDATE usage_logs SET count = count + 1 WHERE user_id = ? AND date = ?", (user_id, today))
    
    conn.commit()
    conn.close()

# --- تعریف تم ---
colors.steel_blue = colors.Color(
    name="steel_blue",
    c50="#EBF3F8",
    c100="#D3E5F0",
    c200="#A8CCE1",
    c300="#7DB3D2",
    c400="#529AC3",
    c500="#4682B4",
    c600="#3E72A0",
    c700="#36638C",
    c800="#2E5378",
    c900="#264364",
    c950="#1E3450",
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# --- بارگذاری سیستم امنیتی دوگانه (Dual Safety System) ---
print("Loading Safety Checkers...")

# مدل اول: سریع و استاندارد
safety_classifier_1 = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=-1)

# مدل دوم: بسیار دقیق و سخت‌گیر (ViT)
safety_classifier_2 = pipeline("image-classification", model="AdamCodd/vit-base-nsfw-detector", device=-1)

def is_image_nsfw(image):
    if image is None: return False
    try:
        # بررسی با مدل اول
        results1 = safety_classifier_1(image)
        for result in results1:
            if result['label'] == 'nsfw' and result['score'] > 0.5:
                print(f"Safety Check 1 Failed: {result['score']}")
                return True

        # بررسی با مدل دوم
        results2 = safety_classifier_2(image)
        for result in results2:
            label = result['label'].lower()
            score = result['score']
            
            if label == 'nsfw' and score > 0.3:
                print(f"Safety Check 2 (NSFW) Failed: {score}")
                return True
            
            if label in ['sexy', 'porn', 'hentai'] and score > 0.4:
                print(f"Safety Check 2 (Partial) Failed: {label} - {score}")
                return True

        return False
    except Exception as e:
        print(f"Safety check error: {e}")
        return True

# --- لیست کلمات ممنوعه ---
BANNED_WORDS = [
    "nsfw", "nude", "naked", "sex", "porn", "erotic", "xxx", "18+", "uncensored",
    "breast", "nipple", "areola", "cleavage", "topless", "open chest",
    "genital", "vagina", "penis", "dick", "cock", "pussy", "ass", "butt", "anus",
    "lingerie", "bikini", "swimwear", "underwear", "panties", "bra",
    "fetish", "bdsm", "bondage", "exhibitionism", "voyeur",
    "hentai", "ecchi", "ahegao", "paizuri",
    "undressed", "stripping", "naked body", "exposed skin", "sheer", "see-through",
    "rape", "violence", "blood", "gore", "sexual"
]

def check_text_safety(text):
    if not text: return True
    text_lower = text.lower()
    for word in BANNED_WORDS:
        if word in text_lower: 
            print(f"Banned word found: {word}")
            return False
    return True

def translate_prompt(text):
    if not text:
        return ""
    try:
        translated = GoogleTranslator(source='auto', target='en').translate(text)
        return translated
    except Exception as e:
        print(f"Translation Error: {e}")
        return text

# --- بارگذاری مدل اصلی ---
from diffusers import FlowMatchEulerDiscreteScheduler
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

print("Loading Generation Pipeline...")
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2509",
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "linoyts/Qwen-Image-Edit-Rapid-AIO",
        subfolder='transformer',
        torch_dtype=dtype,
        device_map='cuda'
    ),
    torch_dtype=dtype
).to(device)

# بارگذاری LoRA ها
pipe.load_lora_weights("autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime", weight_name="Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors", adapter_name="anime")
pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multiple-angles", weight_name="镜头转换.safetensors", adapter_name="multiple-angles")
pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Light_restoration", weight_name="移除光影.safetensors", adapter_name="light-restoration")
pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Relight", weight_name="Qwen-Edit-Relight.safetensors", adapter_name="relight")
pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multi-Angle-Lighting", weight_name="多角度灯光-251116.safetensors", adapter_name="multi-angle-lighting")
pipe.load_lora_weights("tlennon-ie/qwen-edit-skin", weight_name="qwen-edit-skin_1.1_000002750.safetensors", adapter_name="edit-skin")
pipe.load_lora_weights("lovis93/next-scene-qwen-image-lora-2509", weight_name="next-scene_lora-v2-3000.safetensors", adapter_name="next-scene")
pipe.load_lora_weights("vafipas663/Qwen-Edit-2509-Upscale-LoRA", weight_name="qwen-edit-enhance_64-v3_000001000.safetensors", adapter_name="upscale-image")

pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
MAX_SEED = np.iinfo(np.int32).max

LORA_MAPPING = {
    "تبدیل عکس به انیمه": "anime",
    "تغییر زاویه دید": "multiple-angles",
    "اصلاح نور و سایه": "light-restoration",
    "نورپردازی مجدد (Relight)": "relight",
    "نورپردازی چند زاویه‌ای": "multi-angle-lighting",
    "روتوش پوست": "edit-skin",
    "صحنه بعدی (سینمایی)": "next-scene",
    "افزایش کیفیت (Upscale)": "upscale-image"
}

ASPECT_RATIOS_LIST = [
    "خودکار (پیش‌فرض)",
    "۱:۱ (مربع - 1024x1024)",
    "۱۶:۹ (افقی - 1344x768)",
    "۹:۱۶ (عمودی - 768x1344)",
    "شخصی‌سازی (Custom)"
]

ASPECT_RATIOS_MAP = {
    "خودکار (پیش‌فرض)": "Auto",
    "۱:۱ (مربع - 1024x1024)": (1024, 1024),
    "۱۶:۹ (افقی - 1344x768)": (1344, 768),
    "۹:۱۶ (عمودی - 768x1344)": (768, 1344),
    "شخصی‌سازی (Custom)": "Custom"
}

def update_dimensions_on_upload(image):
    if image is None:
        return 1024, 1024
    original_width, original_height = image.size
    if original_width > original_height:
        new_width = 1024
        aspect_ratio = original_height / original_width
        new_height = int(new_width * aspect_ratio)
    else:
        new_height = 1024
        aspect_ratio = original_width / original_height
        new_width = int(new_height * aspect_ratio)
    new_width = (new_width // 8) * 8
    new_height = (new_height // 8) * 8
    return new_width, new_height

def get_error_html(message):
    return f"""
    <div style="background-color: #fee2e2; border: 1px solid #ef4444; color: #b91c1c; padding: 12px; border-radius: 8px; text-align: center; margin-bottom: 10px; font-weight: bold; display: flex; align-items: center; justify-content: center; gap: 8px;">
        <span style="font-size: 1.2em;">⛔</span>
        {message}
    </div>
    """

def get_success_html(message):
    return f"""
    <div style="background-color: #dcfce7; border: 1px solid #22c55e; color: #15803d; padding: 12px; border-radius: 8px; text-align: center; margin-bottom: 10px; font-weight: bold; display: flex; align-items: center; justify-content: center; gap: 8px;">
        <span style="font-size: 1.2em;">✅</span>
        {message}
    </div>
    """

def get_limit_error_html():
    return f"""
    <div style="background-color: #fff3cd; border: 1px solid #ffc107; color: #856404; padding: 16px; border-radius: 12px; text-align: center; margin-bottom: 15px; direction: rtl;">
        <h3 style="margin: 0 0 10px 0;">🚫 محدودیت اعتبار روزانه</h3>
        <p style="margin: 0;">شما از سهمیه ۵ تصویر رایگان امروز خود استفاده کرده‌اید.</p>
        <p style="margin: 8px 0 0 0; font-size: 0.9em;">برای استفاده نامحدود، حساب خود را به نسخه ویژه ارتقا دهید.</p>
        <button onclick="window.parent.postMessage({{type: 'NAVIGATE_TO_PREMIUM'}}, '*')" style="margin-top: 12px; background: linear-gradient(90deg, #FFD700, #FFA500); border: none; padding: 8px 16px; border-radius: 20px; color: black; font-weight: bold; cursor: pointer;">ارتقا به نسخه ویژه ⭐️</button>
    </div>
    """

@spaces.GPU(duration=30)
def infer(
    input_image,
    prompt,
    lora_adapter_persian,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    aspect_ratio_selection,
    custom_width,
    custom_height,
    user_id_hidden, # دریافت شناسه کاربر از فرانت
    is_premium_hidden, # دریافت وضعیت پریمیوم از فرانت
    progress=gr.Progress(track_tqdm=True)
):
    # --- بررسی اعتبار کاربر ---
    is_premium = str(is_premium_hidden).lower() == "true"
    allowed, limit_msg = check_and_update_quota(user_id_hidden, is_premium)
    
    if not allowed:
        return None, seed, get_limit_error_html()

    # --- بررسی وجود تصویر ---
    if input_image is None:
        return None, seed, get_error_html("لطفاً ابتدا یک تصویر بارگذاری کنید.")

    # --- بررسی امنیتی ورودی ---
    if is_image_nsfw(input_image):
        return None, seed, get_error_html("تصویر ورودی دارای محتوای نامناسب است و پردازش نمی‌شود.")

    # --- ترجمه و بررسی متن ---
    english_prompt = translate_prompt(prompt)
    if not check_text_safety(english_prompt):
        return None, seed, get_error_html("متن درخواست شامل کلمات غیرمجاز یا غیراخلاقی است.")

    adapter_internal_name = LORA_MAPPING.get(lora_adapter_persian)
    if adapter_internal_name:
        pipe.set_adapters([adapter_internal_name], adapter_weights=[1.0])
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    
    safety_negative = (
        "nsfw, nude, naked, porn, sexual, xxx, breast, nipple, areola, genital, "
        "vagina, penis, ass, lingerie, bikini, swimwear, underwear, fetish, "
        "topless, open chest, revealing clothes, cleavage, see-through, sheer, "
        "gore, violence, blood, navel, midriff, exposed skin, erotic, ecchi, hentai, "
        "uncensored, stripping"
    )
    base_negative = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"
    final_negative_prompt = f"{safety_negative}, {base_negative}"

    original_image = input_image.convert("RGB")
    
    selection_value = ASPECT_RATIOS_MAP.get(aspect_ratio_selection)
    
    if selection_value == "Custom":
        width = (int(custom_width) // 8) * 8
        height = (int(custom_height) // 8) * 8
    elif selection_value == "Auto" or selection_value is None:
        width, height = update_dimensions_on_upload(original_image)
    else:
        width, height = selection_value

    try:
        result = pipe(
            image=original_image,
            prompt=english_prompt,
            negative_prompt=final_negative_prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            generator=generator,
            true_cfg_scale=guidance_scale,
        ).images[0]

        # --- بررسی امنیتی خروجی ---
        if is_image_nsfw(result):
            return None, seed, get_error_html("تصویر تولید شده حاوی محتوای نامناسب بود و حذف شد.")
        
        # --- کسر اعتبار (فقط برای کاربران رایگان) ---
        increment_usage(user_id_hidden, is_premium)
        
        return result, seed, get_success_html("تصویر با موفقیت ویرایش شد.")

    except Exception as e:
        error_str = str(e)
        if "quota" in error_str.lower() or "exceeded" in error_str.lower():
             raise e 
        return None, seed, get_error_html(f"خطا در پردازش: {error_str}")

@spaces.GPU(duration=30)
def infer_example(input_image, prompt, lora_adapter):
    # برای مثال‌ها، آیدی کاربر فرضی و پریمیوم در نظر گرفته می‌شود تا محدودیت اعمال نشود
    res, s, status = infer(input_image, prompt, lora_adapter, 0, True, 1.0, 4, "خودکار (پیش‌فرض)", 1024, 1024, "example_user", "true")
    return res, s, status

# --- جاوااسکریپت برای دکمه دانلود ---
js_download_func = """
async (image) => {
    if (!image) {
        alert("لطفاً ابتدا تصویر را تولید کنید.");
        return;
    }
    let fileUrl = image.url;
    if (fileUrl && !fileUrl.startsWith('http')) {
       fileUrl = window.location.origin + fileUrl;
    } else if (!fileUrl && image.path) {
       fileUrl = window.location.origin + "/file=" + image.path;
    }
    window.parent.postMessage({
        type: 'DOWNLOAD_REQUEST',
        url: fileUrl
    }, '*');
}
"""

# --- جاوااسکریپت سراسری (شناسایی کاربر، پیام رفع محدودیت، دانلود) ---
js_global_content = """
<script>
document.addEventListener('DOMContentLoaded', () => {
    // --------------------------------------------------------
    // 1. User Identification & Credit Logic (PostMessage)
    // --------------------------------------------------------
    const PREMIUM_PAGE_ID = '1149636'; 
    let userId = "guest_" + Math.random().toString(36).substr(2, 9);
    let isPremium = false;

    // Listen for response from Parent (Android/Iframe)
    window.addEventListener('message', (event) => {
        if (event.data && event.data.type === 'USER_DATA_RESPONSE') {
            const payload = event.data.payload;
            
            if (payload) {
                try {
                    const userObj = JSON.parse(payload);
                    if (userObj.id) userId = userObj.id;
                    
                    // Check Premium Status
                    if (userObj.isLogin && userObj.accessible_pages) {
                        const pages = userObj.accessible_pages;
                        // Check if premium page ID exists in allowed pages
                        if (pages.includes(PREMIUM_PAGE_ID) || pages.includes(parseInt(PREMIUM_PAGE_ID))) {
                            isPremium = true;
                        }
                    }
                } catch (e) {
                    console.error("Error parsing user data:", e);
                }
            }
            updateUserStatusUI();
            updateHiddenInputs();
        }
    });

    // Request User Data on Load
    window.parent.postMessage({ type: 'REQUEST_USER_DATA' }, '*');

    function updateUserStatusUI() {
        const badge = document.getElementById('user-status-badge');
        if (!badge) return;

        if (isPremium) {
            badge.className = 'status-badge premium';
            badge.innerHTML = '⭐️ نسخه نامحدود ویژه';
        } else {
            badge.className = 'status-badge free';
            badge.innerHTML = '👤 نسخه رایگان (۵ تصویر در روز)';
        }
        badge.style.display = 'inline-flex';
    }

    function updateHiddenInputs() {
        // Find hidden textboxes in Gradio (hacky but standard for Gradio)
        // We look for elements with specific IDs assigned in Python
        const idInput = document.querySelector('#hidden_user_id textarea');
        const premiumInput = document.querySelector('#hidden_user_premium textarea');
        
        if (idInput) {
            idInput.value = userId;
            idInput.dispatchEvent(new Event('input', { bubbles: true }));
        }
        if (premiumInput) {
            premiumInput.value = isPremium ? "true" : "false";
            premiumInput.dispatchEvent(new Event('input', { bubbles: true }));
        }
    }
    
    // Retry finding elements if Gradio hasn't fully loaded them
    setTimeout(updateHiddenInputs, 2000);
    setTimeout(updateHiddenInputs, 5000);


    // --------------------------------------------------------
    // 2. Force Light Mode
    // --------------------------------------------------------
    const forceLight = () => {
        const body = document.querySelector('body');
        if (body) {
            body.classList.remove('dark');
            body.style.backgroundColor = '#f5f7fa';
            body.style.color = '#333333';
        }
        document.querySelectorAll('.dark').forEach(el => el.classList.remove('dark'));
    };
    forceLight();
    setInterval(forceLight, 1000);

    // --------------------------------------------------------
    // 3. Retry Functionality
    // --------------------------------------------------------
    window.retryGeneration = function() {
        const modal = document.getElementById('custom-quota-modal');
        if (modal) modal.remove();
        const runBtn = document.getElementById('run-btn');
        if(runBtn) runBtn.click();
    };
    
    window.closeErrorModal = function() {
         const modal = document.getElementById('custom-quota-modal');
         if (modal) modal.remove();
    };

    // --------------------------------------------------------
    // 4. Quota Modal (IP Reset Guide)
    // --------------------------------------------------------
    const showQuotaModal = () => {
        if (document.getElementById('custom-quota-modal')) return;

        const modalHtml = `
            <div id="custom-quota-modal" style="position: fixed; top: 0; left: 0; width: 100%; height: 100%; background: rgba(0,0,0,0.6); backdrop-filter: blur(5px); z-index: 99999; display: flex; align-items: center; justify-content: center; font-family: 'Vazirmatn', sans-serif;">
                <div class="ip-reset-guide-container">
                    <div class="guide-header">
                        <svg class="guide-header-icon" viewbox="0 0 100 100" xmlns="http://www.w3.org/2000/svg"> 
                            <defs><lineargradient id="grad1" x1="0%" y1="0%" x2="100%" y2="100%"><stop offset="0%" style="stop-color: #667eea; stop-opacity: 1;"></stop><stop offset="100%" style="stop-color: #764ba2; stop-opacity: 1;"></stop></lineargradient></defs> 
                            <circle cx="50" cy="50" r="45" fill="url(#grad1)" opacity="0.1"></circle>
                            <circle cx="50" cy="50" r="35" fill="none" stroke="url(#grad1)" stroke-width="2" opacity="0.3"></circle>
                            <path d="M35 50 L45 60 L65 40" stroke="url(#grad1)" stroke-width="4" fill="none" stroke-linecap="round" stroke-linejoin="round"></path>
                            <circle cx="65" cy="35" r="8" fill="#fee140"></circle>
                            <path d="M62 35 L68 35 M65 32 L65 38" stroke="white" stroke-width="2" stroke-linecap="round"></path> 
                        </svg>
                        <div>
                            <h2>یک قدم تا ساخت تصاویر جدید</h2>
                            <p>نیازمند تغییر نقطه دستیابی</p>
                        </div>
                    </div>
                    <div class="guide-content">
                        <div class="info-card">
                            <div class="info-card-header">
                                <svg class="info-card-icon" viewbox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M13 2L3 14h9l-1 8 10-12h-9l1-8z" fill="#667eea" opacity="0.2"></path><path d="M13 2L3 14h9l-1 8 10-12h-9l1-8z" stroke="#667eea" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"></path></svg> 
                                <span class="info-card-title">راه حل سریع</span>
                            </div>
                            <p>طبق ویدیو آموزشی پایین بین نقطه دستیابی جابجا شوید تلاش مجدد بزنید تا تصاویر مجدداً تولید بشه.</p>
                        </div>
                        <div class="summary-section">
                            <div class="summary-header">
                                <svg class="summary-icon" viewbox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><circle cx="12" cy="12" r="10" fill="#56ab2f" opacity="0.2"></circle><path d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2z" stroke="#56ab2f" stroke-width="2"></path><path d="M9 12l2 2 4-4" stroke="#56ab2f" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"></path></svg> 
                                <span class="summary-title">خلاصه راهنما</span>
                            </div>
                            <div class="summary-text">هربار که این صفحه را مشاهده کردید: از اینترنت سیم‌کارت استفاده کنید، VPN را خاموش کرده و طبق ویدیو آموزشی پایین نقطه دستیابی رو تغییر دهید. «تلاش مجدد» کلیک کنید.</div>
                        </div>
                        <div class="video-button-container">
                            <button onclick="parent.postMessage({ type: 'NAVIGATE_TO_URL', url: '#/nav/online/news/getSingle/1149635/eyJpdiI6IjhHVGhPQWJwb3E0cjRXbnFWTW5BaUE9PSIsInZhbHVlIjoiS1V0dTdvT21wbXAwSXZaK1RCTG1pVXZqdlFJa1hXV1RKa2FLem9zU3pXMjd5MmlVOGc2YWY0NVdNR3h3Smp1aSIsIm1hYyI6IjY1NTA5ZDYzMjAzMTJhMGQyMWQ4NjA4ZDgyNGZjZDVlY2MyNjdiMjA2NWYzOWRjY2M4ZmVjYWRlMWNlMWQ3ODEiLCJ0YWciOiIifQ==/21135210' }, '*')" class="elegant-video-button"> 
                                <svg class="elegant-video-button-icon" xmlns="http://www.w3.org/2000/svg" viewbox="0 0 24 24"><path d="M8 5v14l11-7z"></path></svg> 
                                <span>دیدن ویدیو آموزشی استفاده نامحدود</span> 
                            </button>
                        </div>
                    </div>
                    <div class="guide-actions">
                        <button class="action-button back-button" onclick="window.closeErrorModal()"> 
                            <svg class="action-button-icon" viewbox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M19 12H5M12 19l-7-7 7-7" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"></path></svg> 
                            <span>بازگشت</span> 
                        </button>
                        <button class="action-button retry-button" onclick="window.retryGeneration()"> 
                            <svg class="action-button-icon" viewbox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M23 4v6h-6M1 20v-6h6" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"></path><path d="M20.49 9A9 9 0 0 0 5.64 5.64L1 10m22 4l-4.64 4.36A9 9 0 0 1 3.51 15" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"></path></svg> 
                            <span>تلاش مجدد</span> 
                        </button>
                    </div>
                </div>
            </div>
        `;
        document.body.insertAdjacentHTML('beforeend', modalHtml);
        setTimeout(() => { window.closeErrorModal(); }, 10000);
    };

    // --------------------------------------------------------
    // 5. Error Scanner
    // --------------------------------------------------------
    setInterval(() => {
        const potentialErrors = document.querySelectorAll('.toast-body, .error, .toast-wrap, .eta-bar, div[class*="error"]');
        potentialErrors.forEach(el => {
            const text = el.innerText || "";
            if (text.toLowerCase().includes('quota') || text.toLowerCase().includes('exceeded')) {
                 showQuotaModal();
                 el.style.display = 'none';
                 el.style.opacity = '0';
                 el.innerText = '';
                 const parentWrap = el.closest('.toast-wrap');
                 if(parentWrap) parentWrap.style.display = 'none';
            }
        });
    }, 100);
});
</script>
"""

# --- CSS Updated ---
css_code = """
<style>
@import url('https://fonts.googleapis.com/css2?family=Vazirmatn:wght@300;400;500;700&display=swap');

:root, .dark, body, .gradio-container {
    --body-background-fill: #f5f7fa !important;
    --body-text-color: #1f2937 !important;
    font-family: 'Vazirmatn', sans-serif !important;
}

/* User Status Badge CSS */
.status-badge {
    display: none; /* Initially hidden until JS detects user */
    align-items: center;
    justify-content: center;
    padding: 8px 16px;
    border-radius: 50px;
    font-size: 0.95em;
    font-weight: 700;
    margin: 10px auto 20px auto;
    width: fit-content;
    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);
    animation: fadeInBadge 0.5s ease-out;
}

.status-badge.premium {
    background: linear-gradient(135deg, #FFD700 0%, #FDB931 100%);
    color: #333;
    border: 2px solid #fff;
    box-shadow: 0 0 15px rgba(255, 215, 0, 0.4);
}

.status-badge.free {
    background: #e2e8f0;
    color: #475569;
    border: 1px solid #cbd5e1;
}

@keyframes fadeInBadge {
    from { opacity: 0; transform: translateY(-10px); }
    to { opacity: 1; transform: translateY(0); }
}

/* Existing CSS for Guide, Buttons, etc. */
.ip-reset-guide-container {
    text-align: right; direction: rtl; background: rgba(255, 255, 255, 0.98); 
    backdrop-filter: blur(10px); padding: 20px; border-radius: 16px;
    box-shadow: 0 20px 25px -5px rgba(0,0,0,0.1); border: 1px solid rgba(102, 126, 234, 0.2); 
    width: 90%; max-width: 420px; max-height: 90vh; overflow-y: auto; 
    font-family: 'Vazirmatn', sans-serif !important;
}
.guide-header { display: flex; align-items: center; margin-bottom: 15px; }
.guide-header-icon { width: 45px; height: 45px; margin-left: 15px; }
.guide-header h2 { font-size: 1.2rem; margin: 0; }
.guide-content { font-size: 0.9rem; line-height: 1.6; }
.info-card, .summary-section { margin-top: 12px; padding: 12px; border-radius: 12px; }
.info-card { background: #f0f4ff; border: 1px solid #c3dafe; }
.summary-section { background: #f0fff4; border: 1px solid #c6f6d5; }
.video-button-container { text-align: center; margin-top: 15px; }
.elegant-video-button { padding: 8px 16px; border-radius: 20px; border: 1px solid #ddd; background: white; cursor: pointer; display: inline-flex; align-items: center; }
.guide-actions { display: flex; gap: 10px; margin-top: 20px; }
.action-button { flex: 1; padding: 10px; border-radius: 10px; border: none; cursor: pointer; display: flex; align-items: center; justify-content: center; font-weight: bold; }
.back-button { background: white; border: 1px solid #ddd; }
.retry-button { background: #667eea; color: white; }

/* Main UI Tweaks */
#col-container { max-width: 980px; margin: 0 auto; direction: rtl; text-align: right; background: white; padding: 20px; border-radius: 20px; }
#main-title h1 { text-align: center; color: #2563eb; }
#main-description { text-align: center; color: #4b5563; }
.primary-btn { background: linear-gradient(135deg, #10b981, #059669); color: white; border: none; }
footer { display: none !important; }
</style>
"""

combined_html = css_code + js_global_content

with gr.Blocks() as demo:
    gr.HTML(combined_html)
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# **ویرایشگر هوشمند آلفا**", elem_id="main-title")
        gr.Markdown("با هوش مصنوعی آلفا تصاویر تونو به مدل های مختلف ویرایش کنید.", elem_id="main-description")
        
        # بج وضعیت کاربر (HTML خالی که توسط JS پر می‌شود)
        gr.HTML('<div id="user-status-badge" class="status-badge"></div>')

        # فیلدهای مخفی برای انتقال اطلاعات کاربر از JS به پایتون
        user_id_storage = gr.Textbox(visible=False, elem_id="hidden_user_id")
        premium_status_storage = gr.Textbox(visible=False, elem_id="hidden_user_premium")

        with gr.Row(equal_height=True):
            with gr.Column():
                input_image = gr.Image(label="بارگذاری تصویر", type="pil", height=320)
                prompt = gr.Text(label="دستور ویرایش (به فارسی)", placeholder="مثال: تصویر را به سبک انیمه تبدیل کن...", rtl=True, lines=3)
                status_box = gr.HTML(label="وضعیت")
                run_button = gr.Button("✨ شروع پردازش", variant="primary", elem_classes="primary-btn", elem_id="run-btn")

            with gr.Column():
                output_image = gr.Image(label="تصویر نهایی", interactive=False, format="png", height=380)
                download_button = gr.Button("📥 دانلود و ذخیره تصویر", variant="secondary", elem_id="download-btn", elem_classes="primary-btn")
                
                with gr.Row():
                    lora_adapter = gr.Dropdown(label="انتخاب سبک", choices=list(LORA_MAPPING.keys()), value="تبدیل عکس به انیمه")
                
                with gr.Accordion("تنظیمات پیشرفته", open=False):
                    aspect_ratio_selection = gr.Dropdown(label="ابعاد", choices=ASPECT_RATIOS_LIST, value="خودکار (پیش‌فرض)")
                    with gr.Row(visible=False) as custom_dims_row:
                        custom_width = gr.Slider(label="عرض", minimum=256, maximum=2048, step=8, value=1024)
                        custom_height = gr.Slider(label="ارتفاع", minimum=256, maximum=2048, step=8, value=1024)
                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                    randomize_seed = gr.Checkbox(label="Seed تصادفی", value=True)
                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=4)
        
        # لاجیک نمایش تنظیمات ابعاد
        def toggle_row(choice):
            return gr.update(visible=True) if choice == "شخصی‌سازی (Custom)" else gr.update(visible=False)

        aspect_ratio_selection.change(fn=toggle_row, inputs=aspect_ratio_selection, outputs=custom_dims_row)

        # اجرای مدل
        run_button.click(
            fn=infer,
            inputs=[
                input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps, 
                aspect_ratio_selection, custom_width, custom_height,
                user_id_storage, premium_status_storage  # ارسال اطلاعات کاربر به تابع infer
            ],
            outputs=[output_image, seed, status_box],
            api_name="predict"
        )

        download_button.click(fn=None, inputs=[output_image], outputs=None, js=js_download_func)
        
        # نمونه‌ها
        gr.Examples(
            examples=[
                ["examples/1.jpg", "تبدیل به انیمه کن.", "تبدیل عکس به انیمه"],
                ["examples/5.jpg", "سایه‌ها را حذف کن.", "اصلاح نور و سایه"],
            ],
            inputs=[input_image, prompt, lora_adapter],
            outputs=[output_image, seed, status_box],
            fn=infer_example,
            cache_examples=False,
            label="نمونه‌ها"
        )

if __name__ == "__main__":
    demo.queue(max_size=30).launch(show_error=True)