from flask import Flask, request, jsonify
from flask_cors import CORS
from llama_cpp import Llama
import os

# Configuración del servidor
app = Flask(__name__)
CORS(app) # Permite que tu página web (PHP/JS) hable con este script

# --- CONFIGURACIÓN DEL MODELO ---
# Ajusta 'n_ctx' según tu RAM. 2048 es seguro para 8GB RAM. 
# Si tienes más, puedes subirlo a 4096 o 8192.
MODEL_PATH = "./models/gemma-3-4b-it-q4_k_m.gguf"

print(f"🔄 Cargando Gemma 3 desde {MODEL_PATH}...")
print("Esto puede tardar unos segundos...")

try:
    # n_gpu_layers=-1 intenta usar GPU si hay drivers instalados. Si no, usa CPU automáticamente.
    llm = Llama(
        model_path=MODEL_PATH,
        n_ctx=2048,          # Ventana de contexto (cuánto recuerda)
        n_threads=4,         # Núcleos de CPU a usar (ajusta según tu servidor)
        n_gpu_layers=-1,     # Aceleración GPU (opcional)
        verbose=False        # Mantiene la consola limpia
    )
    print("✅ ¡Modelo cargado correctamente!")
except Exception as e:
    print(f"❌ Error fatal cargando el modelo: {e}")
    exit(1)

# --- DEFINICIÓN DEL SISTEMA (PERSONALIDAD) ---
SYSTEM_PROMPT = {
    "role": "system",
    "content": "Eres un asistente educativo útil, amable y conciso de una escuela. Responde siempre en español de forma clara. Si no sabes algo, admítelo. No des información peligrosa o inapropiada."
}

@app.route('/chat', methods=['POST'])
def chat():
    data = request.json
    user_input = data.get('message', '')

    if not user_input:
        return jsonify({"error": "No enviaste ningún mensaje"}), 400

    # Preparamos el historial para Gemma
    # Si quieres memoria conversacional, aquí deberías recibir el historial completo del frontend
    messages = [
        SYSTEM_PROMPT,
        {"role": "user", "content": user_input}
    ]

    try:
        # Generación de respuesta
        output = llm.create_chat_completion(
            messages=messages,
            max_tokens=512,       # Límite de largo de respuesta (para no saturar el server)
            temperature=0.7,      # Creatividad (0.7 es equilibrado)
            top_p=0.9,
            stop=["<end_of_turn>"] # Importante para que no hable solo
        )

        # Extraer solo el texto de la respuesta
        ai_response = output['choices'][0]['message']['content']
        
        return jsonify({
            "response": ai_response,
            "status": "success"
        })

    except Exception as e:
        print(f"Error generando respuesta: {e}")
        return jsonify({"error": "Error interno del servidor IA"}), 500

if __name__ == '__main__':
    # '0.0.0.0' hace que sea visible en la red local, no solo en localhost
    app.run(host='0.0.0.0', port=5000)