# orionai_app.py
import gradio as gr
import numpy as np
import random
import torch
from diffusers import DiffusionPipeline
import traceback
import time
import psycopg2
from psycopg2.extras import RealDictCursor
import hashlib
import secrets
from datetime import datetime, timedelta, timezone
import os
from io import BytesIO
import base64
from PIL import Image

print("🚀 Démarrage de l'application OrionAI...")

# ===================================================================
# CONFIGURATION DATABASE NEON
# ===================================================================
DATABASE_URL = "postgresql://neondb_owner:npg_AtKIqbxoN7s1@ep-rough-fog-adc0kzr6-pooler.c-2.us-east-1.aws.neon.tech/neondb?sslmode=require"

def get_db_connection():
    """Connexion au database PostgreSQL Neon"""
    try:
        conn = psycopg2.connect(DATABASE_URL)
        return conn
    except Exception as e:
        print(f"❌ Errore connessione DB: {e}")
        return None

def init_database():
    """Initialise les tables du database"""
    conn = get_db_connection()
    if not conn:
        print("❌ Impossibile connettersi al database!")
        return False

    try:
        cur = conn.cursor()
        # users
        cur.execute("""
            CREATE TABLE IF NOT EXISTS users (
                id SERIAL PRIMARY KEY,
                username VARCHAR(50) UNIQUE NOT NULL,
                email VARCHAR(100) UNIQUE NOT NULL,
                password_hash VARCHAR(255) NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        # image_generate
        cur.execute("""
            CREATE TABLE IF NOT EXISTS image_generate (
                id SERIAL PRIMARY KEY,
                user_id INTEGER REFERENCES users(id) ON DELETE CASCADE,
                prompt TEXT NOT NULL,
                image_base64 TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        # user_reset_password
        # NOTE: expires_at is NOT NULL but has default = current_timestamp + interval '1 hour'
        cur.execute("""
            CREATE TABLE IF NOT EXISTS user_reset_password (
                id SERIAL PRIMARY KEY,
                user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE,
                token VARCHAR(255),
                reset_token VARCHAR(255),
                expires_at TIMESTAMP NOT NULL DEFAULT (CURRENT_TIMESTAMP + INTERVAL '1 hour'),
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                used BOOLEAN DEFAULT FALSE
            )
        """)

        conn.commit()
        cur.close()
        conn.close()
        print("✅ Database inizializzato con successo!")
        return True
    except Exception as e:
        print(f"❌ Errore inizializzazione DB: {e}")
        print(traceback.format_exc())
        try:
            conn.rollback()
        except:
            pass
        if conn:
            conn.close()
        return False

# initialize DB at startup
init_database()

# ===================================================================
# UTILITAIRES
# ===================================================================
def hash_password(password: str) -> str:
    return hashlib.sha256(password.encode('utf-8')).hexdigest()

def now_utc():
    return datetime.now(timezone.utc)

# convert PIL.Image -> base64 (PNG)
def pil_to_base64(img: Image.Image) -> str:
    buffer = BytesIO()
    img.save(buffer, format="PNG")
    b = buffer.getvalue()
    return base64.b64encode(b).decode('utf-8')  # store without data: prefix

# convert base64 (stored) -> PIL.Image
def base64_to_pil(b64: str) -> Image.Image:
    b = base64.b64decode(b64)
    return Image.open(BytesIO(b)).convert("RGBA")

# ===================================================================
# AUTHENTICATION / USER FUNCTIONS
# ===================================================================
def register_user(username, email, password):
    """Register new user"""
    print(f"🔵 Tentativo registrazione: {username}, {email}")
    if not username or not email or not password:
        return False, "❌ Tutti i campi sono obbligatori!"
    if len(password) < 6:
        return False, "❌ La password deve essere almeno 6 caratteri!"

    conn = get_db_connection()
    if not conn:
        return False, "❌ Errore connessione database!"
    try:
        cur = conn.cursor()
        password_hash = hash_password(password)
        cur.execute(
            "INSERT INTO users (username, email, password_hash) VALUES (%s, %s, %s) RETURNING id",
            (username, email, password_hash)
        )
        user_id = cur.fetchone()[0]
        conn.commit()
        cur.close()
        conn.close()
        print(f"✅ Utente registrato con ID: {user_id}")
        return True, f"✅ Registrazione completata!\n\nUsername: **{username}**\n\nOra puoi accedere."
    except psycopg2.IntegrityError as e:
        print(f"❌ IntegrityError: {e}")
        try:
            conn.rollback()
        except:
            pass
        if conn:
            conn.close()
        return False, "❌ Username o email già esistenti!"
    except Exception as e:
        print(f"❌ Errore registrazione: {e}")
        print(traceback.format_exc())
        if conn:
            conn.close()
        return False, f"❌ Errore: {str(e)}"

def login_user(username, password):
    """Login"""
    print(f"🔵 Tentativo login: {username}")
    if not username or not password:
        return None, "❌ Inserisci username e password!"
    conn = get_db_connection()
    if not conn:
        return None, "❌ Errore connessione database!"
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        password_hash = hash_password(password)
        cur.execute("SELECT id, username, password_hash FROM users WHERE username = %s", (username,))
        user = cur.fetchone()
        if not user:
            cur.close()
            conn.close()
            return None, "❌ Username non esistente!"
        if user['password_hash'] == password_hash:
            user_dict = {'id': user['id'], 'username': user['username']}
            cur.close()
            conn.close()
            print(f"✅ Login riuscito per: {username}")
            return user_dict, "✅ Login effettuato!"
        else:
            cur.close()
            conn.close()
            return None, "❌ Password errata!"
    except Exception as e:
        print(f"❌ Errore login: {e}")
        print(traceback.format_exc())
        if conn:
            conn.close()
        return None, f"❌ Errore: {str(e)}"

# Password reset: request
def request_password_reset(email, expires_hours=1):
    """Generate a reset token and store with expires_at"""
    conn = get_db_connection()
    if not conn:
        return False, "❌ Errore connessione database!"
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute("SELECT id FROM users WHERE email = %s", (email,))
        user = cur.fetchone()
        if not user:
            cur.close()
            conn.close()
            return False, "❌ Email non trovata!"

        reset_token = secrets.token_urlsafe(32)
        expires_at = datetime.now() + timedelta(hours=expires_hours)  # naive local; DB default is UTC-ish, but acceptable
        # Insert token + explicit expires_at (to avoid NOT NULL violation)
        cur.execute(
            "INSERT INTO user_reset_password (user_id, token, reset_token, expires_at) VALUES (%s, %s, %s, %s)",
            (user['id'], reset_token, reset_token, expires_at)
        )
        conn.commit()
        cur.close()
        conn.close()
        print(f"✅ Reset token generated for user_id {user['id']} (expires at {expires_at})")
        # In real app, you'd email the token; here we return it for debug/demo
        return True, f"✅ Token di reset generato!\n\n**Token:** {reset_token}\n\n⚠️ Copia questo token per resettare la password!"
    except Exception as e:
        print(f"❌ Errore reset request: {e}")
        print(traceback.format_exc())
        try:
            conn.rollback()
        except:
            pass
        if conn:
            conn.close()
        return False, f"❌ Errore: {str(e)}"

def reset_password(reset_token, new_password):
    """Reset password using token; checks expiry and used flag"""
    if len(new_password) < 6:
        return False, "❌ La password deve essere almeno 6 caratteri!"
    conn = get_db_connection()
    if not conn:
        return False, "❌ Errore connessione database!"
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        # fetch record, check used & expiry
        cur.execute(
            "SELECT id, user_id, expires_at, used FROM user_reset_password WHERE (token = %s OR reset_token = %s)",
            (reset_token, reset_token)
        )
        rec = cur.fetchone()
        if not rec:
            cur.close()
            conn.close()
            return False, "❌ Token non valido!"
        if rec['used']:
            cur.close()
            conn.close()
            return False, "❌ Token già utilizzato!"
        # expires_at can be a datetime
        if rec['expires_at'] and datetime.now() > rec['expires_at']:
            cur.close()
            conn.close()
            return False, "❌ Token scaduto!"
        # update password
        password_hash = hash_password(new_password)
        cur.execute("UPDATE users SET password_hash = %s WHERE id = %s", (password_hash, rec['user_id']))
        # mark token used
        cur.execute("UPDATE user_reset_password SET used = TRUE WHERE id = %s", (rec['id'],))
        conn.commit()
        cur.close()
        conn.close()
        return True, "✅ Password resettata con successo! Ora puoi accedere."
    except Exception as e:
        print(f"❌ Errore reset password: {e}")
        print(traceback.format_exc())
        try:
            conn.rollback()
        except:
            pass
        if conn:
            conn.close()
        return False, f"❌ Errore: {str(e)}"

# ===================================================================
# GESTIONE IMMAGINI (DB store: base64)
# ===================================================================
def save_image_to_db(user_id, prompt, pil_image: Image.Image):
    """Save image in DB as base64 (no data: prefix)"""
    try:
        img_b64 = pil_to_base64(pil_image)
        conn = get_db_connection()
        if not conn:
            return False
        cur = conn.cursor()
        cur.execute(
            "INSERT INTO image_generate (user_id, prompt, image_base64) VALUES (%s, %s, %s)",
            (user_id, prompt, img_b64)
        )
        conn.commit()
        cur.close()
        conn.close()
        print(f"✅ Immagine salvata per user_id: {user_id}")
        return True
    except Exception as e:
        print(f"❌ Errore salvataggio immagine: {e}")
        print(traceback.format_exc())
        try:
            conn.rollback()
        except:
            pass
        if conn:
            conn.close()
        return False

def get_user_gallery(user_id, limit=50):
    """Return list of dicts with id, prompt, image_base64, created_at"""
    conn = get_db_connection()
    if not conn:
        return []
    try:
        cur = conn.cursor(cursor_factory=RealDictCursor)
        cur.execute(
            "SELECT id, prompt, image_base64, created_at FROM image_generate WHERE user_id = %s ORDER BY created_at DESC LIMIT %s",
            (user_id, limit)
        )
        rows = cur.fetchall()
        cur.close()
        conn.close()
        return [dict(r) for r in rows]
    except Exception as e:
        print(f"❌ Errore recupero galleria: {e}")
        print(traceback.format_exc())
        if conn:
            conn.close()
        return []

# ===================================================================
# LOAD DIFFUSION MODEL
# ===================================================================
device = "cpu"
model_repo_id = "Muyumba/orion_ai"

print(f"💻 Device: CPU")

pipe = None
try:
    print(f"📦 Caricamento modello {model_repo_id}...")
    pipe = DiffusionPipeline.from_pretrained(
        model_repo_id,
        torch_dtype=torch.float32,
        trust_remote_code=True,
        use_safetensors=True,
        low_cpu_mem_usage=True
    )
    pipe = pipe.to("cpu")
    if hasattr(pipe, "enable_attention_slicing"):
        pipe.enable_attention_slicing(1)
    print("✅ Modello caricato!")
except Exception as e:
    print(f"❌ Errore caricamento modello: {e}")
    print(traceback.format_exc())

MAX_SEED = np.iinfo(np.int32).max

# ===================================================================
# IMAGE GENERATION
# ===================================================================
def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, steps, user_id):
    if pipe is None:
        return None, "❌ Modello non caricato!"
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    try:
        generator = torch.Generator(device="cpu").manual_seed(seed)
        start = time.time()
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            num_inference_steps=steps,
            generator=generator
        )
        elapsed = time.time() - start
        image = result.images[0].convert("RGBA")  # PIL Image
        save_status = ""
        if user_id:
            if save_image_to_db(user_id, prompt, image):
                save_status = " | ✅ Salvata in galleria"
            else:
                save_status = " | ⚠️ Errore salvataggio in galleria"
        return image, f"✅ Immagine generata in {elapsed:.1f}s! (Seed: {seed}){save_status}"
    except Exception as e:
        print(f"❌ Errore generazione: {e}")
        print(traceback.format_exc())
        return None, f"❌ Errore: {str(e)}"

