//! WebSocket 连接管理

use dashmap::DashMap;
use futures_util::stream::SplitSink;
use futures_util::SinkExt;
use std::sync::Arc;
use tokio::net::TcpStream;
use tokio::sync::broadcast;
use tokio_rustls::server::TlsStream;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::WebSocketStream;

use crate::models::ChatMessage;

/// WebSocket 连接事件
#[derive(Debug, Clone)]
pub enum WsConnectionEvent {
    /// 新连接建立
    Connected { device_id: String },
    /// 连接断开
    Disconnected { device_id: String, reason: Option<String> },
    /// 收到消息
    MessageReceived { message: ChatMessage },
}

/// WebSocket 写入端类型别名 (服务端TLS连接)
pub type WsWriter = SplitSink<WebSocketStream<TlsStream<TcpStream>>, Message>;

/// 包装后的写入端类型
pub type WsWriterArc = Arc<tokio::sync::Mutex<WsWriter>>;

/// 连接管理器
pub struct ConnectionManager {
    /// 活跃连接 (device_id -> writer)
    connections: Arc<DashMap<String, WsWriterArc>>,
    /// 事件发送器
    event_tx: broadcast::Sender<WsConnectionEvent>,
}

impl ConnectionManager {
    /// 创建新的连接管理器
    pub fn new() -> Self {
        let (event_tx, _) = broadcast::channel(100);
        Self {
            connections: Arc::new(DashMap::new()),
            event_tx,
        }
    }

    /// 获取事件订阅
    pub fn subscribe(&self) -> broadcast::Receiver<WsConnectionEvent> {
        self.event_tx.subscribe()
    }

    /// 注册新连接
    pub fn register(&self, device_id: String, writer: WsWriter) {
        let writer = Arc::new(tokio::sync::Mutex::new(writer));
        self.connections.insert(device_id.clone(), writer);
        let _ = self.event_tx.send(WsConnectionEvent::Connected { device_id });
    }

    /// 注册新连接 (使用已包装的 Arc)
    pub fn register_arc(&self, device_id: String, writer: WsWriterArc) {
        self.connections.insert(device_id.clone(), writer);
        let _ = self.event_tx.send(WsConnectionEvent::Connected { device_id });
    }

    /// 移除连接
    pub fn unregister(&self, device_id: &str, reason: Option<String>) {
        if self.connections.remove(device_id).is_some() {
            let _ = self.event_tx.send(WsConnectionEvent::Disconnected {
                device_id: device_id.to_string(),
                reason,
            });
        }
    }

    /// 检查设备是否已连接
    pub fn is_connected(&self, device_id: &str) -> bool {
        self.connections.contains_key(device_id)
    }

    /// 获取所有已连接的设备 ID
    pub fn connected_devices(&self) -> Vec<String> {
        self.connections.iter().map(|e| e.key().clone()).collect()
    }

    /// 向指定设备发送消息
    pub async fn send_to(&self, device_id: &str, message: &ChatMessage) -> Result<(), String> {
        if let Some(writer) = self.connections.get(device_id) {
            let data = serde_json::to_string(message)
                .map_err(|e| format!("Serialization error: {}", e))?;
            
            let mut writer = writer.lock().await;
            writer
                .send(Message::Text(data))
                .await
                .map_err(|e| format!("Send error: {}", e))?;
            Ok(())
        } else {
            Err(format!("Device {} not connected", device_id))
        }
    }

    /// 广播消息到所有连接
    pub async fn broadcast(&self, message: &ChatMessage) {
        let data = match serde_json::to_string(message) {
            Ok(d) => d,
            Err(e) => {
                log::error!("Failed to serialize message: {}", e);
                return;
            }
        };

        for entry in self.connections.iter() {
            let mut writer = entry.value().lock().await;
            if let Err(e) = writer.send(Message::Text(data.clone())).await {
                log::error!("Failed to send to {}: {}", entry.key(), e);
            }
        }
    }

    /// 发送消息接收事件
    pub fn emit_message_received(&self, message: ChatMessage) {
        let _ = self.event_tx.send(WsConnectionEvent::MessageReceived { message });
    }

    /// 获取连接数
    pub fn connection_count(&self) -> usize {
        self.connections.len()
    }

    /// 关闭所有连接
    pub async fn close_all(&self) {
        for entry in self.connections.iter() {
            let mut writer = entry.value().lock().await;
            let _ = writer.close().await;
        }
        self.connections.clear();
    }
}

impl Default for ConnectionManager {
    fn default() -> Self {
        Self::new()
    }
}

impl Clone for ConnectionManager {
    fn clone(&self) -> Self {
        Self {
            connections: self.connections.clone(),
            event_tx: self.event_tx.clone(),
        }
    }
}
