1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
use axum::extract::ws::Message;
use futures_util::{future::join_all, stream::SplitSink, SinkExt};
use std::{sync::Arc, vec};
use tokio::{sync::Mutex, sync::RwLock};
use tracing::{debug, error};

use crate::relay::appstate::AppState;
use crate::relay::room::Room;
use crate::relay::RequestPacket;
use crate::relay::ResponsePacket;
use uuid::Uuid;

/// Type alias for a synchronized WebSocket sender.
/// 
/// This is used to send messages to a WebSocket connection.
type Sender = Arc<Mutex<SplitSink<axum::extract::ws::WebSocket, Message>>>;

/// Struct representing a WebSocket client.
/// 
/// This struct contains a WebSocket sender and an optional room ID.
/// The sender is used to send messages to the WebSocket connection,
/// while the room ID is used to identify the client's room.
#[derive(Debug)]
pub struct Client {
    /// The WebSocket sender for sending messages.
    sender: Sender,
    /// The optional room ID of the client.
    /// 
    /// This is used to identify the client's room.
    room_id: Option<String>,
}

impl Client {
    /// Creates a new WebSocket client.
    ///
    /// # Arguments
    ///
    /// * `sender` - A synchronized WebSocket sender.
    ///
    /// # Returns
    ///
    /// A new WebSocket client instance.
    pub fn new(sender: Sender) -> Client {
        Client {
            sender, // The WebSocket sender for sending messages.
            room_id: None, // The optional room ID of the client. This is used to identify the client's room.
        }
    }

    /// Sends a message to the WebSocket connection.
    ///
    /// # Arguments
    ///
    /// * `sender` - A synchronized WebSocket sender.
    /// * `message` - The message to send.
    ///
    /// # Errors
    ///
    /// If the message fails to be sent.
    async fn send(&self, sender: Sender, message: Message) {
        let mut sender = sender.lock().await; // Acquires a lock on the sender.
        if let Err(error) = sender.send(message).await { // Sends the message.
            error!("Failed to send message to the client: {}", error); // Logs the error if the message fails to be sent.
        }
    }

    /// Sends a serialized packet to the WebSocket connection.
    ///
    /// # Arguments
    ///
    /// * `sender` - A synchronized WebSocket sender.
    /// * `packet` - The packet to send.
    ///
    /// # Errors
    ///
    /// If the serialized packet fails to be sent.
    async fn send_packet(&self, sender: Sender, packet: ResponsePacket) {
        // Serialize the packet to a string.
        let serialized_packet = serde_json::to_string(&packet).unwrap();

        // Send the serialized packet as a text message.
        self.send(sender, Message::Text(serialized_packet)).await;
    }

    /// Sends an error message to the WebSocket connection.
    ///
    /// # Arguments
    ///
    /// * `sender` - A synchronized WebSocket sender.
    /// * `message` - The error message to send.
    ///
    /// # Errors
    ///
    /// If the error message fails to be sent.
    async fn send_error_packet(&self, sender: Sender, message: String) {
        // Create an error packet with the given message.
        let error_packet = ResponsePacket::Error { message };

        // Send the error packet.
        self.send_packet(sender, error_packet).await;
    }

    /// Handles the "create_room" request from a client.
    ///
    /// # Arguments
    ///
    /// * `server` - A lock guard of the `AppState`.
    /// * `id` - An optional string representing the room identifier.
    ///
    /// # Errors
    ///
    /// If the room already exists or if the room creation fails.
    async fn handle_create_room(&mut self, server: &RwLock<AppState>, id: Option<String>) {
        // Acquire a write lock on the server state.
        let mut server = server.write().await;

        // Check if the client is already in a room.
        if server.rooms.iter().any(|(_, room)| {
            room.senders
                .iter()
                .any(|sender| Arc::ptr_eq(sender, &self.sender))
        }) {
            return;
        }

        // Set the room size and generate a room identifier if none is provided.
        let size = Room::DEFAULT_ROOM_SIZE;
        let room_id = match id {
            Some(id) => id,
            None => Uuid::new_v4().to_string(),
        };

        // Check if the room identifier already exists.
        if server.rooms.contains_key(&room_id) {
            drop(server); // Release the lock before returning.

            return self
                .send_error_packet(
                    self.sender.clone(),
                    "A room with that identifier already exists.".to_string(),
                )
                .await;
        }

        // Create a new room and add the client to it.
        let mut room = Room::new(size);
        room.senders.push(self.sender.clone());

        // Insert the room into the server state.
        server.rooms.insert(room_id.clone(), room);

        self.room_id = Some(room_id.clone()); // Store the room identifier.

        drop(server); // Release the lock before returning.

        debug!("Room created");
        // Send the response packet to the client.
        self.send_packet(self.sender.clone(), ResponsePacket::Create { id: room_id })
            .await
    }

