// [TSM.ID].[11031972] -- All Rights Reserved. Proprietary & Confidential.
use futures_util::{SinkExt, StreamExt};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::net::TcpListener;
use tokio::sync::Mutex;
use tokio_tungstenite::accept_async;
use tokio_tungstenite::tungstenite::Message;
use tracing::{info, error};

type Clients = Arc<Mutex<HashMap<String, futures_util::stream::SplitSink<tokio_tungstenite::WebSocketStream<tokio::net::TcpStream>, Message>>>>;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt::init();
    
    let addr = std::env::var("NEURAL_CHAT_BIND").unwrap_or_else(|_| "0.0.0.0:8443".to_string());
    let listener = TcpListener::bind(&addr).await?;
    info!("XCU NEURAL MESH (MoQ) ignited on: {}", addr);
    info!("Bypassing DB. Operating strictly in memory via Zero-Copy channels.");

    let clients: Clients = Arc::new(Mutex::new(HashMap::new()));

    while let Ok((stream, peer)) = listener.accept().await {
        let clients = clients.clone();
        tokio::spawn(async move {
            info!("Incoming Neural Connection from: {}", peer);
            match accept_async(stream).await {
                Ok(ws_stream) => {
                    let (sender, mut receiver) = ws_stream.split();
                    let mut client_id = String::new();

                    while let Some(msg) = receiver.next().await {
                        match msg {
                            Ok(Message::Text(text)) => {
                                // Menerima payload tersandika Base64 dari native WASM
                                if let Ok(json) = serde_json::from_str::<serde_json::Value>(&text) {
                                    if let Some(c_id) = json["client_id"].as_str() {
                                        if client_id.is_empty() {
                                            client_id = c_id.to_string();
                                            let mut map = clients.lock().await;
                                            map.insert(client_id.clone(), sender);
                                            info!("Client Registered: {}", client_id);
                                            // Lanjutkan iterasi karena sender sudah berpindah kepemilikan
                                            break;
                                        }
                                    }
                                }
                            }
                            Ok(Message::Close(_)) => break,
                            Err(e) => {
                                error!("Websocket error: {}", e);
                                break;
                            }
                            _ => {}
                        }
                    }
                    
                    // Lanjutkan menerima pesan untuk di-broadcast
                    if !client_id.is_empty() {
                        while let Some(msg) = receiver.next().await {
                            match msg {
                                Ok(Message::Text(text)) => {
                                    if let Ok(json) = serde_json::from_str::<serde_json::Value>(&text) {
                                        let _event = json["event"].as_str().unwrap_or("");
                                        // Broadcast the event to all other clients
                                        let mut map = clients.lock().await;
                                        for (id, tx) in map.iter_mut() {
                                            if *id != client_id {
                                                let _ = tx.send(Message::Text(text.clone())).await;
                                            }
                                        }
                                    }
                                }
                                Ok(Message::Close(_)) => break,
                                Err(_) => break,
                                _ => {}
                            }
                        }
                        info!("Client Disconnected: {}", client_id);
                        let mut map = clients.lock().await;
                        map.remove(&client_id);
                    }
                }
                Err(e) => error!("Failed to accept connection: {}", e),
            }
        });
    }

    Ok(())
}
