#![allow(dead_code, unused_must_use)]
use std::collections::HashMap;
use std::time::SystemTime;
use thiserror::Error;
use crate::common::serverstate::ClientId;

use async_std::sync::{Arc, RwLock};
use futures::{stream, StreamExt};
use std::pin::pin;

pub const MAX_ROOMS: usize = 128;

pub enum AreaType {
    Room,
    Lobby,
}


#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct LobbyId(pub usize);

#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, derive_more::Display)]
pub struct RoomId(pub usize);

impl LobbyId {
    pub fn id(&self) -> u8 {
        self.0 as u8
    }
}


#[derive(Error, Debug, PartialEq, Eq)]
pub enum CreateRoomError {
    #[error("no open slots")]
    NoOpenSlots,
    #[error("client already in area")]
    ClientInAreaAlready,
    #[error("join error")]
    JoinError,
}

#[derive(Error, Debug, PartialEq, Eq)]
pub enum JoinRoomError {
    #[error("room does not exist")]
    RoomDoesNotExist,
    #[error("room is full")]
    RoomFull,
    #[error("client already in area")]
    ClientInAreaAlready,
}

#[derive(Error, Debug, PartialEq, Eq)]
pub enum JoinLobbyError {
    #[error("lobby does not exist")]
    LobbyDoesNotExist,
    #[error("lobby is full")]
    LobbyFull,
    #[error("client already in area")]
    ClientInAreaAlready,
}

#[derive(Error, Debug, PartialEq, Eq)]
pub enum GetAreaError {
    #[error("not in a room")]
    NotInRoom,
    #[error("not in a lobby")]
    NotInLobby,
    #[error("get area: invalid client")]
    InvalidClient,
}

#[derive(Error, Debug, PartialEq, Eq)]
pub enum ClientRemovalError {
    #[error("client removal: client not in area")]
    ClientNotInArea,
    #[error("client removal: invalid area")]
    InvalidArea,
}

#[derive(Error, Debug, PartialEq, Eq)]
pub enum GetClientsError {
    #[error("invalid client")]
    InvalidClient,
    #[error("invalid area")]
    InvalidArea,
}

#[derive(Error, Debug, PartialEq, Eq)]
pub enum GetNeighborError {
    #[error("get neighbor: invalid client")]
    InvalidClient,
    #[error("get neighbor: invalid area")]
    InvalidArea,
}

#[derive(Error, Debug, PartialEq, Eq)]
pub enum GetLeaderError {
    #[error("get leader: invalid client")]
    InvalidClient,
    #[error("get leader: invalid area")]
    InvalidArea,
    #[error("get leader: client not in area")]
    NoClientInArea,
}

