import gradio as gr
import sqlite3
import secrets
import string
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- CONFIGURATION DU MODÈLE ---
MODEL_ID = "Finisha-F-scratch/Charlotte-amity"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, 
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

# --- GESTION DE LA BASE DE DONNÉES (SQLITE) ---
DB_NAME = "charlotte_api.db"

def init_db():
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    # Table pour stocker les clés et le nombre de requêtes
    c.execute('''CREATE TABLE IF NOT EXISTS api_keys 
                 (key TEXT PRIMARY KEY, requests_count INTEGER)''')
    conn.commit()
    conn.close()

init_db()

def create_key():
    """Génère une clé tn-xxxx0000 et l'enregistre"""
    letters = ''.join(secrets.choice(string.ascii_lowercase) for _ in range(4))
    digits = ''.join(secrets.choice(string.digits) for _ in range(4))
    new_key = f"tn-{letters}{digits}"
    
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("INSERT INTO api_keys (key, requests_count) VALUES (?, ?)", (new_key, 0))
        conn.commit()
        return new_key
    except:
        return "Erreur lors de la génération. Réessayez."
    finally:
        conn.close()

def validate_and_log(key):
    """Vérifie la clé et incrémente le compteur"""
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    c.execute("SELECT requests_count FROM api_keys WHERE key=?", (key,))
    result = c.fetchone()
    if result is not None:
        c.execute("UPDATE api_keys SET requests_count = requests_count + 1 WHERE key=?", (key,))
        conn.commit()
        conn.close()
        return True
    conn.close()
    return False

# --- FONCTION D'INFÉRENCE (L'API) ---
def charlotte_api(api_key, message):
    """
    Point d'entrée principal pour les appels externes.
    """
    if not validate_and_log(api_key):
        return "❌ Erreur : Clé API 'tn-' invalide ou non trouvée."

    # Limitation stricte au contexte de 128 tokens
    inputs = tokenizer(message, return_tensors="pt", truncation=True, max_length=128).to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=100,
            do_sample=True,
            temperature=0.8,
            pad_token_id=tokenizer.eos_token_id
        )
    
    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # On renvoie uniquement la réponse générée
    response = full_text[len(message):].strip()
    return response

# --- INTERFACE GRADIO ---
with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    gr.Markdown("# 🌸 Charlotte-Amity : Système API Personnel")
    
    with gr.Tab("🔑 Gérer mes clés"):
        gr.Markdown("Cliquez ci-dessous pour générer une clé d'accès unique.")
        key_output = gr.Textbox(label="Votre nouvelle clé API", interactive=False)
        generate_btn = gr.Button("Générer ma clé tn-")
        generate_btn.click(create_key, outputs=key_output)
        gr.Info("Note: Gardez précieusement votre clé pour vos appels externes.")

    with gr.Tab("💬 Test Rapide"):
        api_input = gr.Textbox(label="Clé API (tn-...)")
        msg_input = gr.Textbox(label="Votre message")
        text_output = gr.Textbox(label="Réponse de Charlotte")
        run_btn = gr.Button("Envoyer")
        run_btn.click(charlotte_api, inputs=[api_input, msg_input], outputs=text_output)

    # L'EXPOSITION DE L'API POUR L'EXTÉRIEUR
    # Cette ligne permet d'appeler la fonction via l'URL du Space
    api_interface = gr.Interface(
        fn=charlotte_api,
        inputs=[gr.Textbox(label="api_key"), gr.Textbox(label="message")],
        outputs="text",
        api_name="chat" # C'est le nom à utiliser dans le client Python
    )

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