//! WebSocket event channels for real-time updates.
//!
//! Provides real-time streaming of blockchain events:
//! - Block events (new blocks, confirmations)
//! - Transaction events (mempool, confirmations)
//! - Address events (balance changes, UTXOs)
//! - Contract events (logs, state changes)
//! - Market events (price updates, trades)

use std::collections::HashMap;
use std::sync::Arc;

use axum::{
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        Query, State,
    },
    response::IntoResponse,
    routing::get,
    Router,
};
use futures::{sink::SinkExt, stream::StreamExt};
use serde::{Deserialize, Serialize};
use tokio::sync::{broadcast, RwLock};
use tracing::{debug, error, info, warn};
use uuid::Uuid;

use crate::auth::{AuthContext, OptionalAuth};
use crate::routes::AppState;

// ============================================================================
// Event Types
// ============================================================================

/// Event categories for subscription.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum EventCategory {
    /// Block-related events
    Blocks,
    /// Transaction events
    Transactions,
    /// Address-specific events
    Addresses,
    /// Smart contract events
    Contracts,
    /// DEX market events
    Markets,
    /// Mining events
    Mining,
    /// Network/peer events
    Network,
}

/// Block event data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BlockEvent {
    /// Block hash
    pub hash: String,
    /// Block height
    pub height: u64,
    /// Timestamp
    pub timestamp: u64,
    /// Number of transactions
    pub tx_count: u32,
    /// Blue score (DAG metric)
    pub blue_score: u64,
    /// Parent hashes
    pub parent_hashes: Vec<String>,
}

/// Transaction event data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransactionEvent {
    /// Transaction hash
    pub txid: String,
    /// Event type (mempool, confirmed, orphaned)
    pub event_type: String,
    /// Block hash (if confirmed)
    pub block_hash: Option<String>,
    /// Confirmations
    pub confirmations: u32,
    /// Input addresses
    pub inputs: Vec<String>,
    /// Output addresses
    pub outputs: Vec<String>,
    /// Total value
    pub value: String,
}

/// Address event data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AddressEvent {
    /// Address
    pub address: String,
    /// Event type (received, sent, balance_change)
    pub event_type: String,
    /// Transaction hash
    pub txid: String,
    /// Amount changed
    pub amount: String,
    /// New balance
    pub balance: String,
}

/// Contract event data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContractEvent {
    /// Contract address
    pub contract_id: String,
    /// Event name
    pub event_name: String,
    /// Event topics
    pub topics: Vec<String>,
    /// Event data (hex)
    pub data: String,
    /// Transaction hash
    pub txid: String,
    /// Block hash
    pub block_hash: String,
    /// Log index
    pub log_index: u32,
}

/// Market event data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MarketEvent {
    /// Market symbol
    pub symbol: String,
    /// Event type (trade, orderbook_update, ticker)
    pub event_type: String,
    /// Price
    pub price: Option<String>,
    /// Volume
    pub volume: Option<String>,
    /// Additional data
    pub data: serde_json::Value,
}

/// Mining event data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MiningEvent {
    /// Event type (new_template, difficulty_change, hashrate_update)
    pub event_type: String,
    /// Current difficulty
    pub difficulty: Option<String>,
    /// Network hashrate
    pub hashrate: Option<String>,
    /// Block template hash
    pub template_hash: Option<String>,
}

/// Network event data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkEvent {
    /// Event type (peer_connected, peer_disconnected, sync_progress)
    pub event_type: String,
    /// Peer ID
    pub peer_id: Option<String>,
    /// Peer address
    pub peer_address: Option<String>,
    /// Connection count
    pub connections: Option<u32>,
}

/// Unified event wrapper.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "category", content = "data")]
#[serde(rename_all = "snake_case")]
pub enum Event {
    Block(BlockEvent),
    Transaction(TransactionEvent),
    Address(AddressEvent),
    Contract(ContractEvent),
    Market(MarketEvent),
    Mining(MiningEvent),
    Network(NetworkEvent),
}

