app.py

import gradio as gr
import requests
import json
import os
import time
import sqlite3
import hashlib
import secrets
from datetime import datetime, timedelta

# ============================================
# 환경변수
# ============================================
FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY", "")
FAL_KEY = os.getenv("FAL_KEY", "")
ADMIN_EMAIL = os.getenv("ADMIN_EMAIL", "arxivgpt@gmail.com")
SECRET_KEY = os.getenv("SECRET_KEY", "change-this-secret-key-in-production")

# Persistent Storage 경로 (HuggingFace Spaces)
DATA_DIR = "/data" if os.path.exists("/data") else "./data"
DB_PATH = os.path.join(DATA_DIR, "app.db")

# 설정
DAILY_LIMIT_FREE = 10
SESSION_EXPIRE_HOURS = 24 * 7

# ============================================
# 데이터베이스
# ============================================
def get_db():
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    os.makedirs(DATA_DIR, exist_ok=True)
    conn = get_db()
    cursor = conn.cursor()
    
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            email TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            is_admin INTEGER DEFAULT 0,
            is_active INTEGER DEFAULT 1,
            daily_limit INTEGER DEFAULT 10,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            last_login TIMESTAMP
        )
    ''')
    
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS sessions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            token TEXT UNIQUE NOT NULL,
            expires_at TIMESTAMP NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id)
        )
    ''')
    
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS generations (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            original_prompt TEXT,
            enhanced_prompt TEXT,
            image_url TEXT,
            status TEXT DEFAULT 'success',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id)
        )
    ''')
    
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS daily_usage (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            date DATE NOT NULL,
            count INTEGER DEFAULT 0,
            UNIQUE(user_id, date),
            FOREIGN KEY (user_id) REFERENCES users (id)
        )
    ''')
    
    conn.commit()
    
    # 기본 관리자 계정
    cursor.execute("SELECT id FROM users WHERE email = ?", (ADMIN_EMAIL,))
    if not cursor.fetchone():
        admin_password = secrets.token_urlsafe(12)
        password_hash = hash_password(admin_password)
        cursor.execute(
            "INSERT INTO users (email, password_hash, is_admin, daily_limit) VALUES (?, ?, 1, 9999)",
            (ADMIN_EMAIL, password_hash)
        )
        conn.commit()
        print(f"[ADMIN] 관리자 계정: {ADMIN_EMAIL} / {admin_password}")
    
    conn.close()

# ============================================
# 인증 유틸리티
# ============================================
def hash_password(password: str) -> str:
    return hashlib.sha256((password + SECRET_KEY).encode()).hexdigest()

def verify_password(password: str, password_hash: str) -> bool:
    return hash_password(password) == password_hash

def generate_session_token() -> str:
    return secrets.token_urlsafe(32)

def create_session(user_id: int) -> str:
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute("DELETE FROM sessions WHERE user_id = ?", (user_id,))
    token = generate_session_token()
    expires_at = datetime.now() + timedelta(hours=SESSION_EXPIRE_HOURS)
    cursor.execute(
        "INSERT INTO sessions (user_id, token, expires_at) VALUES (?, ?, ?)",
        (user_id, token, expires_at)
    )
    cursor.execute("UPDATE users SET last_login = ? WHERE id = ?", (datetime.now(), user_id))
    conn.commit()
    conn.close()
    return token

def validate_session(token: str) -> dict:
    if not token:
        return None
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('''
        SELECT u.id, u.email, u.is_admin, u.daily_limit, u.is_active
        FROM sessions s JOIN users u ON s.user_id = u.id
        WHERE s.token = ? AND s.expires_at > ? AND u.is_active = 1
    ''', (token, datetime.now()))
    row = cursor.fetchone()
    conn.close()
    if row:
        return {"id": row["id"], "email": row["email"], "is_admin": bool(row["is_admin"]), "daily_limit": row["daily_limit"]}
    return None

def delete_session(token: str):
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute("DELETE FROM sessions WHERE token = ?", (token,))
    conn.commit()
    conn.close()

# ============================================
# 사용량 관리
# ============================================
def get_daily_usage(user_id: int) -> int:
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute("SELECT count FROM daily_usage WHERE user_id = ? AND date = ?", (user_id, datetime.now().date()))
    row = cursor.fetchone()
    conn.close()
    return row["count"] if row else 0

