from flask import Flask, request, jsonify
import requests
import random
import time

app = Flask(__name__)

WORKERS = [f"https://aeuph-edge-worker-{i}.hf.space" for i in range(2, 11)]
tasks = {}
next_worker_index = 0

@app.route('/')
def hello_world():
    return 'Hello, I am the master node!'

@app.route('/submit', methods=['POST'])
def submit_task():
    global next_worker_index
    task_data = request.json
    if not task_data or 'data' not in task_data:
        return jsonify({"error": "Invalid task data"}), 400

    task_id = f"task_{random.randint(1000, 9999)}"
    
    # Simple round-robin worker selection
    worker_url = WORKERS[next_worker_index]
    next_worker_index = (next_worker_index + 1) % len(WORKERS)

    for attempt in range(5):
        try:
            response = requests.post(f"{worker_url}/execute", json={"task_data": task_data['data']}, timeout=15)
            response.raise_for_status()
            tasks[task_id] = {"status": "completed", "result": response.json(), "worker": worker_url}
            return jsonify({"task_id": task_id})
        except requests.exceptions.RequestException as e:
            if attempt < 4:
                time.sleep(10) # Wait for 10 seconds before retrying
            else:
                tasks[task_id] = {"status": "failed", "error": str(e), "worker": worker_url}

    return jsonify({"task_id": task_id})

@app.route('/result/<task_id>', methods=['GET'])
def get_result(task_id):
    task = tasks.get(task_id)
    if not task:
        return jsonify({"error": "Task not found"}), 404
    return jsonify(task)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)