/// WebSocket message from client.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ClientMessage {
    /// Subscribe to events
    Subscribe {
        /// Channel to subscribe to
        channel: String,
        /// Optional filter parameters
        #[serde(default)]
        params: HashMap<String, String>,
    },
    /// Unsubscribe from events
    Unsubscribe {
        /// Channel to unsubscribe from
        channel: String,
    },
    /// Ping message
    Ping {
        /// Ping ID for correlation
        id: Option<String>,
    },
    /// Request current state
    GetState {
        /// What state to fetch
        target: String,
    },
}

/// WebSocket message to client.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ServerMessage {
    /// Welcome message on connection
    Welcome {
        /// Session ID
        session_id: String,
        /// Server version
        version: String,
    },
    /// Subscription confirmed
    Subscribed {
        /// Channel name
        channel: String,
        /// Subscription ID
        subscription_id: String,
    },
    /// Unsubscription confirmed
    Unsubscribed {
        /// Channel name
        channel: String,
    },
    /// Event notification
    Event {
        /// Channel name
        channel: String,
        /// Event data
        event: Event,
        /// Timestamp
        timestamp: u64,
    },
    /// Pong response
    Pong {
        /// Ping ID for correlation
        id: Option<String>,
    },
    /// Error message
    Error {
        /// Error code
        code: String,
        /// Error message
        message: String,
    },
    /// State response
    State {
        /// Target
        target: String,
        /// State data
        data: serde_json::Value,
    },
}

// ============================================================================
// Subscription Management
// ============================================================================

/// Subscription entry.
#[derive(Debug, Clone)]
pub struct Subscription {
    /// Subscription ID
    pub id: String,
    /// Channel name
    pub channel: String,
    /// Filter parameters
    pub params: HashMap<String, String>,
}

/// Client session state.
#[derive(Debug)]
pub struct ClientSession {
    /// Session ID
    pub id: String,
    /// Active subscriptions
    pub subscriptions: HashMap<String, Subscription>,
    /// Authentication context (if authenticated)
    pub auth: Option<AuthContext>,
}

impl ClientSession {
    /// Create a new session.
    pub fn new(auth: Option<AuthContext>) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            subscriptions: HashMap::new(),
            auth,
        }
    }
}

// ============================================================================
// Event Broadcaster
// ============================================================================

/// Broadcast channel capacity.
const BROADCAST_CAPACITY: usize = 1024;

/// Event broadcaster for publishing events to subscribers.
#[derive(Clone)]
pub struct EventBroadcaster {
    /// Block events channel
    blocks: broadcast::Sender<Event>,
    /// Transaction events channel
    transactions: broadcast::Sender<Event>,
    /// Address events channel
    addresses: broadcast::Sender<Event>,
    /// Contract events channel
    contracts: broadcast::Sender<Event>,
    /// Market events channel
    markets: broadcast::Sender<Event>,
    /// Mining events channel
    mining: broadcast::Sender<Event>,
    /// Network events channel
    network: broadcast::Sender<Event>,
    /// Active connections counter
    connections: Arc<RwLock<u64>>,
}

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

impl EventBroadcaster {
    /// Create a new event broadcaster.
    pub fn new() -> Self {
        Self {
            blocks: broadcast::channel(BROADCAST_CAPACITY).0,
            transactions: broadcast::channel(BROADCAST_CAPACITY).0,
            addresses: broadcast::channel(BROADCAST_CAPACITY).0,
            contracts: broadcast::channel(BROADCAST_CAPACITY).0,
            markets: broadcast::channel(BROADCAST_CAPACITY).0,
            mining: broadcast::channel(BROADCAST_CAPACITY).0,
            network: broadcast::channel(BROADCAST_CAPACITY).0,
            connections: Arc::new(RwLock::new(0)),
        }
    }