def increment_usage(user_id: int) -> int:
    conn = get_db()
    cursor = conn.cursor()
    today = datetime.now().date()
    cursor.execute('''
        INSERT INTO daily_usage (user_id, date, count) VALUES (?, ?, 1)
        ON CONFLICT(user_id, date) DO UPDATE SET count = count + 1
    ''', (user_id, today))
    conn.commit()
    cursor.execute("SELECT count FROM daily_usage WHERE user_id = ? AND date = ?", (user_id, today))
    row = cursor.fetchone()
    conn.close()
    return row["count"]

def check_usage_limit(user_id: int, daily_limit: int) -> tuple:
    current = get_daily_usage(user_id)
    return (current < daily_limit, current, daily_limit)

# ============================================
# 생성 이력
# ============================================
def save_generation(user_id: int, original: str, enhanced: str, image_url: str, status: str = "success"):
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('''
        INSERT INTO generations (user_id, original_prompt, enhanced_prompt, image_url, status)
        VALUES (?, ?, ?, ?, ?)
    ''', (user_id, original, enhanced, image_url, status))
    conn.commit()
    conn.close()

def get_user_generations(user_id: int, limit: int = 20) -> list:
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('''
        SELECT original_prompt, enhanced_prompt, image_url, status, created_at
        FROM generations WHERE user_id = ? ORDER BY created_at DESC LIMIT ?
    ''', (user_id, limit))
    rows = cursor.fetchall()
    conn.close()
    return [dict(row) for row in rows]

# ============================================
# 관리자 기능
# ============================================
def get_all_users() -> list:
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute('SELECT id, email, is_admin, is_active, daily_limit, created_at, last_login FROM users ORDER BY created_at DESC')
    rows = cursor.fetchall()
    conn.close()
    return [dict(row) for row in rows]

def get_stats() -> dict:
    conn = get_db()
    cursor = conn.cursor()
    today = datetime.now().date()
    
    cursor.execute("SELECT COUNT(*) as count FROM users")
    total_users = cursor.fetchone()["count"]
    
    cursor.execute("SELECT COUNT(DISTINCT user_id) as count FROM daily_usage WHERE date = ?", (today,))
    active_today = cursor.fetchone()["count"]
    
    cursor.execute("SELECT COUNT(*) as count FROM generations")
    total_generations = cursor.fetchone()["count"]
    
    cursor.execute("SELECT COUNT(*) as count FROM generations WHERE DATE(created_at) = ?", (today,))
    generations_today = cursor.fetchone()["count"]
    
    conn.close()
    return {"total_users": total_users, "active_today": active_today, "total_generations": total_generations, "generations_today": generations_today}

def update_user_limit(user_id: int, daily_limit: int) -> bool:
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute("UPDATE users SET daily_limit = ? WHERE id = ?", (daily_limit, user_id))
    conn.commit()
    conn.close()
    return True

# ============================================
# 회원가입 / 로그인
# ============================================
def register_user(email: str, password: str, confirm_password: str) -> tuple:
    if not email or not password:
        return None, "❌ 이메일과 비밀번호를 입력해주세요."
    if password != confirm_password:
        return None, "❌ 비밀번호가 일치하지 않습니다."
    if len(password) < 6:
        return None, "❌ 비밀번호는 6자 이상이어야 합니다."
    if "@" not in email:
        return None, "❌ 올바른 이메일 형식이 아닙니다."
    
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute("SELECT id FROM users WHERE email = ?", (email,))
    if cursor.fetchone():
        conn.close()
        return None, "❌ 이미 등록된 이메일입니다."
    
    password_hash = hash_password(password)
    cursor.execute("INSERT INTO users (email, password_hash, daily_limit) VALUES (?, ?, ?)", (email, password_hash, DAILY_LIMIT_FREE))
    conn.commit()
    user_id = cursor.lastrowid
    conn.close()
    
    token = create_session(user_id)
    return token, f"✅ 회원가입 완료! 환영합니다, {email}"

def login_user(email: str, password: str) -> tuple:
    if not email or not password:
        return None, "❌ 이메일과 비밀번호를 입력해주세요."
    
    conn = get_db()
    cursor = conn.cursor()
    cursor.execute("SELECT id, password_hash, is_active FROM users WHERE email = ?", (email,))
    row = cursor.fetchone()
    conn.close()
    
    if not row:
        return None, "❌ 등록되지 않은 이메일입니다."
    if not row["is_active"]:
        return None, "❌ 비활성화된 계정입니다."
    if not verify_password(password, row["password_hash"]):
        return None, "❌ 비밀번호가 올바르지 않습니다."
    
    token = create_session(row["id"])
    return token, "✅ 로그인 성공!"