# ===================================================================
# CSS & GRADIO UI
# ===================================================================
css = """
.container {max-width: 900px; margin: auto; padding: 20px;}
.login-box {max-width: 450px; margin: 50px auto; padding: 40px; border-radius: 15px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); box-shadow: 0 10px 25px rgba(0,0,0,0.2);}
.gr-button-primary {background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; font-weight: bold;}
"""

with gr.Blocks(css=css, title="OrionAI") as demo:
    user_state = gr.State(None)

    # Login screen
    with gr.Column(visible=True) as login_screen:
        with gr.Column(elem_classes="login-box"):
            gr.Markdown("# 🔐 OrionAI - Login\n### Generatore di Immagini IA")
            login_username = gr.Textbox(label="👤 Username", placeholder="Inserisci username")
            login_password = gr.Textbox(label="🔑 Password", type="password", placeholder="Inserisci password")
            with gr.Row():
                login_btn = gr.Button("🚀 Accedi", variant="primary", size="lg")
                register_link = gr.Button("📝 Registrati", variant="secondary", size="lg")
            forgot_password_link = gr.Button("🔄 Password dimenticata?", variant="secondary", size="sm")
            login_msg = gr.Markdown("")

    # Register screen
    with gr.Column(visible=False) as register_screen:
        with gr.Column(elem_classes="login-box"):
            gr.Markdown("# 📝 Registrazione\n### Crea il tuo account")
            reg_username = gr.Textbox(label="👤 Username", placeholder="Scegli un username")
            reg_email = gr.Textbox(label="📧 Email", placeholder="tuaemail@example.com")
            reg_password = gr.Textbox(label="🔑 Password", type="password", placeholder="Min. 6 caratteri")
            reg_password_confirm = gr.Textbox(label="🔑 Conferma Password", type="password", placeholder="Ripeti la password")
            with gr.Row():
                register_btn = gr.Button("✅ Registrati", variant="primary", size="lg")
                back_to_login = gr.Button("← Torna al Login", variant="secondary", size="lg")
            register_msg = gr.Markdown("")

    # Reset screen
    with gr.Column(visible=False) as reset_screen:
        with gr.Column(elem_classes="login-box"):
            gr.Markdown("# 🔄 Reset Password\n### Recupera il tuo account")
            with gr.Tab("1️⃣ Richiedi Reset"):
                reset_email = gr.Textbox(label="📧 Email", placeholder="tuaemail@example.com")
                request_reset_btn = gr.Button("📧 Invia Token Reset", variant="primary")
                reset_msg1 = gr.Markdown("")
            with gr.Tab("2️⃣ Nuova Password"):
                reset_token = gr.Textbox(label="🎫 Token Reset", placeholder="Incolla il token ricevuto")
                new_password = gr.Textbox(label="🔑 Nuova Password", type="password", placeholder="Min. 6 caratteri")
                reset_password_btn = gr.Button("✅ Resetta Password", variant="primary")
                reset_msg2 = gr.Markdown("")
            back_to_login2 = gr.Button("← Torna al Login", variant="secondary")

    # Main app
    with gr.Column(visible=False) as main_app:
        gr.Markdown("# 🌟 OrionAI - Generatore di Immagini IA")
        user_info = gr.Markdown("")
        with gr.Tabs():
            with gr.Tab("🎨 Genera Immagine"):
                with gr.Row():
                    prompt = gr.Textbox(label="📝 Descrivi l'immagine", placeholder="A beautiful landscape...", lines=3)
                    generate_btn = gr.Button("🎨 Genera", variant="primary", scale=0)
                result_image = gr.Image(label="🖼️ Immagine Generata", height=400)
                result_info = gr.Textbox(label="ℹ️ Info", interactive=False)
                with gr.Accordion("⚙️ Parametri", open=False):
                    negative_prompt = gr.Textbox(label="❌ Elementi da evitare", value="blurry, bad quality, distorted")
                    with gr.Row():
                        seed = gr.Slider(0, MAX_SEED, 0, label="🎲 Seed")
                        randomize = gr.Checkbox(True, label="🔄 Random")
                    with gr.Row():
                        width = gr.Slider(256, 768, 512, 64, label="📐 Larghezza")
                        height = gr.Slider(256, 768, 512, 64, label="📏 Altezza")
                    with gr.Row():
                        guidance = gr.Slider(1.0, 15.0, 7.5, 0.5, label="🎯 Guidance")
                        steps = gr.Slider(10, 30, 20, 5, label="🔧 Steps")
                gr.Markdown("**💡 Suggerimento:** Per risultati veloci su CPU, usa 512x512 con 20 steps")

            with gr.Tab("📚 La Mia Galleria"):
                gr.Markdown("**💡 Info:** La galleria si aggiorna automaticamente dopo ogni generazione!")
                refresh_gallery_btn = gr.Button("🔄 Aggiorna Manualmente", variant="secondary")
                gallery = gr.Gallery(label="Le Tue Immagini", columns=3, height=600)

        logout_btn = gr.Button("🚪 Logout", variant="secondary")

    # ========== EVENTS ==========
    def handle_login(username, password):
        user, msg = login_user(username, password)
        if user:
            # load gallery images as PIL images
            images = get_user_gallery(user['id'])
            gallery_images = []
            for img in images:
                try:
                    pil = base64_to_pil(img['image_base64'])
                    caption = img['prompt'][:50] if img['prompt'] else ""
                    gallery_images.append((pil, caption))
                except Exception as e:
                    print(f"❌ Errore decode image for gallery: {e}")
            return {
                login_screen: gr.update(visible=False),
                main_app: gr.update(visible=True),
                user_state: user,
                user_info: gr.update(value=f"👤 **Benvenuto, {user['username']}!**"),
                login_msg: gr.update(value=""),
                login_username: gr.update(value=""),
                login_password: gr.update(value=""),
                gallery: gallery_images
            }
        return {
            login_screen: gr.update(visible=True),
            main_app: gr.update(visible=False),
            user_state: None,
            login_msg: gr.update(value=msg),
            gallery: []
        }

    login_btn.click(
        handle_login,
        [login_username, login_password],
        [login_screen, main_app, user_state, user_info, login_msg, login_username, login_password, gallery]
    )

    def show_register():
        return {login_screen: gr.update(visible=False), register_screen: gr.update(visible=True), register_msg: gr.update(value="")}

    register_link.click(show_register, outputs=[login_screen, register_screen, register_msg])

    def handle_register(username, email, password, password_confirm):
        if password != password_confirm:
            return {register_msg: gr.update(value="❌ Le password non corrispondono!")}
        success, msg = register_user(username, email, password)
        if success:
            return {
                register_screen: gr.update(visible=False),
                login_screen: gr.update(visible=True),
                register_msg: gr.update(value=""),
                login_msg: gr.update(value=msg),
                reg_username: gr.update(value=""),
                reg_email: gr.update(value=""),
                reg_password: gr.update(value=""),
                reg_password_confirm: gr.update(value="")
            }
        return {register_msg: gr.update(value=msg)}

    register_btn.click(
        handle_register,
        [reg_username, reg_email, reg_password, reg_password_confirm],
        [register_screen, login_screen, register_msg, login_msg, reg_username, reg_email, reg_password, reg_password_confirm]
    )

    back_to_login.click(
        lambda: {register_screen: gr.update(visible=False), login_screen: gr.update(visible=True)},
        outputs=[register_screen, login_screen]
    )

    forgot_password_link.click(
        lambda: {login_screen: gr.update(visible=False), reset_screen: gr.update(visible=True)},
        outputs=[login_screen, reset_screen]
    )

    request_reset_btn.click(
        lambda email: request_password_reset(email)[1],
        reset_email,
        reset_msg1
    )

    reset_password_btn.click(
        lambda token, pwd: reset_password(token, pwd)[1],
        [reset_token, new_password],
        reset_msg2
    )

    back_to_login2.click(
        lambda: {reset_screen: gr.update(visible=False), login_screen: gr.update(visible=True)},
        outputs=[reset_screen, login_screen]
    )

    def handle_generate(prompt_text, neg, seed_val, rand, w, h, guid, st, user):
        if not user:
            return None, "❌ Errore: utente non autenticato!", []
        image, info = generate_image(prompt_text, neg, seed_val, rand, w, h, guid, st, user['id'])
        gallery_images = []
        if user:
            images = get_user_gallery(user['id'])
            for img in images:
                try:
                    pil = base64_to_pil(img['image_base64'])
                    caption = img['prompt'][:50] if img['prompt'] else ""
                    gallery_images.append((pil, caption))
                except Exception as e:
                    print(f"❌ Errore decode image for gallery after generation: {e}")
        return image, info, gallery_images

    generate_btn.click(
        handle_generate,
        [prompt, negative_prompt, seed, randomize, width, height, guidance, steps, user_state],
        [result_image, result_info, gallery]
    )

    def load_gallery(user):
        if not user:
            return []
        images = get_user_gallery(user['id'])
        gallery_list = []
        for img in images:
            try:
                pil = base64_to_pil(img['image_base64'])
                caption = img['prompt'][:50] if img['prompt'] else ""
                gallery_list.append((pil, caption))
            except Exception as e:
                print(f"❌ Errore decode gallery image on refresh: {e}")
        return gallery_list

    refresh_gallery_btn.click(load_gallery, user_state, gallery)

    def handle_logout():
        return {
            main_app: gr.update(visible=False),
            login_screen: gr.update(visible=True),
            user_state: None,
            login_username: gr.update(value=""),
            login_password: gr.update(value="")
        }

    logout_btn.click(
        handle_logout,
        outputs=[main_app, login_screen, user_state, login_username, login_password]
    )

# ===================================================================
# LAUNCH
# ===================================================================
if __name__ == "__main__":
    print("🌐 Avvio OrionAI...")
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)