    /// Publish an event.
    pub fn publish(&self, event: Event) {
        let sender = match &event {
            Event::Block(_) => &self.blocks,
            Event::Transaction(_) => &self.transactions,
            Event::Address(_) => &self.addresses,
            Event::Contract(_) => &self.contracts,
            Event::Market(_) => &self.markets,
            Event::Mining(_) => &self.mining,
            Event::Network(_) => &self.network,
        };

        // Don't error if no subscribers
        let _ = sender.send(event);
    }

    /// Subscribe to a category.
    pub fn subscribe(&self, category: EventCategory) -> broadcast::Receiver<Event> {
        match category {
            EventCategory::Blocks => self.blocks.subscribe(),
            EventCategory::Transactions => self.transactions.subscribe(),
            EventCategory::Addresses => self.addresses.subscribe(),
            EventCategory::Contracts => self.contracts.subscribe(),
            EventCategory::Markets => self.markets.subscribe(),
            EventCategory::Mining => self.mining.subscribe(),
            EventCategory::Network => self.network.subscribe(),
        }
    }

    /// Get active connection count.
    pub async fn connection_count(&self) -> u64 {
        *self.connections.read().await
    }

    /// Increment connection count.
    async fn add_connection(&self) {
        let mut count = self.connections.write().await;
        *count += 1;
    }

    /// Decrement connection count.
    async fn remove_connection(&self) {
        let mut count = self.connections.write().await;
        *count = count.saturating_sub(1);
    }
}

// ============================================================================
// WebSocket State
// ============================================================================

/// WebSocket module state.
#[derive(Clone)]
pub struct WebSocketState {
    /// Event broadcaster
    pub broadcaster: EventBroadcaster,
}

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

impl WebSocketState {
    /// Create new WebSocket state.
    pub fn new() -> Self {
        Self {
            broadcaster: EventBroadcaster::new(),
        }
    }
}

// ============================================================================
// Routes
// ============================================================================

/// Query parameters for WebSocket connection.
#[derive(Debug, Deserialize)]
pub struct WsQueryParams {
    /// Optional API key for authentication
    api_key: Option<String>,
    /// Initial subscriptions (comma-separated)
    subscribe: Option<String>,
}

/// Build WebSocket routes.
pub fn router() -> Router<AppState> {
    Router::new()
        .route("/ws", get(ws_handler))
        .route("/ws/blocks", get(ws_blocks_handler))
        .route("/ws/transactions", get(ws_transactions_handler))
        .route("/ws/address/:address", get(ws_address_handler))
        .route("/ws/contract/:contract_id", get(ws_contract_handler))
        .route("/ws/markets", get(ws_markets_handler))
}

/// Main WebSocket handler.
async fn ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    OptionalAuth(auth): OptionalAuth,
    Query(params): Query<WsQueryParams>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_socket(socket, state, auth, params))
}

/// Handle WebSocket connection.
async fn handle_socket(
    socket: WebSocket,
    state: AppState,
    auth: Option<AuthContext>,
    params: WsQueryParams,
) {
    let ws_state = &state.websocket;
    ws_state.broadcaster.add_connection().await;

    let session = ClientSession::new(auth);
    let session_id = session.id.clone();

    info!("WebSocket connected: {}", session_id);

    let (mut sender, mut receiver) = socket.split();

    // Send welcome message
    let welcome = ServerMessage::Welcome {
        session_id: session_id.clone(),
        version: "1.0.0".to_string(),
    };
    if let Ok(msg) = serde_json::to_string(&welcome) {
        let _ = sender.send(Message::Text(msg)).await;
    }

    // Handle initial subscriptions
    if let Some(subs) = params.subscribe {
        for channel in subs.split(',') {
            let subscribed = ServerMessage::Subscribed {
                channel: channel.trim().to_string(),
                subscription_id: Uuid::new_v4().to_string(),
            };
            if let Ok(msg) = serde_json::to_string(&subscribed) {
                let _ = sender.send(Message::Text(msg)).await;
            }
        }
    }

    // Create channels for shutdown coordination
    let (tx, mut rx) = tokio::sync::mpsc::channel::<ServerMessage>(100);

    // Spawn task to forward events to client
    let sender_handle = tokio::spawn(async move {
        loop {
            tokio::select! {
                Some(msg) = rx.recv() => {
                    if let Ok(text) = serde_json::to_string(&msg) {
                        if sender.send(Message::Text(text)).await.is_err() {
                            break;
                        }
                    }
                }
                else => break,
            }
        }
    });

    // Handle incoming messages
    while let Some(msg) = receiver.next().await {
        match msg {
            Ok(Message::Text(text)) => {
                if let Ok(client_msg) = serde_json::from_str::<ClientMessage>(&text) {
                    let response = handle_client_message(client_msg, &state).await;
                    if tx.send(response).await.is_err() {
                        break;
                    }
                }
            }
            Ok(Message::Ping(data)) => {
                let pong = ServerMessage::Pong { id: None };
                if tx.send(pong).await.is_err() {
                    break;
                }
            }
            Ok(Message::Close(_)) => {
                info!("WebSocket closed: {}", session_id);
                break;
            }
            Err(e) => {
                warn!("WebSocket error for {}: {}", session_id, e);
                break;
            }
            _ => {}
        }
    }

    // Cleanup
    sender_handle.abort();
    ws_state.broadcaster.remove_connection().await;
    info!("WebSocket disconnected: {}", session_id);
}