def logout_user(token: str) -> str:
    if token:
        delete_session(token)
    return "✅ 로그아웃되었습니다."

# ============================================
# API 함수들
# ============================================
def enhance_prompt(prompt: str, fireworks_key: str) -> str:
    if not prompt.strip():
        return "❌ 프롬프트를 입력해주세요."
    if not fireworks_key.strip():
        return "❌ Fireworks API 키가 필요합니다."
    
    system_message = """You are a professional prompt engineer for AI image generation.
Enhance prompts with: visual details, lighting, composition, style, mood, camera angles, color palette, quality modifiers.
Output ONLY the enhanced prompt in English, no explanations."""

    url = "https://api.fireworks.ai/inference/v1/chat/completions"
    payload = {
        "model": "accounts/fireworks/models/gpt-oss-120b",
        "max_tokens": 4096,
        "top_p": 1,
        "top_k": 40,
        "presence_penalty": 0,
        "frequency_penalty": 0,
        "temperature": 0.6,
        "messages": [
            {"role": "system", "content": system_message},
            {"role": "user", "content": f"Enhance this prompt:\n\n\"{prompt}\""}
        ]
    }
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {fireworks_key}"
    }
    
    try:
        response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=60)
        response.raise_for_status()
        data = response.json()
        enhanced = data.get("choices", [{}])[0].get("message", {}).get("content", "")
        return enhanced if enhanced else "❌ 증강 실패"
    except requests.exceptions.HTTPError as e:
        return f"❌ HTTP 오류: {e.response.status_code}"
    except Exception as e:
        return f"❌ 오류: {str(e)}"

def generate_image(prompt: str, fal_key: str, aspect_ratio: str = "1:1", resolution: str = "1K") -> str:
    if not prompt.strip() or not fal_key.strip():
        return None
    
    # FAL API - nano-banana-pro
    url = "https://fal.run/fal-ai/nano-banana-pro"
    headers = {"Authorization": f"Key {fal_key}", "Content-Type": "application/json"}
    
    payload = {
        "prompt": prompt,
        "num_images": 1,
        "aspect_ratio": aspect_ratio,
        "output_format": "png",
        "resolution": resolution
    }
    
    try:
        response = requests.post(url, headers=headers, json=payload, timeout=120)
        response.raise_for_status()
        data = response.json()
        if "images" in data and data["images"]:
            return data["images"][0].get("url")
    except Exception as e:
        print(f"이미지 생성 오류: {e}")
    return None

def process_comparison_with_auth(original_prompt, session_token, fireworks_key, fal_key, aspect_ratio, resolution):
    user = validate_session(session_token)
    if not user:
        yield "❌ 로그인이 필요합니다.", None, None, None, ""
        return
    
    allowed, current, limit = check_usage_limit(user["id"], user["daily_limit"])
    if not allowed:
        yield f"❌ 일일 한도({limit}회) 초과", None, None, None, f"{current}/{limit}"
        return
    
    fw_key = fireworks_key if fireworks_key.strip() else FIREWORKS_API_KEY
    fl_key = fal_key if fal_key.strip() else FAL_KEY
    
    yield f"⚡ 처리 중... ({current}/{limit})", None, None, None, f"{current}/{limit}"
    
    enhanced = enhance_prompt(original_prompt, fw_key)
    if enhanced.startswith("❌"):
        yield enhanced, None, None, None, f"{current}/{limit}"
        return
    
    yield "✅ 프롬프트 증강 완료\n⚡ 원본 이미지 생성 중...", enhanced, None, None, f"{current}/{limit}"
    original_image = generate_image(original_prompt, fl_key, aspect_ratio, resolution)
    
    yield "✅ 원본 이미지 완료\n⚡ 증강 이미지 생성 중...", enhanced, original_image, None, f"{current}/{limit}"
    enhanced_image = generate_image(enhanced, fl_key, aspect_ratio, resolution)
    
    new_count = increment_usage(user["id"])
    save_generation(user["id"], original_prompt, enhanced, enhanced_image or original_image, "success" if enhanced_image else "partial")
    
    yield f"✅ 완료! 사용량: {new_count}/{limit}", enhanced, original_image, enhanced_image, f"{new_count}/{limit}"

