use std::collections::HashMap;

use axum::extract::ws::{Message, WebSocket};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use tokio::sync::{RwLock, mpsc};
use uuid::Uuid;

use crate::{AppState, db};

#[derive(Default)]
pub struct VoiceHub {
    rooms: RwLock<HashMap<Uuid, HashMap<Uuid, HashMap<Uuid, ClientHandle>>>>,
}

#[derive(Clone)]
struct ClientHandle {
    display_name: String,
    is_sharing_video: bool,
    is_sharing_screen: bool,
    is_speaking: bool,
    is_muted: bool,
    tx: mpsc::UnboundedSender<ServerEvent>,
}

#[derive(Serialize, Clone)]
pub struct VoiceParticipant {
    pub user_id: Uuid,
    pub display_name: String,
    pub is_sharing_video: bool,
    pub is_sharing_screen: bool,
    pub is_speaking: bool,
    pub is_muted: bool,
}

#[derive(Serialize, Clone)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ServerEvent {
    Peers {
        peers: Vec<VoiceParticipant>,
    },
    PeerJoined {
        user_id: Uuid,
        display_name: String,
    },
    PeerLeft {
        user_id: Uuid,
    },
    VideoStatusChanged {
        user_id: Uuid,
        is_sharing_video: bool,
    },
    ScreenStatusChanged {
        user_id: Uuid,
        is_sharing_screen: bool,
    },
    SpeakingStatusChanged {
        user_id: Uuid,
        is_speaking: bool,
    },
    MuteStatusChanged {
        user_id: Uuid,
        is_muted: bool,
    },
    Signal {
        from_user_id: Uuid,
        kind: String,
        data: serde_json::Value,
    },
    Error {
        message: String,
    },
    PlaySound {
        user_id: Uuid,
        media_url: String,
    },
}

#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ClientEvent {
    Signal {
        to_user_id: Uuid,
        kind: String,
        data: serde_json::Value,
    },
    SetVideoStatus {
        is_sharing_video: bool,
    },
    SetScreenStatus {
        is_sharing_screen: bool,
    },
    SetSpeakingStatus {
        is_speaking: bool,
    },
    SetMuteStatus {
        is_muted: bool,
    },
    PlaySound {
        sound_id: Uuid,
    },
}

#[derive(Default)]
struct VoiceStateChange {
    joined: Option<VoiceParticipant>,
    left_user_id: Option<Uuid>,
    video_changed: Option<(Uuid, bool)>,
    screen_changed: Option<(Uuid, bool)>,
    speaking_changed: Option<(Uuid, bool)>,
    mute_changed: Option<(Uuid, bool)>,
}

impl VoiceHub {
    pub async fn participants(&self, room_id: Uuid) -> Vec<VoiceParticipant> {
        let rooms = self.rooms.read().await;
        let Some(room) = rooms.get(&room_id) else {
            return Vec::new();
        };

        room.iter()
            .filter_map(|(user_id, connections)| aggregate_participant(Some(connections), *user_id))
            .collect()
    }

    async fn join(
        &self,
        room_id: Uuid,
        user_id: Uuid,
        connection_id: Uuid,
        display_name: String,
        tx: mpsc::UnboundedSender<ServerEvent>,
    ) -> (Vec<VoiceParticipant>, VoiceStateChange) {
        let mut rooms = self.rooms.write().await;
        let room = rooms.entry(room_id).or_default();

        let peers = room
            .iter()
            .filter_map(|(peer_id, connections)| aggregate_participant(Some(connections), *peer_id))
            .collect::<Vec<_>>();

        let previous = aggregate_participant(room.get(&user_id), user_id);
        room.entry(user_id).or_default().insert(
            connection_id,
            ClientHandle {
                display_name,
                is_sharing_video: false,
                is_sharing_screen: false,
                is_speaking: false,
                is_muted: false,
                tx,
            },
        );
        let current = aggregate_participant(room.get(&user_id), user_id);

        (peers, diff_voice_state(previous, current))
    }