/// Handle client message.
async fn handle_client_message(msg: ClientMessage, state: &AppState) -> ServerMessage {
    match msg {
        ClientMessage::Subscribe { channel, params } => {
            let subscription_id = Uuid::new_v4().to_string();
            debug!("Subscribe to {}: {:?}", channel, params);
            ServerMessage::Subscribed {
                channel,
                subscription_id,
            }
        }
        ClientMessage::Unsubscribe { channel } => {
            debug!("Unsubscribe from {}", channel);
            ServerMessage::Unsubscribed { channel }
        }
        ClientMessage::Ping { id } => ServerMessage::Pong { id },
        ClientMessage::GetState { target } => {
            // Return mock state for now
            let data = match target.as_str() {
                "chain" => serde_json::json!({
                    "height": 100000,
                    "difficulty": "1234567890",
                    "hashrate": "100 TH/s"
                }),
                "mempool" => serde_json::json!({
                    "size": 100,
                    "bytes": 50000
                }),
                _ => serde_json::json!(null),
            };
            ServerMessage::State { target, data }
        }
    }
}

// ============================================================================
// Specialized WebSocket Handlers
// ============================================================================

/// Blocks-only WebSocket handler.
async fn ws_blocks_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    OptionalAuth(auth): OptionalAuth,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_blocks_socket(socket, state, auth))
}

async fn handle_blocks_socket(
    socket: WebSocket,
    state: AppState,
    _auth: Option<AuthContext>,
) {
    let ws_state = &state.websocket;
    ws_state.broadcaster.add_connection().await;

    let mut receiver = ws_state.broadcaster.subscribe(EventCategory::Blocks);
    let (mut sender, mut ws_receiver) = socket.split();

    // Send subscription confirmation
    let subscribed = ServerMessage::Subscribed {
        channel: "blocks".to_string(),
        subscription_id: Uuid::new_v4().to_string(),
    };
    if let Ok(msg) = serde_json::to_string(&subscribed) {
        let _ = sender.send(Message::Text(msg)).await;
    }

    loop {
        tokio::select! {
            event = receiver.recv() => {
                match event {
                    Ok(e) => {
                        let server_msg = ServerMessage::Event {
                            channel: "blocks".to_string(),
                            event: e,
                            timestamp: chrono::Utc::now().timestamp() as u64,
                        };
                        if let Ok(text) = serde_json::to_string(&server_msg) {
                            if sender.send(Message::Text(text)).await.is_err() {
                                break;
                            }
                        }
                    }
                    Err(_) => continue,
                }
            }
            msg = ws_receiver.next() => {
                match msg {
                    Some(Ok(Message::Close(_))) | None => break,
                    _ => {}
                }
            }
        }
    }

    ws_state.broadcaster.remove_connection().await;
}

