from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import pipeline
import os

app = Flask(__name__)
CORS(app) # Habilita CORS para permitir peticiones desde tu frontend

# Nombre del modelo de Hugging Face
MODEL_ID = "sarvamai/sarvam-m"

# Variable para almacenar el pipeline de la IA
pipe = None

# Función para cargar el modelo
def load_model():
    global pipe
    try:
        print(f"Cargando modelo '{MODEL_ID}'. Esto puede tardar unos minutos...")
        # El token de Hugging Face se obtiene automáticamente de la variable de entorno HF_TOKEN
        # que configuras como un "Secret" en Hugging Face Spaces.
        pipe = pipeline("text-generation", model=MODEL_ID, device=-1) # device=-1 usa CPU
        print(f"Modelo '{MODEL_ID}' cargado con éxito.")
    except Exception as e:
        print(f"Error al cargar el modelo '{MODEL_ID}': {e}")
        pipe = None
        print("El servidor no podrá responder a las solicitudes de IA.")

# Cargar el modelo inmediatamente después de crear la aplicación Flask
# Esto se asegura de que el modelo se cargue al inicio del servidor
with app.app_context():
    load_model()

@app.route('/chat', methods=['POST'])
def chat():
    if pipe is None:
        return jsonify({"error": "El modelo de IA no está cargado. Por favor, revisa el log del servidor."}), 503

    data = request.json
    messages = data.get('messages') # El frontend envía el historial completo

    if not messages:
        return jsonify({"error": "No se proporcionaron mensajes."}), 400

    try:
        prompt = ""
        for msg in messages:
            if msg['role'] == 'user':
                prompt += f"Usuario: {msg['content']}\n"
            elif msg['role'] == 'assistant':
                prompt += f"Asistente: {msg['content']}\n"
        prompt += "Asistente:"

        response = pipe(prompt, max_new_tokens=250, temperature=0.7, do_sample=True, clean_up_tokenization_spaces=True)

        generated_text = response[0]['generated_text']
        ai_response_content = generated_text.split("Asistente:", 1)[-1].strip()

        if not ai_response_content:
            ai_response_content = "Lo siento, no pude generar una respuesta clara. ¿Puedes intentarlo de nuevo?"

        return jsonify({"response": ai_response_content})

    except Exception as e:
        print(f"Error en la inferencia de la IA: {e}")
        return jsonify({"error": f"Error interno del servidor: {str(e)}"}), 500

if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    app.run(debug=False, host='0.0.0.0', port=port)