    async fn leave(&self, room_id: Uuid, user_id: Uuid, connection_id: Uuid) -> VoiceStateChange {
        let mut rooms = self.rooms.write().await;
        let Some(room) = rooms.get_mut(&room_id) else {
            return VoiceStateChange::default();
        };

        let previous = aggregate_participant(room.get(&user_id), user_id);
        if let Some(connections) = room.get_mut(&user_id) {
            connections.remove(&connection_id);
            if connections.is_empty() {
                room.remove(&user_id);
            }
        }
        let current = aggregate_participant(room.get(&user_id), user_id);

        if room.is_empty() {
            rooms.remove(&room_id);
        }

        diff_voice_state(previous, current)
    }

    pub async fn relay_signal(
        &self,
        room_id: Uuid,
        from_user_id: Uuid,
        to_user_id: Uuid,
        kind: String,
        data: serde_json::Value,
    ) {
        let rooms = self.rooms.read().await;
        let Some(room) = rooms.get(&room_id) else {
            return;
        };

        if let Some(target) = room
            .get(&to_user_id)
            .and_then(|connections| connections.values().next())
        {
            let _ = target.tx.send(ServerEvent::Signal {
                from_user_id,
                kind,
                data,
            });
        }
    }

    async fn set_video_status(
        &self,
        room_id: Uuid,
        user_id: Uuid,
        connection_id: Uuid,
        is_sharing_video: bool,
    ) -> VoiceStateChange {
        let mut rooms = self.rooms.write().await;
        let Some(room) = rooms.get_mut(&room_id) else {
            return VoiceStateChange::default();
        };

        let previous = aggregate_participant(room.get(&user_id), user_id);
        if let Some(handle) = room
            .get_mut(&user_id)
            .and_then(|connections| connections.get_mut(&connection_id))
        {
            handle.is_sharing_video = is_sharing_video;
        }
        let current = aggregate_participant(room.get(&user_id), user_id);
        diff_voice_state(previous, current)
    }

    async fn set_screen_status(
        &self,
        room_id: Uuid,
        user_id: Uuid,
        connection_id: Uuid,
        is_sharing_screen: bool,
    ) -> VoiceStateChange {
        let mut rooms = self.rooms.write().await;
        let Some(room) = rooms.get_mut(&room_id) else {
            return VoiceStateChange::default();
        };

        let previous = aggregate_participant(room.get(&user_id), user_id);
        if let Some(handle) = room
            .get_mut(&user_id)
            .and_then(|connections| connections.get_mut(&connection_id))
        {
            handle.is_sharing_screen = is_sharing_screen;
        }
        let current = aggregate_participant(room.get(&user_id), user_id);
        diff_voice_state(previous, current)
    }

    async fn set_speaking_status(
        &self,
        room_id: Uuid,
        user_id: Uuid,
        connection_id: Uuid,
        is_speaking: bool,
    ) -> VoiceStateChange {
        let mut rooms = self.rooms.write().await;
        let Some(room) = rooms.get_mut(&room_id) else {
            return VoiceStateChange::default();
        };

        let previous = aggregate_participant(room.get(&user_id), user_id);
        if let Some(handle) = room
            .get_mut(&user_id)
            .and_then(|connections| connections.get_mut(&connection_id))
        {
            handle.is_speaking = is_speaking;
        }
        let current = aggregate_participant(room.get(&user_id), user_id);
        diff_voice_state(previous, current)
    }

    async fn set_mute_status(
        &self,
        room_id: Uuid,
        user_id: Uuid,
        connection_id: Uuid,
        is_muted: bool,
    ) -> VoiceStateChange {
        let mut rooms = self.rooms.write().await;
        let Some(room) = rooms.get_mut(&room_id) else {
            return VoiceStateChange::default();
        };

        let previous = aggregate_participant(room.get(&user_id), user_id);
        if let Some(handle) = room
            .get_mut(&user_id)
            .and_then(|connections| connections.get_mut(&connection_id))
        {
            handle.is_muted = is_muted;
        }
        let current = aggregate_participant(room.get(&user_id), user_id);
        diff_voice_state(previous, current)
    }

