use std::collections::HashMap;
use std::io;
use std::net::SocketAddr;
use std::sync::LazyLock;

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{Mutex, mpsc};
use tokio::task::JoinHandle;

use super::SERVER_ADDR;

pub(crate) static TCP_SERVER: LazyLock<Mutex<TcpServer>> =
    LazyLock::new(|| Mutex::new(TcpServer::new()));

pub(crate) struct TcpServer {
    is_running: bool,
    clients: HashMap<SocketAddr, JoinHandle<()>>,
    shutdown_tx: Option<mpsc::Sender<()>>,
}

impl TcpServer {
    fn new() -> Self {
        Self {
            is_running: false,
            clients: HashMap::new(),
            shutdown_tx: None,
        }
    }

    pub(crate) async fn start(&mut self) {
        if self.is_running {
            log::warn!("TCP server is already running");
            return;
        }

        match TcpListener::bind(SERVER_ADDR).await {
            Ok(listener) => {
                let (shutdown_tx, shutdown_rx) = mpsc::channel(1);
                self.is_running = true;
                self.shutdown_tx = Some(shutdown_tx);

                tokio::spawn(async move {
                    Self::listen_to_clients(listener, shutdown_rx).await;
                });

                log::info!("TCP Server started on {}", SERVER_ADDR);
            }
            Err(e) => log::error!("Failed to bind to address: {e}"),
        }
    }

    pub(crate) async fn stop(&mut self) {
        if !self.is_running {
            return;
        }

        self.is_running = false;

        if let Some(shutdown_tx) = self.shutdown_tx.take() {
            _ = shutdown_tx.send(()).await;
        }

        for (addr, connection) in self.clients.drain() {
            log::info!("Closing connection to {}", addr);
            connection.abort();
        }
    }

    async fn listen_to_clients(listener: TcpListener, mut shutdown_rx: mpsc::Receiver<()>) {
        loop {
            tokio::select! {
                result = listener.accept() => {
                    match result {
                        Ok((socket, addr)) => {
                            log::info!("New client connected: {addr}");

                            let task_handle = tokio::spawn(async move {
                                if let Err(e) = Self::handle_client(socket, addr).await {
                                    log::error!("Client {addr} error: {e}");
                                }
                                log::info!("Client {addr} disconnected");
                            });

                            let mut server = TCP_SERVER.lock().await;
                            server.clients.insert(addr, task_handle);
                        }
                        Err(e) => log::error!("Couldn't get client: {e}"),
                    }
                }

                _ = shutdown_rx.recv() => {
                    log::info!("TCP Server shutting down");
                    break;
                }
            }
        }
    }

    async fn handle_client(mut socket: TcpStream, addr: SocketAddr) -> io::Result<()> {
        let mut buffer = [0; 1024];

        loop {
            let len = socket.read(&mut buffer).await?;
            if len == 0 {
                break;
            }

            log::debug!("Received {} bytes from {}", len, addr);

            // TODO: Deserialize data
            socket.write_all(&buffer[..len]).await?;
        }

        let mut server = TCP_SERVER.lock().await;
        server.clients.remove(&addr);

        Ok(())
    }
}