mod player;
mod room;

use crate::room::{generate_room_name, Room, RoomHandle};
use anyhow::Result;
use async_tungstenite::WebSocketStream;
use board_network::protocol::ClientMessage;
use futures::channel::mpsc::{unbounded, UnboundedSender};
use futures::future::join;
use futures::{future, SinkExt, StreamExt};
use smol::Async;
use std::collections::HashMap;
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::sync::{Arc, Mutex};
use std::{env, io};
use tungstenite::Message as WebsocketMessage;

type Rooms = Arc<Mutex<HashMap<String, RoomHandle>>>;

async fn run_player(
    player_name: String,
    addr: SocketAddr,
    handle: RoomHandle,
    ws_stream: WebSocketStream<Async<TcpStream>>,
) {
    let (incoming, outgoing) = ws_stream.split();
    let (tx, rx) = unbounded();

    {
        let room = &mut handle.room.lock().unwrap();
        if let Err(e) = room.add_player(addr, player_name.clone(), tx) {
            eprintln!("[{}] Failed to add player: {:?}", room.name, e);
            return;
        }
    }

    let write = handle.write.clone();
    let ra = rx
        .map(|c| {
            serde_json::to_string(&c).unwrap_or_else(|_| panic!("Could not serialize {:?}", c))
        })
        .map(WebsocketMessage::Text)
        .map(Ok)
        .forward(incoming);
    let rb = outgoing
        .map(|m| match m {
            Ok(WebsocketMessage::Text(t)) => serde_json::from_str::<ClientMessage>(&t).ok(),
            _ => None,
        })
        .take_while(|m| future::ready(m.is_some()))
        .map(|m| m.unwrap())
        .chain(futures::stream::once(async { ClientMessage::Disconnected }))
        .map(move |m| Ok((addr, m)))
        .forward(write);

    let (ra, rb) = join(ra, rb).await;
    if let Err(e) = ra {
        eprintln!("[{addr}] Got error {e} from player {player_name}'s rx queue");
    }
    if let Err(e) = rb {
        eprintln!("[{addr}] Got error {e} from player {player_name}'s tx queue");
    }
    println!("[{addr}] Finished session with {player_name}");
}

async fn handle_connection(
    rooms: Rooms,
    raw_stream: Async<TcpStream>,
    addr: SocketAddr,
    mut close_room: UnboundedSender<String>,
) -> Result<()> {
    println!("[{addr}] Incoming TCP connection");
    let mut ws_stream = async_tungstenite::accept_async(raw_stream).await?;
    println!("[{addr}] WebSocket connection established");

    if let Some(Ok(WebsocketMessage::Text(text))) = ws_stream.next().await {
        let msg = serde_json::from_str::<ClientMessage>(&text)?;
        match msg {
            ClientMessage::CreateRoom(player_name) => {
                let (write, read) = unbounded();
                let room = Arc::new(Mutex::new(Room::default()));
                let handle = RoomHandle { write, room };
                let room_name = generate_room_name(&mut rooms.lock().unwrap(), handle.clone());
                println!("[{addr}] Creating room '{room_name}' for player '{player_name}'");

                let mut h = handle.clone();
                join(
                    h.run_room(read),
                    run_player(player_name, addr, handle, ws_stream),
                )
                .await;

                if let Err(e) = close_room.send(room_name.clone()).await {
                    eprintln!("[{room_name}] Failed to close room: {e}");
                }

                return Ok(());
            }
            msg => eprintln!("[{addr}] Received illegal message {:?}", msg),
        }
    }

    println!("[{addr}] Dropping connection");
    Ok(())
}

fn main() -> Result<(), io::Error> {
    let addr = env::args()
        .nth(1)
        .unwrap_or_else(|| "0.0.0.0:8080".to_string())
        .parse::<SocketAddr>()
        .expect("Invalid address");

    let rooms = Rooms::new(Mutex::new(HashMap::new()));

    let close_room = {
        let (tx, mut rx) = unbounded();
        let rooms = rooms.clone();
        smol::spawn(async move {
            while let Some(room_name) = rx.next().await {
                println!("Closing room {room_name}");
                rooms.lock().unwrap().remove(&room_name);
            }
        })
        .detach();
        tx
    };

    smol::block_on(async {
        println!("Listening on: {addr}");
        let listener = Async::<TcpListener>::bind(addr).expect("Could not create listener");

        while let Ok((stream, addr)) = listener.accept().await {
            let close_room = close_room.clone();
            let rooms = rooms.clone();
            smol::spawn(async move {
                if let Err(e) = handle_connection(rooms, stream, addr, close_room).await {
                    eprintln!("Failed to handle connection from {addr}: {e}");
                }
            })
            .detach();
        }
    });
    Ok(())
}