use std::env;
use std::error::Error;
use std::net::SocketAddr;
use std::sync::Arc;
use std::collections::HashMap;
use tokio::io::{AsyncReadExt, AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::Mutex;

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    // Parse command-line arguments
    let args: Vec<String> = env::args().collect();
    if args.len() != 3 {
        println!("Usage: {} [server|client] [addr:port]", args[0]);
        return Ok(());
    }

    let command = &args[1];
    let addr = &args[2];

    match command.as_str() {
        "server" => start_server(addr).await?,
        "client" => start_client(addr).await?,
        _ => println!("Unknown command {}", command),
    }

    Ok(())
}

async fn start_server(addr: &str) -> Result<(), Box<dyn Error>> {
    // Start the TCP listener
    let listener = TcpListener::bind(addr).await?;
    println!("Server listening on {}", addr);

    // Shared data for client connections
    let clients: Arc<Mutex<HashMap<SocketAddr, String>>> = Arc::new(Mutex::new(HashMap::new()));

    loop {
        // Accept a new client connection
        let (mut socket, client_addr) = listener.accept().await?;

        // Read username from the client
        let mut buf = vec![0; 1024];
        let n = socket.read(&mut buf).await?;
        let username = String::from_utf8_lossy(&buf[..n]).trim().to_string();

        println!("Accepted connection from {}", client_addr);

        let clients_shared = Arc::clone(&clients);

        // Insert client into the shared HashMap
        clients_shared.lock().await.insert(client_addr, username.clone());

        // Clone the HashMap and release the lock
        let clients_map = clients_shared.lock().await.clone();

        // Spawn a new task for each client connection
        tokio::spawn(async move {
            let mut buf = vec![0; 1024];

            loop {
                let n = match socket.read(&mut buf).await {
                    Ok(n) if n == 0 => {
                        // Client disconnected
                        let username = clients_shared.lock().await.remove(&client_addr).unwrap_or("Unknown".to_string());
                        println!("Client {} ({}) disconnected!", username, client_addr);
                        return;
                    }
                    Ok(n) => n,
                    Err(e) => {
                        eprintln!("Failed to read from socket; err = {:?}", e);
                        return;
                    }
                };

                let username = clients_map.get(&client_addr).unwrap_or(&"Unknown".to_string()).clone();

                let message = String::from_utf8_lossy(&buf[..n]).to_string();
                println!(
                    "Received message from ({})\n{}: {}",
                    client_addr,
                    username,
                    message
                );
            }
        });
    }
}

async fn start_client(addr: &str) -> Result<(), Box<dyn Error>> {
    // Connect to the server
    let mut stream = TcpStream::connect(addr).await?;

    println!("Enter a username: ");
    let mut username = String::new();
    let mut reader = BufReader::new(tokio::io::stdin());
    reader.read_line(&mut username).await?;

    // Send username to the server
    stream.write_all(username.trim().as_bytes()).await?;

    println!("Connected to server at {}", addr);
    loop {
        let mut input = String::new();
        reader.read_line(&mut input).await?;
        stream.write_all(input.trim().as_bytes()).await?;

        // Print sent message
        println!("Sent message from {}: {}", username.trim(), input.trim());
    }
}