    pub async fn play_sound(&self, room_id: Uuid, user_id: Uuid, media_url: String) {
        let rooms = self.rooms.read().await;
        let Some(room) = rooms.get(&room_id) else {
            return;
        };

        for (peer_id, connections) in room.iter() {
            if *peer_id == user_id {
                continue;
            }
            for peer in connections.values() {
                let _ = peer.tx.send(ServerEvent::PlaySound {
                    user_id,
                    media_url: media_url.clone(),
                });
            }
        }
    }

    async fn emit_change(&self, room_id: Uuid, user_id: Uuid, change: VoiceStateChange) {
        if is_voice_state_change_empty(&change) {
            return;
        }

        let rooms = self.rooms.read().await;
        let Some(room) = rooms.get(&room_id) else {
            return;
        };

        if let Some(participant) = change.joined {
            broadcast_voice_event(
                room,
                user_id,
                ServerEvent::PeerJoined {
                    user_id: participant.user_id,
                    display_name: participant.display_name,
                },
            );
        }
        if let Some(left_user_id) = change.left_user_id {
            broadcast_voice_event(
                room,
                user_id,
                ServerEvent::PeerLeft {
                    user_id: left_user_id,
                },
            );
        }
        if let Some((changed_user_id, is_sharing_video)) = change.video_changed {
            broadcast_voice_event(
                room,
                user_id,
                ServerEvent::VideoStatusChanged {
                    user_id: changed_user_id,
                    is_sharing_video,
                },
            );
        }
        if let Some((changed_user_id, is_sharing_screen)) = change.screen_changed {
            broadcast_voice_event(
                room,
                user_id,
                ServerEvent::ScreenStatusChanged {
                    user_id: changed_user_id,
                    is_sharing_screen,
                },
            );
        }
        if let Some((changed_user_id, is_speaking)) = change.speaking_changed {
            broadcast_voice_event(
                room,
                user_id,
                ServerEvent::SpeakingStatusChanged {
                    user_id: changed_user_id,
                    is_speaking,
                },
            );
        }
        if let Some((changed_user_id, is_muted)) = change.mute_changed {
            broadcast_voice_event(
                room,
                user_id,
                ServerEvent::MuteStatusChanged {
                    user_id: changed_user_id,
                    is_muted,
                },
            );
        }
    }
}

