//! WebSocket 客户端
//! 用于主动连接其他设备
//! 
//! 注意: 客户端使用独立的连接管理，因为 TLS 流类型与服务端不同

use dashmap::DashMap;
use futures_util::stream::SplitSink;
use futures_util::{SinkExt, StreamExt};
use std::sync::Arc;
use tokio::net::TcpStream;
use tokio_rustls::client::TlsStream;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{client_async, WebSocketStream};

use crate::models::{ChatMessage, DeviceInfo};
use crate::tls::client as tls_client;
use crate::utils::AppError;

use super::{ConnectionManager, MessageHandler};

/// 客户端 WebSocket 写入端类型
type ClientWsWriter = SplitSink<WebSocketStream<TlsStream<TcpStream>>, Message>;
type ClientWsWriterArc = Arc<tokio::sync::Mutex<ClientWsWriter>>;

/// WebSocket 客户端
#[derive(Clone)]
pub struct WsClient {
    /// 本机设备 ID
    local_device_id: String,
    /// 连接管理器 (用于事件通知)
    connection_manager: ConnectionManager,
    /// 客户端连接存储
    client_connections: Arc<DashMap<String, ClientWsWriterArc>>,
}

impl WsClient {
    /// 创建新的客户端
    pub fn new(local_device_id: String, connection_manager: ConnectionManager) -> Self {
        Self {
            local_device_id,
            connection_manager,
            client_connections: Arc::new(DashMap::new()),
        }
    }

    /// 连接到指定设备
    pub async fn connect(&self, device: &DeviceInfo) -> Result<(), AppError> {
        let addr = format!("{}:{}", device.ip, device.ws_port);
        let ws_url = format!("wss://{}", addr);

        log::info!("Connecting to {} at {}", device.device_name, ws_url);

        // 建立 TCP 连接
        let tcp_stream = TcpStream::connect(&addr)
            .await
            .map_err(|e| AppError::Connection(format!("TCP connection failed: {}", e)))?;

        // TLS 握手
        let tls_connector = tls_client::build_connector()?;
        let server_name = rustls::ServerName::try_from(device.ip.as_str())
            .map_err(|_| AppError::Tls("Invalid server name".to_string()))?;
        
        let tls_stream = tls_connector
            .connect(server_name, tcp_stream)
            .await
            .map_err(|e| AppError::Tls(format!("TLS handshake failed: {}", e)))?;

        // WebSocket 握手
        let (ws_stream, _) = client_async(&ws_url, tls_stream)
            .await
            .map_err(|e| AppError::WebSocket(format!("WebSocket handshake failed: {}", e)))?;

        let (writer, mut reader) = ws_stream.split();
        let writer_arc = Arc::new(tokio::sync::Mutex::new(writer));

        // 存储连接
        let device_id = device.device_id.clone();
        self.client_connections.insert(device_id.clone(), writer_arc.clone());

        // 发送初始消息以标识自己
        let hello = ChatMessage::new_event(
            self.local_device_id.clone(),
            device.device_id.clone(),
            "online",
        );
        
        {
            let data = serde_json::to_string(&hello)
                .map_err(|e| AppError::Serialization(e))?;
            let mut w = writer_arc.lock().await;
            w.send(Message::Text(data)).await
                .map_err(|e| AppError::WebSocket(format!("Send failed: {}", e)))?;
        }

        log::info!("Connected to device {}", device.device_name);

        // 启动读取任务
        let connection_manager = self.connection_manager.clone();
        let local_device_id = self.local_device_id.clone();
        let peer_device_id = device.device_id.clone();
        let client_connections = self.client_connections.clone();

        tokio::spawn(async move {
            while let Some(msg_result) = reader.next().await {
                match msg_result {
                    Ok(msg) => {
                        if msg.is_text() {
                            let text = msg.to_text().unwrap_or("");
                            if let Ok(message) = serde_json::from_str::<ChatMessage>(text) {
                                MessageHandler::handle_message(
                                    message,
                                    &connection_manager,
                                    &local_device_id,
                                ).await;
                            }
                        } else if msg.is_close() {
                            break;
                        }
                    }
                    Err(e) => {
                        log::error!("WebSocket read error: {}", e);
                        break;
                    }
                }
            }

            client_connections.remove(&peer_device_id);
            log::info!("Disconnected from device {}", peer_device_id);
        });

        Ok(())
    }

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

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

    /// 断开与指定设备的连接
    pub async fn disconnect(&self, device_id: &str) {
        // 发送离线通知
        let goodbye = ChatMessage::new_event(
            self.local_device_id.clone(),
            device_id.to_string(),
            "offline",
        );
        
        let _ = self.send_to(device_id, &goodbye).await;
        self.client_connections.remove(device_id);
    }

    /// 获取连接管理器引用
    pub fn connection_manager(&self) -> &ConnectionManager {
        &self.connection_manager
    }
}
