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

model = None
tokenizer = None

def get_system_prompt():
    with open("prompt.txt", "r", encoding="utf-8") as f:
        return f.read().strip()

def load_model_if_needed():
    global model, tokenizer

    if model is None:
        print("🔁 Cargando modelo Falcon-7B-Instruct...")

        model_name = "tiiuae/falcon-7b-instruct"

        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float32,
            trust_remote_code=True,
            low_cpu_mem_usage=True
        )

        model = model.to("cpu")

        gc.collect()
        torch.cuda.empty_cache() if torch.cuda.is_available() else None

        print("✅ Modelo Falcon-7B cargado en CPU")

def generate_response(user_message):
    try:
        load_model_if_needed()

        if not user_message.strip():
            return "Por favor, escribe una pregunta para que pueda ayudarte."

        system_prompt = get_system_prompt()

        prompt = f"{system_prompt}\n\nUsuario: {user_message}\nBITER:"

        inputs = tokenizer(prompt, return_tensors="pt")

        generation_config = {
            "max_new_tokens": 400,
            "temperature": 0.7,
            "top_p": 0.9,
            "do_sample": True,
            "pad_token_id": tokenizer.eos_token_id,
            "num_return_sequences": 1
        }

        with torch.no_grad():
            outputs = model.generate(**inputs, **generation_config)

        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        answer = response.replace(prompt, "").strip()

        gc.collect()
        torch.cuda.empty_cache() if torch.cuda.is_available() else None

        return answer

    except Exception as e:
        print(f"Error: {str(e)}")
        return f"❌ Lo siento, ha ocurrido un error: {str(e)}"

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

if __name__ == "__main__":
    demo.queue(max_size=1).launch(share=False, debug=False)