use std::{collections::HashMap, net::SocketAddr};

use axum::extract::ws::{Message, WebSocket};
use letterbox_shared::WebsocketMessage;
use tracing::{info, warn};

#[derive(Default)]
pub struct ConnectionTracker {
    peers: HashMap<SocketAddr, WebSocket>,
}

impl ConnectionTracker {
    pub async fn add_peer(&mut self, socket: WebSocket, who: SocketAddr) {
        warn!("adding {who:?} to connection tracker");
        self.peers.insert(who, socket);
        self.send_message_all(WebsocketMessage::RefreshMessages)
            .await;
    }
    pub async fn send_message_all(&mut self, msg: WebsocketMessage) {
        info!("send_message_all {msg}");
        let m = serde_json::to_string(&msg).expect("failed to json encode WebsocketMessage");
        let mut bad_peers = Vec::new();
        for (who, socket) in &mut self.peers.iter_mut() {
            if let Err(e) = socket.send(Message::Text(m.clone().into())).await {
                warn!("{:?} is bad, scheduling for removal: {e}", who);
                bad_peers.push(who.clone());
            }
        }

        for b in bad_peers {
            info!("removing bad peer {b:?}");
            self.peers.remove(&b);
        }
    }
}