pub async fn handle_socket(state: AppState, socket: WebSocket, room_id: Uuid, user_id: Uuid) {
    let Some(user) = db::get_user_by_id(&state.db, user_id).await.ok().flatten() else {
        return;
    };
    let Ok(Some(guild_id)) = db::guild_id_for_channel(&state.db, room_id).await else {
        return;
    };

    let (mut ws_sender, mut ws_receiver) = socket.split();
    let (tx, mut rx) = mpsc::unbounded_channel::<ServerEvent>();
    let connection_id = Uuid::new_v4();

    let (peers, change) = state
        .voice
        .join(
            room_id,
            user_id,
            connection_id,
            user.display_name.clone(),
            tx.clone(),
        )
        .await;
    let _ = tx.send(ServerEvent::Peers { peers });
    state.voice.emit_change(room_id, user_id, change).await;

    let send_task = tokio::spawn(async move {
        while let Some(event) = rx.recv().await {
            let Ok(payload) = serde_json::to_string(&event) else {
                continue;
            };
            if ws_sender.send(Message::Text(payload.into())).await.is_err() {
                break;
            }
        }
    });

    while let Some(Ok(msg)) = ws_receiver.next().await {
        match msg {
            Message::Text(text) => {
                let parsed = serde_json::from_str::<ClientEvent>(&text);
                match parsed {
                    Ok(ClientEvent::Signal {
                        to_user_id,
                        kind,
                        data,
                    }) => {
                        state
                            .voice
                            .relay_signal(room_id, user_id, to_user_id, kind, data)
                            .await;
                    }
                    Ok(ClientEvent::SetVideoStatus { is_sharing_video }) => {
                        let change = state
                            .voice
                            .set_video_status(room_id, user_id, connection_id, is_sharing_video)
                            .await;
                        state.voice.emit_change(room_id, user_id, change).await;
                    }
                    Ok(ClientEvent::SetScreenStatus { is_sharing_screen }) => {
                        let change = state
                            .voice
                            .set_screen_status(room_id, user_id, connection_id, is_sharing_screen)
                            .await;
                        state.voice.emit_change(room_id, user_id, change).await;
                    }
                    Ok(ClientEvent::SetSpeakingStatus { is_speaking }) => {
                        let change = state
                            .voice
                            .set_speaking_status(room_id, user_id, connection_id, is_speaking)
                            .await;
                        state.voice.emit_change(room_id, user_id, change).await;
                    }
                    Ok(ClientEvent::SetMuteStatus { is_muted }) => {
                        let change = state
                            .voice
                            .set_mute_status(room_id, user_id, connection_id, is_muted)
                            .await;
                        state.voice.emit_change(room_id, user_id, change).await;
                    }
                    Ok(ClientEvent::PlaySound { sound_id }) => {
                        if let Ok(Some(sound)) = db::get_sound_by_id(&state.db, sound_id).await
                            && sound.guild_id == guild_id
                        {
                            state
                                .voice
                                .play_sound(room_id, user_id, sound.media_url)
                                .await;
                        }
                    }
                    Err(err) => {
                        let _ = tx.send(ServerEvent::Error {
                            message: format!("invalid voice message: {err}"),
                        });
                    }
                }
            }
            Message::Close(_) => break,
            _ => {}
        }
    }

    send_task.abort();
    let change = state.voice.leave(room_id, user_id, connection_id).await;
    state.voice.emit_change(room_id, user_id, change).await;
}

fn aggregate_participant(
    connections: Option<&HashMap<Uuid, ClientHandle>>,
    user_id: Uuid,
) -> Option<VoiceParticipant> {
    let connections = connections?;
    if connections.is_empty() {
        return None;
    }

    let mut handles = connections.values();
    let first = handles.next()?;
    Some(VoiceParticipant {
        user_id,
        display_name: first.display_name.clone(),
        is_sharing_video: connections.values().any(|handle| handle.is_sharing_video),
        is_sharing_screen: connections.values().any(|handle| handle.is_sharing_screen),
        is_speaking: connections.values().any(|handle| handle.is_speaking),
        is_muted: connections.values().all(|handle| handle.is_muted),
    })
}

fn diff_voice_state(
    previous: Option<VoiceParticipant>,
    current: Option<VoiceParticipant>,
) -> VoiceStateChange {
    match (previous, current) {
        (None, None) => VoiceStateChange::default(),
        (None, Some(current)) => VoiceStateChange {
            joined: Some(current),
            ..Default::default()
        },
        (Some(previous), None) => VoiceStateChange {
            left_user_id: Some(previous.user_id),
            ..Default::default()
        },
        (Some(previous), Some(current)) => VoiceStateChange {
            video_changed: (previous.is_sharing_video != current.is_sharing_video)
                .then_some((current.user_id, current.is_sharing_video)),
            screen_changed: (previous.is_sharing_screen != current.is_sharing_screen)
                .then_some((current.user_id, current.is_sharing_screen)),
            speaking_changed: (previous.is_speaking != current.is_speaking)
                .then_some((current.user_id, current.is_speaking)),
            mute_changed: (previous.is_muted != current.is_muted)
                .then_some((current.user_id, current.is_muted)),
            ..Default::default()
        },
    }
}

fn broadcast_voice_event(
    room: &HashMap<Uuid, HashMap<Uuid, ClientHandle>>,
    source_user_id: Uuid,
    event: ServerEvent,
) {
    for (peer_id, connections) in room {
        if *peer_id == source_user_id {
            continue;
        }
        for peer in connections.values() {
            let _ = peer.tx.send(event.clone());
        }
    }
}