/// Transactions-only WebSocket handler.
async fn ws_transactions_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    OptionalAuth(auth): OptionalAuth,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_transactions_socket(socket, state, auth))
}

async fn handle_transactions_socket(
    socket: WebSocket,
    state: AppState,
    _auth: Option<AuthContext>,
) {
    let ws_state = &state.websocket;
    ws_state.broadcaster.add_connection().await;

    let mut receiver = ws_state.broadcaster.subscribe(EventCategory::Transactions);
    let (mut sender, mut ws_receiver) = socket.split();

    let subscribed = ServerMessage::Subscribed {
        channel: "transactions".to_string(),
        subscription_id: Uuid::new_v4().to_string(),
    };
    if let Ok(msg) = serde_json::to_string(&subscribed) {
        let _ = sender.send(Message::Text(msg)).await;
    }

    loop {
        tokio::select! {
            event = receiver.recv() => {
                match event {
                    Ok(e) => {
                        let server_msg = ServerMessage::Event {
                            channel: "transactions".to_string(),
                            event: e,
                            timestamp: chrono::Utc::now().timestamp() as u64,
                        };
                        if let Ok(text) = serde_json::to_string(&server_msg) {
                            if sender.send(Message::Text(text)).await.is_err() {
                                break;
                            }
                        }
                    }
                    Err(_) => continue,
                }
            }
            msg = ws_receiver.next() => {
                match msg {
                    Some(Ok(Message::Close(_))) | None => break,
                    _ => {}
                }
            }
        }
    }

    ws_state.broadcaster.remove_connection().await;
}

/// Address-specific WebSocket handler.
async fn ws_address_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    OptionalAuth(auth): OptionalAuth,
    axum::extract::Path(address): axum::extract::Path<String>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_address_socket(socket, state, auth, address))
}

async fn handle_address_socket(
    socket: WebSocket,
    state: AppState,
    _auth: Option<AuthContext>,
    address: String,
) {
    let ws_state = &state.websocket;
    ws_state.broadcaster.add_connection().await;

    let mut receiver = ws_state.broadcaster.subscribe(EventCategory::Addresses);
    let (mut sender, mut ws_receiver) = socket.split();

    let subscribed = ServerMessage::Subscribed {
        channel: format!("address:{}", address),
        subscription_id: Uuid::new_v4().to_string(),
    };
    if let Ok(msg) = serde_json::to_string(&subscribed) {
        let _ = sender.send(Message::Text(msg)).await;
    }

    loop {
        tokio::select! {
            event = receiver.recv() => {
                match event {
                    Ok(Event::Address(ref addr_event)) if addr_event.address == address => {
                        let server_msg = ServerMessage::Event {
                            channel: format!("address:{}", address),
                            event: event.unwrap(),
                            timestamp: chrono::Utc::now().timestamp() as u64,
                        };
                        if let Ok(text) = serde_json::to_string(&server_msg) {
                            if sender.send(Message::Text(text)).await.is_err() {
                                break;
                            }
                        }
                    }
                    _ => continue,
                }
            }
            msg = ws_receiver.next() => {
                match msg {
                    Some(Ok(Message::Close(_))) | None => break,
                    _ => {}
                }
            }
        }
    }

    ws_state.broadcaster.remove_connection().await;
}

/// Contract-specific WebSocket handler.
async fn ws_contract_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    OptionalAuth(auth): OptionalAuth,
    axum::extract::Path(contract_id): axum::extract::Path<String>,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_contract_socket(socket, state, auth, contract_id))
}

