use std::collections::HashMap;
use std::io::Write;
use std::sync::{Arc, Mutex};
use std::{
    io::Read,
    net::{Shutdown, SocketAddr, TcpListener, TcpStream},
    sync::mpsc::{channel, Sender},
    thread::{self},
};

const BUFFER_SIZE: usize = 1024;

fn main() {
    let listener = TcpListener::bind("127.0.0.1:42069").unwrap();
    let sockets: Arc<Mutex<HashMap<u16, TcpStream>>> = Arc::new(Mutex::new(HashMap::new()));
    let c_sockets = Arc::clone(&sockets);
    let cc_sockets = Arc::clone(&sockets);

    let mut next_socket_id = 0;

    let (tx_broadcast, rx_broadcast) = channel::<Vec<u8>>();
    let (tx_disconnect, rx_disconnect) = channel::<u16>();

    thread::spawn(move || loop {
        let msg = &rx_broadcast.recv().unwrap();
        let lock = sockets.lock().unwrap();
        let sockets = lock.iter();

        for (_id, mut socket) in sockets {
            socket.write(msg).unwrap();
        }
    });

    thread::spawn(move || loop {
        let disconnected_socket_id = &rx_disconnect.recv().unwrap();
        let mut sockets = cc_sockets.lock().unwrap();
        sockets.retain(|id, _socket| id != disconnected_socket_id);
    });

    loop {
        let (socket, addr) = listener.accept().unwrap();
        let socket_id = next_socket_id;

        c_sockets
            .lock()
            .unwrap()
            .insert(socket_id, socket.try_clone().unwrap());
        next_socket_id = next_socket_id + 1;

        let cloned_broadcast_tx = tx_broadcast.clone();
        let cloned_disconnect_tx = tx_disconnect.clone();

        println!("[EVENT] {} has connected", addr);
        thread::spawn(move || {
            handle(
                socket_id,
                socket,
                addr,
                cloned_broadcast_tx,
                cloned_disconnect_tx,
            );
        });
    }
}

fn handle(
    id: u16,
    mut socket: TcpStream,
    addr: SocketAddr,
    sender: Sender<Vec<u8>>,
    disconnect_sender: Sender<u16>,
) {
    loop {
        let mut buf: [u8; BUFFER_SIZE] = [0; BUFFER_SIZE];
        match socket.read(&mut buf) {
            Ok(n) => {
                let msg = String::from_utf8(buf[..n].to_vec())
                    .unwrap()
                    .trim()
                    .to_string();

                if n <= 0 {
                    disconnect_socket(id, socket, disconnect_sender);
                    broadcast_message(&sender, format!("[EVENT] {id} has disconnected."));
                    return;
                }

                if msg.is_empty() {
                    continue;
                }

                broadcast_message(&sender, format!("[MESSAGE] {id}: {msg}"));
            }
            Err(_) => {
                disconnect_socket(id, socket, disconnect_sender);
                broadcast_message(&sender, format!("[EVENT] {id} has disconnected."));
                return;
            }
        }
    }
}

fn disconnect_socket(id: u16, socket: TcpStream, sender: Sender<u16>) {
    sender.send(id).unwrap();
    println!("[EVENT] {} has disconnected", socket.peer_addr().unwrap());
    let _ = socket.shutdown(Shutdown::Both);
}

fn broadcast_message(sender: &Sender<Vec<u8>>, message: String) {
    sender.send(message.as_bytes().to_vec()).unwrap();
}