from flask import Flask, request, jsonify
from typing import Dict, Optional
import uuid
import os
from app import initialize_llm, setup_database_connection, create_agent, gr

app = Flask(__name__)

# Almacenamiento en memoria de los mensajes
message_store: Dict[str, str] = {}

@app.route('/user_message', methods=['POST'])
def handle_user_message():
    try:
        data = request.get_json()
        if not data or 'message' not in data:
            return jsonify({'error': 'Se requiere el campo message'}), 400
            
        user_message = data['message']
        
        # Generar un ID único para este mensaje
        message_id = str(uuid.uuid4())
        
        # Almacenar el mensaje
        message_store[message_id] = user_message
        
        return jsonify({
            'message_id': message_id,
            'status': 'success'
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/ask', methods=['POST'])
def handle_ask():
    try:
        data = request.get_json()
        if not data or 'message_id' not in data:
            return jsonify({'error': 'Se requiere el campo message_id'}), 400
            
        message_id = data['message_id']
        
        # Recuperar el mensaje almacenado
        if message_id not in message_store:
            return jsonify({'error': 'ID de mensaje no encontrado'}), 404
            
        user_message = message_store[message_id]
        
        # Inicializar componentes necesarios
        llm, llm_error = initialize_llm()
        if llm_error:
            return jsonify({'error': f'Error al inicializar LLM: {llm_error}'}), 500
            
        db_connection, db_error = setup_database_connection()
        if db_error:
            return jsonify({'error': f'Error de conexión a la base de datos: {db_error}'}), 500
            
        agent, agent_error = create_agent(llm, db_connection)
        if agent_error:
            return jsonify({'error': f'Error al crear el agente: {agent_error}'}), 500
        
        # Obtener respuesta del agente
        response = agent.invoke({"input": user_message})
        
        # Procesar la respuesta
        if hasattr(response, 'output') and response.output:
            response_text = response.output
        elif isinstance(response, str):
            response_text = response
        elif hasattr(response, 'get') and callable(response.get) and 'output' in response:
            response_text = response['output']
        else:
            response_text = str(response)
        
        # Eliminar el mensaje almacenado después de procesarlo
        del message_store[message_id]
        
        return jsonify({
            'response': response_text,
            'status': 'success'
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

# Integración con Gradio para Hugging Face Spaces
def mount_in_app(gradio_app):
    """Monta la API Flask en la aplicación Gradio."""
    return gradio_app

if __name__ == '__main__':
    # Si se ejecuta directamente, inicia el servidor Flask
    port = int(os.environ.get('PORT', 5000))
    app.run(host='0.0.0.0', port=port)
else:
    # Si se importa como módulo (en Hugging Face Spaces),
    # expone la función para montar en Gradio
    gradio_app = gr.mount_gradio_app(
        app,
        "/api",  # Prefijo para los endpoints de la API
        lambda: True  # Autenticación deshabilitada
    )