async fn handle_contract_socket(
    socket: WebSocket,
    state: AppState,
    _auth: Option<AuthContext>,
    contract_id: String,
) {
    let ws_state = &state.websocket;
    ws_state.broadcaster.add_connection().await;

    let mut receiver = ws_state.broadcaster.subscribe(EventCategory::Contracts);
    let (mut sender, mut ws_receiver) = socket.split();

    let subscribed = ServerMessage::Subscribed {
        channel: format!("contract:{}", contract_id),
        subscription_id: Uuid::new_v4().to_string(),
    };
    if let Ok(msg) = serde_json::to_string(&subscribed) {
        let _ = sender.send(Message::Text(msg)).await;
    }

    loop {
        tokio::select! {
            event = receiver.recv() => {
                match event {
                    Ok(Event::Contract(ref contract_event)) if contract_event.contract_id == contract_id => {
                        let server_msg = ServerMessage::Event {
                            channel: format!("contract:{}", contract_id),
                            event: event.unwrap(),
                            timestamp: chrono::Utc::now().timestamp() as u64,
                        };
                        if let Ok(text) = serde_json::to_string(&server_msg) {
                            if sender.send(Message::Text(text)).await.is_err() {
                                break;
                            }
                        }
                    }
                    _ => continue,
                }
            }
            msg = ws_receiver.next() => {
                match msg {
                    Some(Ok(Message::Close(_))) | None => break,
                    _ => {}
                }
            }
        }
    }

    ws_state.broadcaster.remove_connection().await;
}

/// Markets WebSocket handler.
async fn ws_markets_handler(
    ws: WebSocketUpgrade,
    State(state): State<AppState>,
    OptionalAuth(auth): OptionalAuth,
) -> impl IntoResponse {
    ws.on_upgrade(move |socket| handle_markets_socket(socket, state, auth))
}

async fn handle_markets_socket(
    socket: WebSocket,
    state: AppState,
    _auth: Option<AuthContext>,
) {
    let ws_state = &state.websocket;
    ws_state.broadcaster.add_connection().await;

    let mut receiver = ws_state.broadcaster.subscribe(EventCategory::Markets);
    let (mut sender, mut ws_receiver) = socket.split();

    let subscribed = ServerMessage::Subscribed {
        channel: "markets".to_string(),
        subscription_id: Uuid::new_v4().to_string(),
    };
    if let Ok(msg) = serde_json::to_string(&subscribed) {
        let _ = sender.send(Message::Text(msg)).await;
    }

    loop {
        tokio::select! {
            event = receiver.recv() => {
                match event {
                    Ok(e) => {
                        let server_msg = ServerMessage::Event {
                            channel: "markets".to_string(),
                            event: e,
                            timestamp: chrono::Utc::now().timestamp() as u64,
                        };
                        if let Ok(text) = serde_json::to_string(&server_msg) {
                            if sender.send(Message::Text(text)).await.is_err() {
                                break;
                            }
                        }
                    }
                    Err(_) => continue,
                }
            }
            msg = ws_receiver.next() => {
                match msg {
                    Some(Ok(Message::Close(_))) | None => break,
                    _ => {}
                }
            }
        }
    }

    ws_state.broadcaster.remove_connection().await;
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_event_serialization() {
        let event = Event::Block(BlockEvent {
            hash: "abc123".to_string(),
            height: 100,
            timestamp: 1705312200,
            tx_count: 5,
            blue_score: 50,
            parent_hashes: vec!["parent1".to_string()],
        });

        let json = serde_json::to_string(&event).unwrap();
        assert!(json.contains("blocks"));
        assert!(json.contains("abc123"));
    }

    #[test]
    fn test_client_message_parsing() {
        let json = r#"{"type":"subscribe","channel":"blocks","params":{}}"#;
        let msg: ClientMessage = serde_json::from_str(json).unwrap();

        match msg {
            ClientMessage::Subscribe { channel, .. } => {
                assert_eq!(channel, "blocks");
            }
            _ => panic!("Expected Subscribe message"),
        }
    }

    #[test]
    fn test_server_message_serialization() {
        let msg = ServerMessage::Welcome {
            session_id: "test-session".to_string(),
            version: "1.0.0".to_string(),
        };

        let json = serde_json::to_string(&msg).unwrap();
        assert!(json.contains("welcome"));
        assert!(json.contains("test-session"));
    }
}