#[derive(Error, Debug, PartialEq, Eq)]
pub enum ClientLocationError {
    #[error("create room error {0}")]
    CreateRoomError(#[from] CreateRoomError),
    #[error("join room error {0}")]
    JoinRoomError(#[from] JoinRoomError),
    #[error("join lobby error {0}")]
    JoinLobbyError(#[from] JoinLobbyError),
    #[error("get area error {0}")]
    GetAreaError(#[from] GetAreaError),
    #[error("client removal error {0}")]
    ClientRemovalError(#[from] ClientRemovalError),
    #[error("get clients error {0}")]
    GetClientsError(#[from] GetClientsError),
    #[error("get neighbor error {0}")]
    GetNeighborError(#[from] GetNeighborError),
    #[error("get leader error {0}")]
    GetLeaderError(#[from] GetLeaderError)
}


#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct LocalClientId(usize);

impl LocalClientId {
    pub fn id(&self) -> u8 {
        self.0 as u8
    }
}

impl PartialEq<u8> for LocalClientId {
    fn eq(&self, other: &u8) -> bool {
        self.0 == *other as usize
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct AreaClient {
    pub client: ClientId,
    pub local_client: LocalClientId,
    time_join: SystemTime,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
struct Lobby([Option<AreaClient>; 12]);
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
struct Room([Option<AreaClient>; 4]);

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum RoomLobby {
    Room(RoomId),
    Lobby(LobbyId),
}

#[derive(Clone, Debug)]
pub struct ClientLocation {
    lobbies: [Arc<RwLock<Lobby>>; 15],
    rooms: [Arc<RwLock<Option<Room>>>; MAX_ROOMS],
    client_location: Arc<RwLock<HashMap<ClientId, RoomLobby>>>,
}

impl Default for ClientLocation {
    fn default() -> ClientLocation {
        ClientLocation {
            lobbies: core::array::from_fn(|_| Arc::new(RwLock::new(Lobby([None; 12])))),
            rooms: core::array::from_fn(|_| Arc::new(RwLock::new(None))),
            client_location: Arc::new(RwLock::new(HashMap::new())),
        }
    }
}


impl ClientLocation {
    pub async fn add_client_to_lobby(&self, id: ClientId, lobby_id: LobbyId) -> Result<(), JoinLobbyError> {
        {
            let lobby = self.lobbies
                .get(lobby_id.0)
                .ok_or(JoinLobbyError::LobbyDoesNotExist)?
                .read()
                .await;
            if lobby.0.iter().all(|c| c.is_some()) {
                return Err(JoinLobbyError::LobbyFull);
            }
        }
        self.remove_client_from_area(id).await;

        let mut lobby = self.lobbies
            .get(lobby_id.0)
            .ok_or(JoinLobbyError::LobbyDoesNotExist)?
            .write()
            .await;
        let (index, empty_slot) = lobby.0.iter_mut()
            .enumerate()
            .find(|(_, k)| k.is_none())
            .ok_or(JoinLobbyError::LobbyFull)?;
        *empty_slot = Some(AreaClient {
            client: id,
            local_client: LocalClientId(index),
            time_join: SystemTime::now(),
        });
        self.client_location
            .write()
            .await
            .insert(id, RoomLobby::Lobby(lobby_id));
        Ok(())
    }

    pub async fn add_client_to_next_available_lobby(&self, id: ClientId, lobby: LobbyId) -> Result<LobbyId, JoinLobbyError> {
        pin!(stream::iter(0..15)
            .filter_map(|lobby_index| async move {
                let new_lobby = LobbyId((lobby.0 + lobby_index) % 15);
                Some((new_lobby, self.add_client_to_lobby(id, new_lobby).await.ok()?))
            }))
            .next()
            .await
            .map(|l| l.0)
            .ok_or(JoinLobbyError::LobbyFull)
    }

    pub async fn create_new_room(&mut self, id: ClientId) -> Result<RoomId, CreateRoomError> {
        let (index, empty_slot) = Box::pin(stream::iter(self.rooms.iter())
            .enumerate()
            .filter(|(_, r)| async {r.read().await.is_none()}))
            .next()
            .await
            .ok_or(CreateRoomError::NoOpenSlots)?;
        *empty_slot.write().await = Some(Room([None; 4]));
        self.add_client_to_room(id, RoomId(index))
            .await
            .map_err(|_err| CreateRoomError::JoinError)?;
        Ok(RoomId(index))
    }

    pub async fn add_client_to_room(&mut self, id: ClientId, room: RoomId) -> Result<(), JoinRoomError> {
        let mut r = self.rooms.get(room.0)
            .ok_or(JoinRoomError::RoomDoesNotExist)?
            .as_ref()
            .write()
            .await;
        let r = r.as_mut()
            .ok_or(JoinRoomError::RoomDoesNotExist)?;
        let (index, empty_slot) = r.0.iter_mut()
            .enumerate()
            .find(|(_, k)| k.is_none())
            .ok_or(JoinRoomError::RoomFull)?;
        *empty_slot = Some(AreaClient {
            client: id,
            local_client: LocalClientId(index),
            time_join: SystemTime::now(),
        });
        self.remove_client_from_area(id).await;
        self.client_location
            .write()
            .await
            .insert(id, RoomLobby::Room(room));
        Ok(())
    }

    pub async fn get_all_clients_by_client(&self, id: ClientId) -> Result<Vec<AreaClient>, GetNeighborError> {
        let area = self.client_location
            .read()
            .await;
        let area = area
            .get(&id)
            .ok_or(GetNeighborError::InvalidClient)?;
        match area {
            RoomLobby::Room(room) => {
                Ok(self.get_clients_in_room(*room).await.map_err(|_| GetNeighborError::InvalidArea)?
                   .into_iter()
                   .collect())
            },
            RoomLobby::Lobby(lobby) => {
                Ok(self.get_clients_in_lobby(*lobby).await.map_err(|_| GetNeighborError::InvalidArea)?
                   .into_iter()
                   .collect())
            }
        }
    }

    pub async fn get_client_neighbors(&self, id: ClientId) -> Result<Vec<AreaClient>, GetNeighborError> {
        let area = self.client_location
            .read()
            .await;
        let area = area
            .get(&id)
            .ok_or(GetNeighborError::InvalidClient)?;
        match area {
            RoomLobby::Room(room) => {
                Ok(self.get_clients_in_room(*room).await.map_err(|_| GetNeighborError::InvalidArea)?
                   .into_iter()
                   .filter(|c| c.client != id)
                   .collect())
            },
            RoomLobby::Lobby(lobby) => {
                Ok(self.get_clients_in_lobby(*lobby).await.map_err(|_| GetNeighborError::InvalidArea)?
                   .into_iter()
                   .filter(|c| c.client != id)
                   .collect())
            }
        }
    }

    pub async fn get_room_leader(&self, room: RoomId) -> Result<AreaClient, GetLeaderError> {
        let r = self.rooms[room.0]
            .as_ref()
            .read()
            .await
            .ok_or(GetLeaderError::InvalidArea)?;
        let mut r = r
            .0
            .iter()
            .flatten()
            .collect::<Vec<_>>();
        r.sort_by_key(|k| k.time_join);
        let c = r.get(0)
            .ok_or(GetLeaderError::NoClientInArea)?;
        Ok(**c)
    }

    pub async fn get_lobby_leader(&self, lobby: LobbyId) -> Result<AreaClient, GetLeaderError> {
        let l = self.lobbies[lobby.0]
            .read()
            .await;
        let mut l = l
            .0
            .iter()
            .flatten()
            .collect::<Vec<_>>();
        l.sort_by_key(|k| k.time_join);
        let c = l.get(0).ok_or(GetLeaderError::NoClientInArea)?;
        Ok(**c)
    }

    pub async fn get_area_leader(&self, roomlobby: RoomLobby) -> Result<AreaClient, GetLeaderError> {
        match roomlobby {
            RoomLobby::Room(room) => {
                self.get_room_leader(room).await
            },
            RoomLobby::Lobby(lobby) => {
                self.get_lobby_leader(lobby).await
            }
        }
    }

    pub async fn get_leader_by_client(&self, id: ClientId) -> Result<AreaClient, GetLeaderError> {
        let area = self.client_location
            .read()
            .await;
        let area = area
            .get(&id)
            .ok_or(GetLeaderError::InvalidClient)?;
        match area {
            RoomLobby::Room(room) => {
                self.get_room_leader(*room).await
            },
            RoomLobby::Lobby(lobby) => {
                self.get_lobby_leader(*lobby).await
            }
        }
    }

    pub async fn get_clients_in_lobby(&self, lobby: LobbyId) -> Result<Vec<AreaClient>, GetClientsError> {
        Ok(self.lobbies
           .get(lobby.0)
           .ok_or(GetClientsError::InvalidArea)?
           .read()
           .await
           .0
           .iter()
           .filter_map(|client| {
               client.map(|c| {
                   c
               })
           }).collect())
    }

    pub async fn get_clients_in_room(&self, room: RoomId) -> Result<Vec<AreaClient>, GetClientsError> {
        Ok(self.rooms.get(room.0)
           .ok_or(GetClientsError::InvalidArea)?
           .as_ref()
           .read()
           .await
           .ok_or(GetClientsError::InvalidArea)?
           .0
           .iter()
           .filter_map(|client| {
               client.map(|c| {
                   c
               })
           }).collect())
    }

    pub async fn get_local_client(&self, id: ClientId) -> Result<AreaClient, GetClientsError> {
        let area = self.client_location
            .read()
            .await;
        let area = area
            .get(&id)
            .ok_or(GetClientsError::InvalidClient)?;
        match area {
            RoomLobby::Room(room) => {
                self.get_clients_in_room(*room)
                    .await
                    .map_err(|_| GetClientsError::InvalidArea)?
                    .into_iter()
                    .find(|c| c.client == id)
                    .ok_or(GetClientsError::InvalidClient)
            },
            RoomLobby::Lobby(lobby) => {
                self.get_clients_in_lobby(*lobby)
                    .await
                    .map_err(|_| GetClientsError::InvalidArea)?
                    .into_iter()
                    .find(|c| c.client == id)
                    .ok_or(GetClientsError::InvalidClient)
            }
        }
    }

    pub async fn get_area(&self, id: ClientId) -> Result<RoomLobby, GetAreaError> {
        self.client_location
            .read()
            .await
            .get(&id)
            .ok_or(GetAreaError::InvalidClient)
            .map(Clone::clone)
    }

    pub async fn get_room(&self, id: ClientId) -> Result<RoomId, GetAreaError> {
        if let RoomLobby::Room(room) = self.client_location.read().await.get(&id).ok_or(GetAreaError::InvalidClient)? {
            Ok(*room)
        }
        else {
            Err(GetAreaError::NotInRoom)
        }
    }

    pub async fn get_lobby(&self, id: ClientId) -> Result<LobbyId, GetAreaError> {
        if let RoomLobby::Lobby(lobby) = self.client_location.read().await.get(&id).ok_or(GetAreaError::InvalidClient)? {
            Ok(*lobby)
        }
        else {
            Err(GetAreaError::NotInLobby)
        }
    }

    pub async fn remove_client_from_area(&self, id: ClientId) -> Result<(), ClientRemovalError> {
        fn remove_client<const N: usize>(id: ClientId, client_list : &mut [Option<AreaClient>; N]) {
            client_list
                .iter_mut()
                .filter(|client| {
                    client.map_or(false, |c| {
                        c.client == id
                    })
                })
                .for_each(|client| {
                    *client = None
                });
        }

        let area = self.client_location
            .read()
            .await;
        let area = area
            .get(&id)
            .ok_or(ClientRemovalError::ClientNotInArea)?;
        match area {
            RoomLobby::Room(room) => {
                let mut r = self.rooms.get(room.0)
                    .ok_or(ClientRemovalError::InvalidArea)?
                    .as_ref()
                    .write()
                    .await;
                if let Some(r) = r.as_mut() {
                    remove_client(id, &mut r.0)
                }
                else {
                    return Err(ClientRemovalError::InvalidArea)
                }
            },
            RoomLobby::Lobby(lobby) => {
                remove_client(id, &mut self.lobbies[lobby.0].write().await.0)
            }
        };

        Ok(())
    }
}








#[cfg(test)]
mod test {
    use super::*;

    #[async_std::test]
    async fn test_add_client_to_lobby() {
        let cl = ClientLocation::default();
        cl.add_client_to_lobby(ClientId(12), LobbyId(0)).await.unwrap();
        cl.add_client_to_lobby(ClientId(13), LobbyId(1)).await.unwrap();
        cl.add_client_to_lobby(ClientId(14), LobbyId(0)).await.unwrap();

        assert!(cl.get_clients_in_lobby(LobbyId(0)).await.into_iter().flatten().map(|c| (c.client, c.local_client)).collect::<Vec<_>>() == vec![
            (ClientId(12), LocalClientId(0)),
            (ClientId(14), LocalClientId(1)),
        ]);
    }

    #[async_std::test]
    async fn test_add_client_to_full_lobby() {
        let cl = ClientLocation::default();
        for i in 0..12 {
            cl.add_client_to_lobby(ClientId(i), LobbyId(0)).await.unwrap();
        }
        assert!(cl.add_client_to_lobby(ClientId(99), LobbyId(0)).await == Err(JoinLobbyError::LobbyFull));
    }

    #[async_std::test]
    async fn test_add_client_to_next_available_lobby() {
        let cl = ClientLocation::default();
        for lobby in 1..4 {
            for i in 0..12 {
                cl.add_client_to_lobby(ClientId(lobby*12+i), LobbyId(lobby)).await.unwrap();
            }
        }
        assert!(cl.add_client_to_next_available_lobby(ClientId(99), LobbyId(1)).await == Ok(LobbyId(4)));
    }

    #[async_std::test]
    async fn test_add_to_lobby_when_all_are_full() {
        let cl = ClientLocation::default();
        for lobby in 0..15 {
            for i in 0..12 {
                cl.add_client_to_lobby(ClientId(lobby*12+i), LobbyId(lobby)).await.unwrap();
            }
        }
        assert_eq!(cl.add_client_to_next_available_lobby(ClientId(99), LobbyId(1)).await, Err(JoinLobbyError::LobbyFull));
    }

    #[async_std::test]
    async fn test_new_room() {
        let mut cl = ClientLocation::default();
        assert!(cl.create_new_room(ClientId(12)).await == Ok(RoomId(0)));
    }

    #[async_std::test]
    async fn test_add_client_to_room() {
        let mut cl = ClientLocation::default();
        let room = cl.create_new_room(ClientId(12)).await.unwrap();
        assert!(cl.add_client_to_room(ClientId(234), room).await == Ok(()));
        assert!(cl.get_clients_in_room(room).await.unwrap().into_iter().map(|c| (c.client, c.local_client)).collect::<Vec<_>>() == vec![
            (ClientId(12), LocalClientId(0)),
            (ClientId(234), LocalClientId(1)),
        ]);
    }

    #[async_std::test]
    async fn test_no_new_room_slots() {
        let mut cl = ClientLocation::default();
        for i in 0..128 {
            cl.create_new_room(ClientId(i)).await;
        }
        assert!(cl.create_new_room(ClientId(234)).await == Err(CreateRoomError::NoOpenSlots));
    }

    #[async_std::test]
    async fn test_joining_full_room() {
        let mut cl = ClientLocation::default();
        let room = cl.create_new_room(ClientId(0)).await.unwrap();
        assert!(cl.add_client_to_room(ClientId(1), room).await == Ok(()));
        assert!(cl.add_client_to_room(ClientId(2), room).await == Ok(()));
        assert!(cl.add_client_to_room(ClientId(3), room).await == Ok(()));
        assert!(cl.add_client_to_room(ClientId(234), room).await == Err(JoinRoomError::RoomFull));
    }

    #[async_std::test]
    async fn test_adding_client_to_room_removes_from_lobby() {
        let mut cl = ClientLocation::default();
        cl.add_client_to_lobby(ClientId(93), LobbyId(0)).await;
        cl.add_client_to_lobby(ClientId(23), LobbyId(0)).await;
        cl.add_client_to_lobby(ClientId(51), LobbyId(0)).await;
        cl.add_client_to_lobby(ClientId(12), LobbyId(0)).await;

        let room = cl.create_new_room(ClientId(51)).await.unwrap();
        assert!(cl.add_client_to_room(ClientId(93), room).await == Ok(()));
        assert!(cl.get_clients_in_lobby(LobbyId(0)).await.unwrap().into_iter().map(|c| (c.client, c.local_client)).collect::<Vec<_>>() == vec![
            (ClientId(23), LocalClientId(1)),
            (ClientId(12), LocalClientId(3)),
        ]);
        assert!(cl.get_clients_in_room(room).await.unwrap().into_iter().map(|c| (c.client, c.local_client)).collect::<Vec<_>>() == vec![
            (ClientId(51), LocalClientId(0)),
            (ClientId(93), LocalClientId(1)),
        ]);
    }

    #[async_std::test]
    async fn test_getting_neighbors() {
        let cl = ClientLocation::default();
        cl.add_client_to_lobby(ClientId(93), LobbyId(0)).await.unwrap();
        cl.add_client_to_lobby(ClientId(23), LobbyId(0)).await.unwrap();
        cl.add_client_to_lobby(ClientId(51), LobbyId(0)).await.unwrap();
        cl.add_client_to_lobby(ClientId(12), LobbyId(0)).await.unwrap();

        assert!(cl.get_client_neighbors(ClientId(23)).await.unwrap().into_iter().map(|c| (c.client, c.local_client)).collect::<Vec<_>>() == vec![
            (ClientId(93), LocalClientId(0)),
            (ClientId(51), LocalClientId(2)),
            (ClientId(12), LocalClientId(3)),
        ]);
    }

    #[async_std::test]
    async fn test_failing_to_join_lobby_does_not_remove_from_current_area() {
        let cl = ClientLocation::default();
        for i in 0..12 {
            cl.add_client_to_lobby(ClientId(i), LobbyId(0)).await.unwrap();
        }
        assert!(cl.add_client_to_lobby(ClientId(99), LobbyId(1)).await.is_ok());
        assert!(cl.add_client_to_lobby(ClientId(99), LobbyId(0)).await.is_err());
        assert_eq!(cl.get_clients_in_lobby(LobbyId(0)).await.unwrap().len(), 12);
        assert_eq!(
            cl.get_clients_in_lobby(LobbyId(1)).await.unwrap().into_iter().map(|c| (c.client, c.local_client)).collect::<Vec<_>>(),
            vec![(ClientId(99), LocalClientId(0))]
        );
    }

    #[async_std::test]
    async fn test_get_leader() {
        let cl = ClientLocation::default();
        cl.add_client_to_lobby(ClientId(93), LobbyId(0)).await;
        cl.add_client_to_lobby(ClientId(23), LobbyId(0)).await;
        cl.add_client_to_lobby(ClientId(51), LobbyId(0)).await;
        cl.add_client_to_lobby(ClientId(12), LobbyId(0)).await;

        assert!(cl.get_leader_by_client(ClientId(51)).await.map(|c| (c.client, c.local_client)) == Ok((ClientId(93), LocalClientId(0))));
    }

    #[async_std::test]
    async fn test_remove_client_from_room() {
        let mut cl = ClientLocation::default();
        let room = cl.create_new_room(ClientId(51)).await.unwrap();
        cl.add_client_to_room(ClientId(93), room).await;
        cl.add_client_to_room(ClientId(23), room).await;
        cl.remove_client_from_area(ClientId(51)).await;
        cl.add_client_to_room(ClientId(12), room).await;

        assert!(cl.get_clients_in_room(room).await.unwrap().into_iter().map(|c| (c.client, c.local_client)).collect::<Vec<_>>() == vec![
            (ClientId(12), LocalClientId(0)),
            (ClientId(93), LocalClientId(1)),
            (ClientId(23), LocalClientId(2)),
        ]);
    }

    #[async_std::test]
    async fn test_leader_changes_on_leader_leaving() {
        let mut cl = ClientLocation::default();
        let room = cl.create_new_room(ClientId(51)).await.unwrap();
        cl.add_client_to_room(ClientId(93), room).await.unwrap();
        cl.add_client_to_room(ClientId(23), room).await.unwrap();
        cl.remove_client_from_area(ClientId(51)).await.unwrap();
        cl.add_client_to_room(ClientId(12), room).await.unwrap();
        cl.remove_client_from_area(ClientId(23)).await.unwrap();
        cl.add_client_to_room(ClientId(99), room).await.unwrap();

        assert!(cl.get_leader_by_client(ClientId(12)).await.map(|c| (c.client, c.local_client)) == Ok((ClientId(93), LocalClientId(1))));
    }
}