"""
Versión usando Gemma-2B, un modelo pequeño de Google.
Renombra este archivo a app.py si prefieres usar este modelo.
"""

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
import gc
import os

# Función para cargar el prompt desde el archivo
def get_system_prompt():
    with open("prompt.txt", "r", encoding="utf-8") as f:
        return f.read().strip()

# Función para generar respuestas
def generate_response(user_message):
    try:
        if not user_message.strip():
            return "Por favor, escribe una pregunta para que pueda ayudarte."
        
        # Obtener el prompt del sistema
        system_prompt = get_system_prompt()
        
        # Crear el prompt completo para Gemma
        prompt = f"<start_of_turn>system\n{system_prompt}<end_of_turn>\n<start_of_turn>user\n{user_message}<end_of_turn>\n<start_of_turn>model\n"
        
        # Generar respuesta usando el pipeline
        response = generator(
            prompt,
            max_new_tokens=256,  # Reducido para mayor velocidad
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            num_return_sequences=1
        )[0]["generated_text"]
        
        # Extraer solo la respuesta del modelo (después del último <start_of_turn>model)
        assistant_response = response.split("<start_of_turn>model\n")[-1].split("<end_of_turn>")[0].strip()
        
        # Forzar recolección de basura
        gc.collect()
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        
        return assistant_response
    
    except Exception as e:
        print(f"Error: {str(e)}")
        return f"Lo siento, ha ocurrido un error: {str(e)}"

# Cargar el modelo y crear el pipeline
print("Cargando modelo Gemma-2B...")

# Usar Gemma-2B, un modelo pequeño (2B parámetros) pero potente
model_name = "google/gemma-2b-it"

# Configuración para reducir el uso de memoria y aumentar velocidad
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
    # Cargamos en 8-bit para reducir memoria y aumentar velocidad
    load_in_8bit=True
)

# Crear el pipeline de generación de texto
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer
)

print("Modelo cargado correctamente!")

# Crear la interfaz de Gradio
demo = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(
        placeholder="Escribe tu pregunta sobre emprendimiento aquí...",
        label="Tu pregunta"
    ),
    outputs=gr.Textbox(label="Respuesta de BITER"),
    title="BITER - Tu Mentor en Tiempo Real para Decisiones de Negocio",
    description="BITER es un asistente de IA que responde dudas de emprendedores como si fuera un CEO experimentado.",
    examples=[
        ["¿Cómo puedo validar mi idea de negocio con poco presupuesto?"],
        ["¿Cuál es la mejor estrategia para conseguir mis primeros clientes?"],
        ["¿Debería invertir en publicidad en redes sociales o en SEO?"]
    ],
    allow_flagging="never"
)

# Lanzar la aplicación con configuración para ahorrar memoria
if __name__ == "__main__":
    # Configurar menos workers para ahorrar memoria y aumentar velocidad
    demo.queue(max_size=1).launch(share=False, debug=False)