//! WebSocket 服务端

use futures_util::StreamExt;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::TcpListener;
use tokio_rustls::TlsAcceptor;
use tokio_tungstenite::accept_async;

use crate::models::ChatMessage;
use crate::tls::{CertificateManager, TlsConfigBuilder};
use crate::utils::AppError;

use super::{ConnectionManager, MessageHandler};

/// WebSocket 服务端
pub struct WsServer {
    /// 监听端口
    port: u16,
    /// 本机设备 ID
    local_device_id: String,
    /// 连接管理器
    connection_manager: ConnectionManager,
    /// 停止信号
    stop_tx: Option<tokio::sync::oneshot::Sender<()>>,
}

impl WsServer {
    /// 创建新的 WebSocket 服务端
    pub fn new(port: u16, local_device_id: String, connection_manager: ConnectionManager) -> Self {
        Self {
            port,
            local_device_id,
            connection_manager,
            stop_tx: None,
        }
    }

    /// 启动服务
    pub async fn start(&mut self) -> Result<(), AppError> {
        // 获取 TLS 配置
        let cert_bundle = CertificateManager::get()
            .ok_or_else(|| AppError::Tls("Certificate not initialized".to_string()))?;
        
        let tls_acceptor = TlsConfigBuilder::build_acceptor(&cert_bundle)?;

        // 绑定端口
        let addr = SocketAddr::from(([0, 0, 0, 0], self.port));
        let listener = TcpListener::bind(addr).await?;
        
        log::info!("WebSocket server started on port {}", self.port);

        // 创建停止信号
        let (stop_tx, mut stop_rx) = tokio::sync::oneshot::channel();
        self.stop_tx = Some(stop_tx);

        let connection_manager = self.connection_manager.clone();
        let local_device_id = self.local_device_id.clone();

        tokio::spawn(async move {
            loop {
                tokio::select! {
                    result = listener.accept() => {
                        match result {
                            Ok((stream, peer_addr)) => {
                                let tls_acceptor = tls_acceptor.clone();
                                let conn_mgr = connection_manager.clone();
                                let device_id = local_device_id.clone();
                                
                                tokio::spawn(async move {
                                    if let Err(e) = Self::handle_connection(
                                        stream,
                                        peer_addr,
                                        tls_acceptor,
                                        conn_mgr,
                                        device_id,
                                    ).await {
                                        log::error!("Connection error from {}: {}", peer_addr, e);
                                    }
                                });
                            }
                            Err(e) => {
                                log::error!("Accept error: {}", e);
                            }
                        }
                    }
                    _ = &mut stop_rx => {
                        log::info!("WebSocket server stopping");
                        break;
                    }
                }
            }
        });

        Ok(())
    }

    /// 处理单个连接
    async fn handle_connection(
        stream: tokio::net::TcpStream,
        peer_addr: SocketAddr,
        tls_acceptor: TlsAcceptor,
        connection_manager: ConnectionManager,
        local_device_id: String,
    ) -> Result<(), AppError> {
        log::info!("New connection from {}", peer_addr);

        // TLS 握手
        let tls_stream = tls_acceptor
            .accept(stream)
            .await
            .map_err(|e| AppError::Tls(format!("TLS handshake failed: {}", e)))?;

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

        let (write_half, mut read_half) = ws_stream.split();
        
        // 使用 Arc<Mutex> 包装 writer 以便在读取第一条消息后注册
        let writer = Arc::new(tokio::sync::Mutex::new(write_half));
        let mut peer_device_id: Option<String> = None;
        let mut registered = false;

        while let Some(msg_result) = read_half.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) {
                            if !registered {
                                // 第一条消息，注册连接
                                peer_device_id = Some(message.from.clone());
                                connection_manager.register_arc(message.from.clone(), writer.clone());
                                registered = true;
                                log::info!("Device {} connected", message.from);
                            }

                            // 处理消息
                            MessageHandler::handle_message(
                                message,
                                &connection_manager,
                                &local_device_id,
                            ).await;
                        }
                    } else if msg.is_close() {
                        break;
                    }
                }
                Err(e) => {
                    log::error!("WebSocket error: {}", e);
                    break;
                }
            }
        }

        // 连接关闭，移除注册
        if let Some(device_id) = peer_device_id {
            connection_manager.unregister(&device_id, Some("Connection closed".to_string()));
            log::info!("Device {} disconnected", device_id);
        }

        Ok(())
    }

    /// 停止服务
    pub fn stop(&mut self) {
        if let Some(stop_tx) = self.stop_tx.take() {
            let _ = stop_tx.send(());
        }
    }

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