import asyncio
import json
import logging
import time
from typing import Dict, List, Set
from fastapi import WebSocket

logger = logging.getLogger("friday.state_sync")

class GridStateSync:
    """Manages real-time state synchronization across the FRIDAY distributed grid."""
    def __init__(self):
        self.active_connections: Set[WebSocket] = set()
        self.device_stats: Dict[str, dict] = {} # dev_id -> {proximity, last_seen, load}
        
    async def register_device(self, websocket: WebSocket, device_id: str):
        await websocket.accept()
        self.active_connections.add(websocket)
        self.device_stats[device_id] = {
            "proximity": 0.0,
            "last_seen": time.time(),
            "status": "online"
        }
        logger.info(f"Grid: Device '{device_id}' joined the sovereign mesh.")
        
    async def unregister_device(self, websocket: WebSocket, device_id: str):
        self.active_connections.remove(websocket)
        if device_id in self.device_stats:
            self.device_stats[device_id]["status"] = "offline"
        logger.info(f"Grid: Device '{device_id}' left the mesh.")

    async def broadcast(self, message: dict):
        """Broadcasts a message to all connected devices in the grid."""
        if not self.active_connections: return
        
        payload = json.dumps(message)
        disconnected = []
        # Iterate over a copy to prevent "Set changed size during iteration" crash
        for ws in list(self.active_connections):
            try:
                await ws.send_text(payload)
            except Exception:
                disconnected.append(ws)
        
        for ws in disconnected:
            self.active_connections.remove(ws)

    def update_proximity(self, device_id: str, rssi_or_volume: float):
        """Updates the proximity score for a device to enable Proximity Logic."""
        if device_id in self.device_stats:
            self.device_stats[device_id]["proximity"] = rssi_or_volume
            self.device_stats[device_id]["last_seen"] = time.time()

    def get_best_output_device(self) -> str:
        """Determines the device closest to the user (Highest Proximity)."""
        if not self.device_stats: return "core"
        
        # Sort by proximity, then by recency
        sorted_devices = sorted(
            self.device_stats.items(), 
            key=lambda x: (x[1]["proximity"], x[1]["last_seen"]), 
            reverse=True
        )
        return sorted_devices[0][0] if sorted_devices else "core"

# Global Singleton
grid_sync = GridStateSync()