"""
Versión alternativa usando un modelo aún más pequeño y optimizado.
Renombra este archivo a app.py si el modelo Phi-2 sigue siendo demasiado grande.
"""

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
        prompt = f"{system_prompt}\n\nPregunta: {user_message}\n\nRespuesta:"
        
        # Generar respuesta usando el pipeline
        response = generator(
            prompt,
            max_new_tokens=256,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            num_return_sequences=1
        )[0]["generated_text"]
        
        # Extraer solo la respuesta (después de "Respuesta:")
        assistant_response = response.split("Respuesta:")[-1].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 ultra ligero...")

# Usar un modelo muy pequeño con soporte para español
model_name = "PlanTL-GOB-ES/gpt2-base-bne"  # Modelo de ~125M parámetros, extremadamente ligero

# Configuración para reducir el uso de memoria
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
    device_map="auto"
)

# 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
    demo.queue(max_size=1).launch(share=False, debug=False)