from flask import Flask, render_template, request, jsonify
from gradio_client import Client
import threading

app = Flask(__name__)

# Initialize the Gradio client
client = Client("m-ric/open_Deep-Research")

def interact_with_agent(messages):
    try:
        response = client.predict(
            messages=messages,
            api_name="/interact_with_agent"
        )
        return response
    except Exception as e:
        return {"error": str(e)}

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/search', methods=['POST'])
def search():
    data = request.get_json()
    query = data.get('query', '').strip()
    if not query:
        return jsonify({"error": "No query provided."}), 400

    # Initialize the conversation with the user's query
    conversation = [
        {"role": "user", "content": query}
    ]

    def background_task():
        try:
            # First, get the initial response from the API
            initial_response = client.predict(
                text_input=query,
                api_name="/log_user_message"
            )

            # Append the assistant's initial response to the conversation
            conversation.append({"role": "assistant", "content": initial_response})

            # Now, interact with the agent using the conversation
            final_response = interact_with_agent(conversation)

            # Extract the final answer from the response
            final_answer = None
            for message in final_response:
                if message.get('role') == 'assistant' and 'Final answer' in message.get('content', ''):
                    final_answer = message.get('content', '').split('Final answer:')[-1].strip()
                    break

            # Send the result back to the client
            socketio.emit('result', {'query': query, 'result': final_answer})
        except Exception as e:
            socketio.emit('result', {'query': query, 'result': f"An error occurred: {str(e)}"})

    # Start the background thread
    thread = threading.Thread(target=background_task)
    thread.start()

    return jsonify({"status": "Processing..."}), 200

if __name__ == '__main__':
    from flask_socketio import SocketIO

    # Initialize SocketIO for real-time updates
    socketio = SocketIO(app, cors_allowed_origins="*")

    @socketio.on('connect')
    def handle_connect():
        print('Client connected')

    @socketio.on('disconnect')
    def handle_disconnect():
        print('Client disconnected')

    # Run the Flask app with SocketIO
    socketio.run(app, host='0.0.0.0', port=7860, allow_unsafe_werkzeug=True)