# ============================================
# CSS - 밝은 프로페셔널 테마
# ============================================
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=Poppins:wght@400;500;600;700&display=swap');

:root {
    --primary-blue: #2563eb;
    --primary-blue-light: #3b82f6;
    --primary-blue-dark: #1d4ed8;
    --accent-purple: #7c3aed;
    --accent-green: #10b981;
    --accent-orange: #f59e0b;
    --bg-white: #ffffff;
    --bg-gray-50: #f9fafb;
    --bg-gray-100: #f3f4f6;
    --bg-gray-200: #e5e7eb;
    --text-primary: #111827;
    --text-secondary: #4b5563;
    --text-muted: #9ca3af;
    --border-color: #e5e7eb;
    --shadow-sm: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
    --shadow-md: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
    --shadow-lg: 0 10px 15px -3px rgba(0, 0, 0, 0.1), 0 4px 6px -2px rgba(0, 0, 0, 0.05);
    --radius: 12px;
}

/* 메인 컨테이너 - 밝은 배경 */
.gradio-container {
    background: linear-gradient(135deg, var(--bg-gray-50) 0%, var(--bg-white) 50%, var(--bg-gray-100) 100%) !important;
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
    min-height: 100vh;
}

/* 타이틀 스타일 */
.cyber-title {
    font-family: 'Poppins', sans-serif !important;
    font-size: 2.5rem !important;
    font-weight: 700 !important;
    background: linear-gradient(135deg, var(--primary-blue) 0%, var(--accent-purple) 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    text-align: center;
    letter-spacing: -0.5px;
}

/* 입력 필드 - 밝은 스타일 */
.gradio-container textarea,
.gradio-container input[type="text"],
.gradio-container input[type="password"],
.gradio-container input[type="email"] {
    background: var(--bg-white) !important;
    border: 2px solid var(--border-color) !important;
    border-radius: var(--radius) !important;
    color: var(--text-primary) !important;
    font-family: 'Inter', sans-serif !important;
    font-size: 0.95rem !important;
    padding: 12px 16px !important;
    transition: all 0.2s ease !important;
}

.gradio-container textarea:focus,
.gradio-container input:focus {
    border-color: var(--primary-blue) !important;
    box-shadow: 0 0 0 3px rgba(37, 99, 235, 0.1) !important;
    outline: none !important;
}

.gradio-container textarea::placeholder,
.gradio-container input::placeholder {
    color: var(--text-muted) !important;
}

/* 라벨 스타일 */
.gradio-container label {
    font-family: 'Inter', sans-serif !important;
    color: var(--text-primary) !important;
    font-weight: 600 !important;
    font-size: 0.9rem !important;
    margin-bottom: 6px !important;
}

/* 드롭다운 */
.gradio-container select,
.gradio-container .wrap-inner {
    background: var(--bg-white) !important;
    border: 2px solid var(--border-color) !important;
    border-radius: var(--radius) !important;
    color: var(--text-primary) !important;
}

/* 메인 버튼 - 블루 그라디언트 */
.neon-button {
    background: linear-gradient(135deg, var(--primary-blue) 0%, var(--primary-blue-dark) 100%) !important;
    border: none !important;
    border-radius: var(--radius) !important;
    color: #ffffff !important;
    font-family: 'Inter', sans-serif !important;
    font-size: 0.95rem !important;
    font-weight: 600 !important;
    padding: 12px 24px !important;
    box-shadow: var(--shadow-md), 0 4px 14px rgba(37, 99, 235, 0.25) !important;
    transition: all 0.2s ease !important;
}

.neon-button:hover {
    background: linear-gradient(135deg, var(--primary-blue-light) 0%, var(--primary-blue) 100%) !important;
    box-shadow: var(--shadow-lg), 0 6px 20px rgba(37, 99, 235, 0.35) !important;
    transform: translateY(-2px) !important;
}

/* 일반 버튼 */
.gradio-container button {
    border-radius: var(--radius) !important;
    font-family: 'Inter', sans-serif !important;
    transition: all 0.2s ease !important;
}

/* 탭 스타일 */
.gradio-container .tabs {
    background: var(--bg-white) !important;
    border-radius: 16px !important;
    padding: 8px !important;
    box-shadow: var(--shadow-sm) !important;
}

.gradio-container .tab-nav button {
    background: transparent !important;
    color: var(--text-secondary) !important;
    border-radius: 10px !important;
    font-weight: 500 !important;
    padding: 10px 20px !important;
}

.gradio-container .tab-nav button.selected {
    background: var(--primary-blue) !important;
    color: #ffffff !important;
}

/* 아코디언 */
.gradio-container .accordion {
    background: var(--bg-white) !important;
    border: 1px solid var(--border-color) !important;
    border-radius: var(--radius) !important;
    box-shadow: var(--shadow-sm) !important;
}

/* 마크다운 텍스트 */
.gradio-container .markdown-text {
    color: var(--text-primary) !important;
}

.gradio-container .markdown-text h1,
.gradio-container .markdown-text h2,
.gradio-container .markdown-text h3 {
    color: var(--text-primary) !important;
}

/* 테이블 */
.gradio-container table {
    background: var(--bg-white) !important;
    border-radius: var(--radius) !important;
    overflow: hidden !important;
}

.gradio-container th {
    background: var(--bg-gray-100) !important;
    color: var(--text-primary) !important;
    font-weight: 600 !important;
}

.gradio-container td {
    color: var(--text-secondary) !important;
    border-color: var(--border-color) !important;
}

/* 이미지 컨테이너 */
.gradio-container .image-container {
    background: var(--bg-white) !important;
    border-radius: var(--radius) !important;
    box-shadow: var(--shadow-md) !important;
}

/* 카드 스타일 박스 */
.gradio-container .block {
    background: var(--bg-white) !important;
    border: 1px solid var(--border-color) !important;
    border-radius: var(--radius) !important;
}

/* 슬라이더 */
.gradio-container input[type="range"] {
    accent-color: var(--primary-blue) !important;
}

/* 체크박스 */
.gradio-container input[type="checkbox"] {
    accent-color: var(--primary-blue) !important;
}

/* 상태 메시지 */
.gradio-container .message {
    border-radius: var(--radius) !important;
}

/* 성공/에러 색상 */
.success-text {
    color: var(--accent-green) !important;
}

.error-text {
    color: #ef4444 !important;
}

/* 데이터프레임 */
.gradio-container .dataframe {
    border-radius: var(--radius) !important;
    overflow: hidden !important;
}

/* 스크롤바 */
::-webkit-scrollbar {
    width: 8px;
    height: 8px;
}

::-webkit-scrollbar-track {
    background: var(--bg-gray-100);
    border-radius: 4px;
}

::-webkit-scrollbar-thumb {
    background: var(--bg-gray-300);
    border-radius: 4px;
}

::-webkit-scrollbar-thumb:hover {
    background: var(--text-muted);
}

/* 허깅페이스 배지 숨기기 */
.built-with,
.built-with-badge,
a[href*="huggingface.co/spaces"],
footer,
.footer,
#footer,
.gradio-container > footer,
div[class*="footer"],
.space-info,
.hf-space-header,
[class*="space-header"],
div.wrap.svelte-1rjryqp,
.svelte-1rjryqp:has(a[href*="huggingface"]),
a.svelte-1rjryqp[href*="huggingface"] {
    display: none !important;
    visibility: hidden !important;
    opacity: 0 !important;
    height: 0 !important;
    width: 0 !important;
    position: absolute !important;
    pointer-events: none !important;
}

div[style*="position: fixed"][style*="right"],
div[style*="position: fixed"][style*="top: 0"] {
    display: none !important;
}
"""

# ============================================
# Gradio UI
# ============================================
init_db()

examples = [
    ["한복을 입은 여성이 전통 한옥 마당에서 부채를 들고 있다"],
    ["사이버펑크 도시의 네온사인이 빛나는 밤거리"],
    ["바다 위 일몰과 작은 돛단배"],
]

with gr.Blocks(title="AI PROMPT ENHANCER", css=CUSTOM_CSS) as demo:
    session_token = gr.State(value="")
    
    gr.HTML('''
    <div style="text-align:center; padding:40px 20px; background: linear-gradient(135deg, #ffffff 0%, #f0f9ff 100%); border-radius: 16px; margin-bottom: 20px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);">
        <h1 class="cyber-title">🚀 PROMPT ENHANCER</h1>
        <p style="color: #4b5563; font-size: 1.1rem; margin-top: 8px;">AI-Powered Image Generation Platform</p>
    </div>
    ''')
    
    with gr.Tabs() as tabs:
        # Tab 1: 로그인
        with gr.Tab("🔐 로그인"):
            with gr.Row():
                with gr.Column(scale=1):
                    pass
                with gr.Column(scale=2):
                    with gr.Tabs():
                        with gr.Tab("로그인"):
                            login_email = gr.Textbox(label="이메일", placeholder="your@email.com")
                            login_password = gr.Textbox(label="비밀번호", type="password")
                            login_btn = gr.Button("로그인", elem_classes=["neon-button"])
                            login_status = gr.Markdown("")
                        with gr.Tab("회원가입"):
                            reg_email = gr.Textbox(label="이메일", placeholder="your@email.com")
                            reg_password = gr.Textbox(label="비밀번호", type="password", placeholder="6자 이상")
                            reg_confirm = gr.Textbox(label="비밀번호 확인", type="password")
                            reg_btn = gr.Button("회원가입", elem_classes=["neon-button"])
                            reg_status = gr.Markdown("")
                    current_user_display = gr.Markdown("로그인하지 않음")
                    logout_btn = gr.Button("로그아웃", visible=False)
                with gr.Column(scale=1):
                    pass
        
        # Tab 2: 메인 서비스
        with gr.Tab("✨ 이미지 생성"):
            usage_display = gr.Markdown("사용량: 로그인 필요")
            
            with gr.Accordion("⚙️ API 설정", open=False):
                with gr.Row():
                    fireworks_key_input = gr.Textbox(label="FIREWORKS API KEY", placeholder="환경변수 사용", type="password")
                    fal_key_input = gr.Textbox(label="FAL API KEY", placeholder="환경변수 사용", type="password")
            
            with gr.Row():
                with gr.Column(scale=2):
                    prompt_input = gr.Textbox(label="프롬프트 입력", placeholder="생성할 이미지를 설명해주세요...", lines=4)
                    gr.Examples(examples=examples, inputs=prompt_input, label="예시")
                with gr.Column(scale=1):
                    aspect_ratio = gr.Dropdown(label="비율", choices=["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "21:9"], value="1:1")
                    resolution = gr.Dropdown(label="해상도", choices=["1K", "2K", "4K"], value="1K")
                    generate_btn = gr.Button("⚡ 생성하기", elem_classes=["neon-button"], size="lg")
            
            status_text = gr.Textbox(label="상태", interactive=False, lines=2)
            enhanced_output = gr.Textbox(label="증강된 프롬프트", lines=3)
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 📌 원본")
                    original_image_output = gr.Image(label="", height=400)
                with gr.Column():
                    gr.Markdown("### ✨ 증강")
                    enhanced_image_output = gr.Image(label="", height=400)
        
        # Tab 3: 내 계정
        with gr.Tab("👤 내 계정") as account_tab:
            account_info = gr.Markdown("로그인이 필요합니다.")
            history_display = gr.Dataframe(headers=["프롬프트", "생성일시", "상태"], interactive=False)
            refresh_history_btn = gr.Button("🔄 새로고침")
        
        # Tab 4: 관리자
        with gr.Tab("🛠️ 관리자") as admin_tab:
            admin_auth_status = gr.Markdown("관리자 권한이 필요합니다.")
            with gr.Row(visible=False) as admin_panel:
                with gr.Column():
                    gr.Markdown("### 📊 통계")
                    stats_display = gr.Markdown("")
                    refresh_stats_btn = gr.Button("통계 새로고침")
                with gr.Column():
                    gr.Markdown("### 👥 사용자 관리")
                    users_table = gr.Dataframe(headers=["ID", "이메일", "관리자", "활성", "일일제한", "가입일"], interactive=False)
                    refresh_users_btn = gr.Button("사용자 새로고침")
            with gr.Row(visible=False) as admin_actions:
                user_id_input = gr.Number(label="사용자 ID", precision=0)
                toggle_active_btn = gr.Button("활성화/비활성화")
                new_limit = gr.Number(label="새 일일 제한", value=10, precision=0)
                update_limit_btn = gr.Button("제한 변경")
                admin_action_status = gr.Markdown("")

    # ========== 이벤트 핸들러 ==========
    def do_login(email, password):
        token, msg = login_user(email, password)
        if token:
            user = validate_session(token)
            user_info = f"✅ **{user['email']}**" + (" (관리자)" if user['is_admin'] else "")
            return token, msg, user_info, gr.update(visible=True)
        return "", msg, "로그인하지 않음", gr.update(visible=False)
    
    def do_register(email, password, confirm):
        token, msg = register_user(email, password, confirm)
        if token:
            user = validate_session(token)
            return token, msg, f"✅ **{user['email']}**", gr.update(visible=True)
        return "", msg, "로그인하지 않음", gr.update(visible=False)
    
    def do_logout(token):
        logout_user(token)
        return "", "로그인하지 않음", gr.update(visible=False), "✅ 로그아웃됨"
    
    def load_account_info(token):
        user = validate_session(token)
        if not user:
            return "로그인이 필요합니다.", []
        info = f"### 👤 계정 정보\n- **이메일**: {user['email']}\n- **등급**: {'관리자' if user['is_admin'] else '일반'}\n- **일일 제한**: {user['daily_limit']}회\n- **오늘 사용량**: {get_daily_usage(user['id'])}회"
        history = get_user_generations(user["id"])
        history_data = [[h["original_prompt"][:50]+"...", h["created_at"], h["status"]] for h in history]
        return info, history_data
    
    def load_admin_panel(token):
        user = validate_session(token)
        if not user or not user["is_admin"]:
            return "❌ 관리자 권한이 필요합니다.", gr.update(visible=False), gr.update(visible=False), "", []
        stats = get_stats()
        stats_md = f"| 지표 | 값 |\n|---|---|\n| 총 사용자 | {stats['total_users']} |\n| 오늘 활성 | {stats['active_today']} |\n| 총 생성 | {stats['total_generations']} |\n| 오늘 생성 | {stats['generations_today']} |"
        users = get_all_users()
        users_data = [[u["id"], u["email"], "✅" if u["is_admin"] else "", "✅" if u["is_active"] else "❌", u["daily_limit"], u["created_at"]] for u in users]
        return "✅ 관리자 패널", gr.update(visible=True), gr.update(visible=True), stats_md, users_data
    
    def toggle_user_active(token, user_id):
        user = validate_session(token)
        if not user or not user["is_admin"]:
            return "❌ 권한 없음"
        conn = get_db()
        cursor = conn.cursor()
        cursor.execute("UPDATE users SET is_active = NOT is_active WHERE id = ?", (int(user_id),))
        conn.commit()
        conn.close()
        return "✅ 상태 변경됨"
    
    def change_user_limit(token, user_id, new_limit):
        user = validate_session(token)
        if not user or not user["is_admin"]:
            return "❌ 권한 없음"
        update_user_limit(int(user_id), int(new_limit))
        return f"✅ 제한 변경됨: {int(new_limit)}"

    # 이벤트 바인딩
    login_btn.click(do_login, [login_email, login_password], [session_token, login_status, current_user_display, logout_btn])
    reg_btn.click(do_register, [reg_email, reg_password, reg_confirm], [session_token, reg_status, current_user_display, logout_btn])
    logout_btn.click(do_logout, [session_token], [session_token, current_user_display, logout_btn, login_status])
    generate_btn.click(process_comparison_with_auth, [prompt_input, session_token, fireworks_key_input, fal_key_input, aspect_ratio, resolution], [status_text, enhanced_output, original_image_output, enhanced_image_output, usage_display])
    refresh_history_btn.click(load_account_info, [session_token], [account_info, history_display])
    refresh_stats_btn.click(load_admin_panel, [session_token], [admin_auth_status, admin_panel, admin_actions, stats_display, users_table])
    refresh_users_btn.click(load_admin_panel, [session_token], [admin_auth_status, admin_panel, admin_actions, stats_display, users_table])
    toggle_active_btn.click(toggle_user_active, [session_token, user_id_input], [admin_action_status])
    update_limit_btn.click(change_user_limit, [session_token, user_id_input, new_limit], [admin_action_status])
    
    # 탭 선택 시 자동 로드
    admin_tab.select(load_admin_panel, [session_token], [admin_auth_status, admin_panel, admin_actions, stats_display, users_table])
    account_tab.select(load_account_info, [session_token], [account_info, history_display])

if __name__ == "__main__":
    demo.launch()