from runtime.rooms.room_manager import room_manager
import json
import uuid

class TaskDistributor:
    """
    Distributes tasks among available peers in the grid computing swarm.
    """
    def __init__(self):
        pass

    def dispatch_subtasks(self, tasks, peers):
        """
        Distributes tasks among available peers using a simple round-robin approach.
        """
        if not peers:
            return {}
            
        assignments = {peer['node_id']: [] for peer in peers}
        peer_ids = [peer['node_id'] for peer in peers]
        
        for i, task in enumerate(tasks):
            peer_id = peer_ids[i % len(peer_ids)]
            assignments[peer_id].append(task)
            
        return assignments

    def dispatch_to_room(self, tasks):
        """Room-aware distribution logic. Dispatches directly to connected P2P peers."""
        if not room_manager.active_room:
            # Fallback to local
            return {"local": tasks}

        peers = room_manager.active_room.get("members", [])
        if not peers:
            # Fallback to just the host
            host_id = room_manager.active_room.get("host_id", "local")
            return {host_id: tasks}

        assignments = {peer['user_id']: [] for peer in peers}
        peer_ids = [peer['user_id'] for peer in peers]
        
        for i, task in enumerate(tasks):
            peer_id = peer_ids[i % len(peer_ids)]
            assignments[peer_id].append(task)
            
            # Announce over WS
            if room_manager.is_host and room_manager.host_service:
                import asyncio
                asyncio.create_task(room_manager.host_service.broadcast({
                    "type": "task_event",
                    "room_id": room_manager.active_room["room_id"],
                    "task_id": task.get("task_id", str(uuid.uuid4())),
                    "assigned_to": peer_id,
                    "status": "dispatched"
                }))
                
        return assignments