fn is_voice_state_change_empty(change: &VoiceStateChange) -> bool {
    change.joined.is_none()
        && change.left_user_id.is_none()
        && change.video_changed.is_none()
        && change.screen_changed.is_none()
        && change.speaking_changed.is_none()
        && change.mute_changed.is_none()
}

#[cfg(test)]
mod tests {
    use super::{VoiceHub, aggregate_participant};
    use std::collections::HashMap;
    use tokio::sync::mpsc;
    use uuid::Uuid;

    #[tokio::test]
    async fn multiple_connections_keep_voice_participant_until_last_leave() {
        let hub = VoiceHub::default();
        let room_id = Uuid::new_v4();
        let user_id = Uuid::new_v4();
        let first_connection = Uuid::new_v4();
        let second_connection = Uuid::new_v4();
        let (tx1, _rx1) = mpsc::unbounded_channel();
        let (tx2, _rx2) = mpsc::unbounded_channel();

        let (_peers, first_join) = hub
            .join(room_id, user_id, first_connection, "User".to_string(), tx1)
            .await;
        let (_peers, second_join) = hub
            .join(room_id, user_id, second_connection, "User".to_string(), tx2)
            .await;
        let first_leave = hub.leave(room_id, user_id, first_connection).await;
        let second_leave = hub.leave(room_id, user_id, second_connection).await;

        assert!(first_join.joined.is_some());
        assert!(second_join.joined.is_none());
        assert!(first_leave.left_user_id.is_none());
        assert_eq!(second_leave.left_user_id, Some(user_id));
    }

    #[tokio::test]
    async fn voice_mute_state_only_flips_when_all_connections_are_muted() {
        let hub = VoiceHub::default();
        let room_id = Uuid::new_v4();
        let user_id = Uuid::new_v4();
        let first_connection = Uuid::new_v4();
        let second_connection = Uuid::new_v4();
        let (tx1, _rx1) = mpsc::unbounded_channel();
        let (tx2, _rx2) = mpsc::unbounded_channel();

        let _ = hub
            .join(room_id, user_id, first_connection, "User".to_string(), tx1)
            .await;
        let _ = hub
            .join(room_id, user_id, second_connection, "User".to_string(), tx2)
            .await;

        let first_mute = hub
            .set_mute_status(room_id, user_id, first_connection, true)
            .await;
        let second_mute = hub
            .set_mute_status(room_id, user_id, second_connection, true)
            .await;
        let unmute = hub
            .set_mute_status(room_id, user_id, first_connection, false)
            .await;

        assert!(first_mute.mute_changed.is_none());
        assert_eq!(second_mute.mute_changed, Some((user_id, true)));
        assert_eq!(unmute.mute_changed, Some((user_id, false)));
    }

    #[test]
    fn aggregate_participant_combines_connection_state() {
        let user_id = Uuid::new_v4();
        let first_connection = Uuid::new_v4();
        let second_connection = Uuid::new_v4();
        let mut connections = HashMap::new();
        let (tx1, _rx1) = mpsc::unbounded_channel();
        let (tx2, _rx2) = mpsc::unbounded_channel();

        connections.insert(
            first_connection,
            super::ClientHandle {
                display_name: "User".to_string(),
                is_sharing_video: true,
                is_sharing_screen: false,
                is_speaking: false,
                is_muted: true,
                tx: tx1,
            },
        );
        connections.insert(
            second_connection,
            super::ClientHandle {
                display_name: "User".to_string(),
                is_sharing_video: false,
                is_sharing_screen: true,
                is_speaking: true,
                is_muted: false,
                tx: tx2,
            },
        );

        let participant = aggregate_participant(Some(&connections), user_id).unwrap();
        assert!(participant.is_sharing_video);
        assert!(participant.is_sharing_screen);
        assert!(participant.is_speaking);
        assert!(!participant.is_muted);
    }
}