signer-remote 0.4.1

use futures::{SinkExt, StreamExt};
use serde_json;
use signer_core::{SignerCrypted, SignerUser};
use std::sync::Arc;
use tokio::sync::{broadcast, Mutex}; // Added broadcast
use tokio_tungstenite::{connect_async, tungstenite::protocol::Message, WebSocketStream};

use super::signer_transfer_message::{ChannelDevice, ClientToTransferMessage, TransferToClientMessage}; // Import ChannelDevice
use crate::{
    error::{RemoteError, RemoteResult},
    SignerSummary,
};

// --- New Event Enum ---
#[derive(Clone, Debug, serde::Serialize)]
#[serde(tag = "type", content = "data")]
pub enum TransferEvent {
    ChannelJoined {
        channel_id: String,
        devices: Vec<ChannelDevice>,
    },
    DeviceListUpdate {
        devices: Vec<ChannelDevice>,
    },
    TransferRequestReceived {
        from_device_id: String,
        encrypted_summary: SignerCrypted<SignerSummary>,
    },
    TransferStarted {
        target_device_id: String,
    },
    TransferCompleted {
        receiver_device_id: String,
    },
    ChannelClosed {
        reason: String,
    },
    ConnectionClosed,
    Error(String),
}

/// `SignerTransferClient` 结构体，用于与传输服务器进行 WebSocket 通信。
pub struct SignerTransferClient {
    ws_sender: Option<Arc<Mutex<futures::stream::SplitSink<WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>, Message>>>>,
    pub channel_id: Arc<Mutex<Option<String>>>,
    // The broadcast sender for events.
    event_sender: broadcast::Sender<TransferEvent>,
    state_mutex: Arc<Mutex<()>>,
}

impl SignerTransferClient {
    /// 创建一个新的 `SignerTransferClient` 实例。
    pub fn new() -> Self {
        let (event_sender, _) = broadcast::channel(32); // Create the channel
        Self {
            ws_sender: None,
            channel_id: Arc::new(Mutex::new(None)),
            event_sender, // Store the sender
            state_mutex: Arc::new(Mutex::new(())),
        }
    }

    /// 订阅事件流
    pub fn events(&self) -> broadcast::Receiver<TransferEvent> {
        self.event_sender.subscribe()
    }

    /// 连接到传输服务器的 WebSocket 服务器并开始监听。
    /// `server_url` 是服务器的基础 URL，例如 "ws://localhost:3000/api/transfer"。
    pub async fn connect_and_listen(&mut self, server_url: &str) -> RemoteResult<()> {
        let _guard = self.state_mutex.lock().await;

        let (ws_stream, _response) = connect_async(server_url)
            .await
            .map_err(|e| RemoteError::NetworkError(format!("连接到传输服务器失败: {}", e)))?;

        let (sink, stream) = ws_stream.split();
        
        self.ws_sender = Some(Arc::new(Mutex::new(sink)));
        
        println!("已连接到传输服务器: {}", server_url);

        // Automatically start the receiving task
        self.start_receiving_task(stream).await;

        Ok(())
    }

    /// 加入一个传输 channel。
    pub async fn join_channel(
        &self, // Changed to &self since it doesn't modify the client itself directly
        channel_id: Option<String>,
        user_info: SignerUser,
    ) -> RemoteResult<()> {
        let message = ClientToTransferMessage::JoinChannel {
            channel_id,
            user_info,
        };
        self.send_message(&message).await?;
        println!("已发送加入 Channel 请求");
        Ok(())
    }

    /// 发送 AcceptTransfer 消息
    pub async fn send_transfer_request(
        &self, // Changed to &self
        target_device_id: String,
        encrypted_summary: SignerCrypted<SignerSummary>,
    ) -> RemoteResult<()> {
        let channel_id_guard = self.channel_id.lock().await;
        if channel_id_guard.is_none() {
            return Err(RemoteError::InvalidOperation(
                "未加入任何 channel，无法发送传输请求".to_string(),
            ));
        }
        drop(channel_id_guard);

        let message = ClientToTransferMessage::AcceptTransfer {
            target_device_id,
            encrypted_summary,
        };
        self.send_message(&message).await?;
        Ok(())
    }

