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

# Variables globales para el modelo y tokenizador
model = None
tokenizer = None

# 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 cargar el modelo (se ejecutará solo cuando sea necesario)
def load_model_if_needed():
    global model, tokenizer
    
    if model is None:
        print("Cargando modelo Mistral-7B-Instruct-v0.1...")
        
        # Cargar el tokenizador
        model_name = "mistralai/Mistral-7B-Instruct-v0.1"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        # Configuración para CPU con optimizaciones de memoria
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float32,  # Usar float32 para CPU
            low_cpu_mem_usage=True,     # Optimización para CPU con poca memoria
        )
        
        # Mover el modelo a CPU explícitamente
        model = model.to("cpu")
        
        # Forzar recolección de basura para liberar memoria
        gc.collect()
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        
        print("Modelo cargado correctamente en CPU!")

# Función principal que procesa las preguntas del usuario
def generate_response(user_message):
    try:
        # Cargar el modelo si aún no está cargado
        load_model_if_needed()
        
        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()
        
        # Formato de prompt para Mistral-7B-Instruct-v0.1
        prompt = f"""<s>[INST] {system_prompt}

Pregunta del usuario: {user_message} [/INST]</s>"""
        
        # Tokenizar el prompt
        inputs = tokenizer(prompt, return_tensors="pt")
        
        # Configuración de generación optimizada para CPU
        generation_config = {
            "max_new_tokens": 512,      # Limitar tokens para ahorrar memoria
            "temperature": 0.7,
            "top_p": 0.9,
            "do_sample": True,
            "pad_token_id": tokenizer.eos_token_id,
            "num_return_sequences": 1
        }
        
        # Generar respuesta
        with torch.no_grad():
            outputs = model.generate(**inputs, **generation_config)
        
        # Decodificar la respuesta
        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extraer solo la respuesta del asistente (después del prompt)
        assistant_response = full_response.replace(prompt.replace("<s>", "").replace("</s>", ""), "").strip()
        
        # Forzar recolección de basura para liberar memoria
        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)}"

# 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)