    /// Handles the "join_room" request from a client.
    ///
    /// # Arguments
    ///
    /// * `server` - A lock guard of the `AppState`.
    /// * `room_id` - A string representing the room identifier.
    ///
    /// # Errors
    ///
    /// If the room does not exist or if the room is full.
    async fn handle_join_room(&mut self, server: &RwLock<AppState>, room_id: String) {
        let mut server = server.write().await;

        // Check if the client is already in a room.
        if server.rooms.iter().any(|(_, room)| {
            room.senders
                .iter()
                .any(|sender| Arc::ptr_eq(sender, &self.sender))
        }) {
            return;
        }

        let Some(room) = server.rooms.get_mut(&room_id) else {
            drop(server);

            // Send an error packet to the client.
            return self
                .send_error_packet(self.sender.clone(), "The room does not exist.".to_string())
                .await;
        };

        // Check if the room is full.
        if room.senders.len() >= room.size {
            drop(server);

            // Send an error packet to the client.
            return self
                .send_error_packet(self.sender.clone(), "The room is full.".to_string())
                .await;
        }

        // Add the client to the room.
        room.senders.push(self.sender.clone());
        self.room_id = Some(room_id);

        let mut futures = vec![];
        for sender in &room.senders {
            // Send a join packet to the client with its position in the room.
            if Arc::ptr_eq(sender, &self.sender) {
                futures.push(self.send_packet(
                    sender.clone(),
                    ResponsePacket::Join {
                        size: Some(room.senders.len() - 1),
                    },
                ));
            } else {
                // Send a join packet to the client without its position in the room.
                futures.push(self.send_packet(sender.clone(), ResponsePacket::Join { size: None }));
            }
        }

        drop(server);
        join_all(futures).await;
    }

    /// Handle the leave room request from the client.
    ///
    /// This function removes the client from the current room and notifies the other
    /// clients in the room about the client's departure.
    ///
    /// # Arguments
    ///
    /// * `server` - A read-write lock guard for the server state.
    ///
    /// # Returns
    ///
    /// This function does not return anything.
    #[allow(clippy::needless_pass_by_value)]
    async fn handle_leave_room(&mut self, server: &RwLock<AppState>) {
        // Acquire a write lock on the server state.
        let mut server = server.write().await;

        // Get the room ID of the current room.
        let Some(room_id) = self.room_id.clone() else {
            return;
        };

        // Get the mutable reference to the room.
        let Some(room) = server.rooms.get_mut(&room_id) else {
            return;
        };

        // Get the index of the client in the room.
        let Some(index) = room
            .senders
            .iter()
            .position(|sender| Arc::ptr_eq(sender, &self.sender))
        else {
            return;
        };

        // Remove the client from the room.
        room.senders.remove(index);

        self.room_id = None;

        let mut futures = vec![];
        for sender in &room.senders {
            // Send a leave packet to the other clients in the room.
            futures.push(self.send_packet(sender.clone(), ResponsePacket::Leave { index }));
        }

        // If the room is empty, remove it from the server state.
        if room.senders.is_empty() {
            server.rooms.remove(&room_id);
        }

        drop(server);

        // Wait for all the futures to complete.
        join_all(futures).await;
    }

    /// Handles incoming messages from the client.
    ///
    /// This function interprets the incoming message and performs the corresponding action.
    ///
    /// # Arguments
    ///
    /// * `server` - A RwLock guard containing the state of the server.
    /// * `message` - The incoming message from the client.
    pub async fn handle_message(&mut self, server: &RwLock<AppState>, message: Message) {
        // Match on the type of the message.
        match message {
            // If the message is text, parse it as a RequestPacket.
            Message::Text(text) => {
                let packet = match serde_json::from_str(&text) {
                    Ok(packet) => packet,
                    Err(_) => return, // Return if the parsing fails.
                };
                // Match on the RequestPacket type and perform the corresponding action.
                match packet {
                    RequestPacket::Create { id } => self.handle_create_room(server, id).await,
                    RequestPacket::Join { id } => self.handle_join_room(server, id).await,
                    RequestPacket::Leave => self.handle_leave_room(server).await,
                }
            }
            // If the message is binary, handle it accordingly.
            Message::Binary(_) => {
                // Acquire a read lock on the server state.
                let server = server.read().await;

                // Get the room ID of the current room.
                let Some(room_id) = &self.room_id else {
                    drop(server);
                    return; // Return if the client is not in a room.
                };

                // Get the room corresponding to the room ID.
                let Some(room) = server.rooms.get(room_id) else {
                    drop(server);
                    return; // Return if the room does not exist.
                };

                // Get the index of the client in the room.
                let Some(index) = room
                    .senders
                    .iter()
                    .position(|sender| Arc::ptr_eq(sender, &self.sender))
                else {
                    drop(server);
                    return; // Return if the client is not in the room.
                };

                // Get the binary data from the message.
                let mut data = message.into_data();
                if data.is_empty() {
                    drop(server);
                    return; // Return if the data is empty.
                }

                // Convert the index to a u8 and assign it as the source.
                let source = u8::try_from(index).unwrap();

                // Get the destination from the first byte of the data.
                let destination = usize::from(data[0]);
                data[0] = source; // Assign the source to the first byte of the data.

                // If the destination is within the range of the room senders, send the data to that sender.
                if destination < room.senders.len() {
                    let sender = room.senders[destination].clone();

                    drop(server);
                    return self.send(sender, Message::Binary(data)).await;
                }

                // If the destination is u8::MAX, send the data to all the room senders except the current one.
                if destination == usize::from(u8::MAX) {
                    let mut futures = vec![];
                    for sender in &room.senders {
                        if Arc::ptr_eq(sender, &self.sender) {
                            continue; // Skip the current client.
                        }

                        futures.push(self.send(sender.clone(), Message::Binary(data.clone())));
                    }

                    drop(server);
                    join_all(futures).await;
                }
            }
            // If the message is Ping, print a message.
            Message::Ping(_) => {
                println!("Got Message Type Ping");
            }
            // If the message is Pong, print a message.
            Message::Pong(_) => {
                println!("Got Message Type Pong");
            }
            // If the message is Close, print a message and handle the close.
            Message::Close(_) => {
                println!("Got Message Type Close");
                self.handle_close(server).await;
            }
        }
    }

    pub async fn handle_close(&mut self, server: &RwLock<AppState>) {
        self.handle_leave_room(server).await
    }
}
// TODO: Add tests
#[cfg(test)]
mod tests {
    // use super::*;
}