    /// 发送传输完成确认消息
    pub async fn send_transfer_completed(
        &self,
        sender_device_id: String,
    ) -> RemoteResult<()> {
        let channel_id_guard = self.channel_id.lock().await;
        if channel_id_guard.is_none() {
            return Err(RemoteError::InvalidOperation(
                "未加入任何 channel，无法发送传输完成确认".to_string(),
            ));
        }
        drop(channel_id_guard);

        let message = ClientToTransferMessage::TransferCompleted {
            sender_device_id,
        };
        self.send_message(&message).await?;
        println!("已发送传输完成确认消息");
        Ok(())
    }

    /// 内部方法：发送消息到 WebSocket 服务器
    async fn send_message(&self, msg: &ClientToTransferMessage) -> RemoteResult<()> {
        let ws_sender = self
            .ws_sender
            .as_ref()
            .ok_or(RemoteError::NetworkError(
                "WebSocket 未连接".to_string(),
            ))?
            .clone();

        let msg_str = serde_json::to_string(msg)
            .map_err(|e| RemoteError::SerializationError(e.to_string()))?;

        let mut sender = ws_sender.lock().await;
        sender
            .send(Message::Text(msg_str.into()))
            .await
            .map_err(|e| RemoteError::NetworkError(format!("发送消息失败: {}", e)))?;
        Ok(())
    }

    /// 启动消息接收任务
    async fn start_receiving_task(&self, mut stream: futures::stream::SplitStream<WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>>) {
        let channel_id = self.channel_id.clone();
        let event_sender = self.event_sender.clone();
        
        tokio::spawn(async move {
            while let Some(msg) = stream.next().await {
                match msg {
                    Ok(Message::Text(text)) => {
                        match serde_json::from_str::<TransferToClientMessage>(&text) {
                            Ok(transfer_msg) => {
                                Self::handle_transfer_message(transfer_msg, &channel_id, &event_sender).await;
                            }
                            Err(e) => {
                                eprintln!("解析消息失败: {}", e);
                                let _ = event_sender.send(TransferEvent::Error(format!("解析消息失败: {}", e)));
                            }
                        }
                    }
                    Ok(Message::Close(_)) => {
                        println!("连接已关闭");
                        let _ = event_sender.send(TransferEvent::ConnectionClosed);
                        break;
                    }
                    Ok(_) => {}
                    Err(e) => {
                        eprintln!("接收消息时出错: {}", e);
                        let _ = event_sender.send(TransferEvent::Error(format!("接收消息时出错: {}", e)));
                        break;
                    }
                }
            }
        });
    }

    /// 处理来自传输服务器的消息
    async fn handle_transfer_message(
        msg: TransferToClientMessage,
        channel_id_arc: &Arc<Mutex<Option<String>>>,
        event_sender: &broadcast::Sender<TransferEvent>,
    ) {
        let event = match msg {
            TransferToClientMessage::ChannelJoined {
                channel_id,
                devices,
            } => {
                println!("成功加入 Channel: {}", channel_id);
                let mut channel_id_guard = channel_id_arc.lock().await;
                *channel_id_guard = Some(channel_id.clone());
                Some(TransferEvent::ChannelJoined { channel_id, devices })
            }
            TransferToClientMessage::DeviceListUpdate { devices } => {
                println!("收到设备列表更新");
                Some(TransferEvent::DeviceListUpdate { devices })
            }
            TransferToClientMessage::TransferStarted { target_device_id } => {
                println!("传输已开始到设备: {}", target_device_id);
                Some(TransferEvent::TransferStarted { target_device_id })
            }
            TransferToClientMessage::ReceiveTransfer {
                from_device_id,
                encrypted_summary,
            } => {
                Some(TransferEvent::TransferRequestReceived { from_device_id, encrypted_summary })
            }
            TransferToClientMessage::TransferCompleted {
                receiver_device_id,
            } => {
                println!("收到传输完成确认，来自设备: {}", receiver_device_id);
                Some(TransferEvent::TransferCompleted { receiver_device_id })
            }
            TransferToClientMessage::ChannelClosed { reason } => {
                println!("Channel 已关闭: {}", reason);
                Some(TransferEvent::ChannelClosed { reason })
            }
            TransferToClientMessage::Error { message } => {
                eprintln!("服务器错误: {}", message);
                Some(TransferEvent::Error(message))
            }
        };

        if let Some(event) = event {
            if let Err(e) = event_sender.send(event) {
                eprintln!("Failed to broadcast event: {}", e);
            }
        }
    }
}

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