// file: kb-onchain-transport/src/ws_pool.rs
// version: 6

//! WebSocket endpoint pool and role-based routing.

/// Pool of standard Solana WebSocket endpoints.
#[derive(Clone, Debug)]
pub struct WsEndpointPool {
    clients: std::vec::Vec<crate::WsClient>,
    next_index: std::sync::Arc<std::sync::atomic::AtomicUsize>,
}

impl crate::WsEndpointPool {
    /// Builds a pool from the active profile WebSocket endpoint list.
    pub fn from_profile(profile: &kb_config::ProfileConfig) -> kb_core::Result<Self> {
        let mut clients = std::vec::Vec::new();
        for endpoint in &profile.solana.ws_endpoints {
            if !endpoint.enabled {
                continue;
            }
            let client = match crate::WsClient::new(endpoint.clone()) {
                std::result::Result::Ok(client) => client,
                std::result::Result::Err(error) => return std::result::Result::Err(error),
            };
            clients.push(client);
        }
        return crate::WsEndpointPool::new(clients);
    }

    /// Creates a pool from already constructed clients.
    pub fn new(clients: std::vec::Vec<crate::WsClient>) -> kb_core::Result<Self> {
        if clients.is_empty() {
            tracing::error!(target: crate::TRACING_TARGET, action = "create_ws_pool", error_code = "ws_pool_empty", "WebSocket endpoint pool has no enabled endpoint");
            return std::result::Result::Err(kb_core::Error::config(
                "ws endpoint pool requires at least one enabled endpoint".to_string(),
            ));
        }
        tracing::debug!(target: crate::TRACING_TARGET, action = "create_ws_pool", endpoint_count = clients.len(), "WebSocket endpoint pool created");
        return std::result::Result::Ok(Self {
            clients,
            next_index: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)),
        });
    }

    /// Returns a serializable snapshot of every endpoint in the pool.
    pub fn snapshot(&self) -> std::vec::Vec<crate::WsPoolClientSnapshot> {
        let mut snapshots = std::vec::Vec::new();
        for client in &self.clients {
            snapshots.push(client.snapshot());
        }
        return snapshots;
    }

    /// Selects one endpoint for the requested role and method.
    pub fn select_client_for_role_and_method(
        &self,
        required_role: &str,
        method: &str,
    ) -> kb_core::Result<crate::WsClient> {
        let request_kind = crate::request_kind_from_method(method);
        return self.select_client_for_role_and_kind(required_role, &request_kind);
    }

    /// Selects one endpoint for the requested role and request kind.
    pub fn select_client_for_role_and_kind(
        &self,
        required_role: &str,
        request_kind: &str,
    ) -> kb_core::Result<crate::WsClient> {
        if self.clients.is_empty() {
            tracing::error!(target: crate::TRACING_TARGET, action = "select_ws_endpoint", required_role, request_kind, error_code = "ws_pool_empty", "WebSocket endpoint pool has no clients");
            return std::result::Result::Err(kb_core::Error::not_connected(
                "ws endpoint pool has no clients".to_string(),
            ));
        }
        let start_index = self.next_index.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let client_count = self.clients.len();
        let mut offset = 0_usize;
        while offset < client_count {
            let index = (start_index + offset) % client_count;
            let client = self.clients[index].clone();
            if client.can_handle(required_role, request_kind) {
                tracing::debug!(target: crate::TRACING_TARGET, action = "select_ws_endpoint", required_role, request_kind, endpoint_name = %client.endpoint_name(), provider = %client.provider(), "selected WebSocket endpoint");
                return std::result::Result::Ok(client);
            }
            offset += 1;
        }
        tracing::error!(target: crate::TRACING_TARGET, action = "select_ws_endpoint", required_role, request_kind, endpoint_count = self.clients.len(), error_code = "ws_endpoint_not_found", "no WebSocket endpoint supports requested role and kind");
        return std::result::Result::Err(kb_core::Error::config(format!(
            "no ws endpoint supports role '{}' and request kind '{}'",
            required_role, request_kind
        )));
    }

    /// Selects one endpoint for an explicitly registered standard WebSocket subscription.
    pub fn select_client_for_standard_subscription(
        &self,
        required_role: &str,
        subscription: &crate::StandardWsSubscriptionSpec,
    ) -> kb_core::Result<crate::WsClient> {
        return self
            .select_client_for_role_and_method(required_role, subscription.subscribe_method);
    }

    /// Executes one short WebSocket JSON-RPC request through the selected endpoint.
    pub async fn execute_json_rpc_once_for_role(
        &self,
        required_role: &str,
        method: std::string::String,
        params: std::vec::Vec<serde_json::Value>,
    ) -> kb_core::Result<crate::JsonRpcResponse> {
        let client = match self.select_client_for_role_and_method(required_role, &method) {
            std::result::Result::Ok(client) => client,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        return client.execute_json_rpc_once(method, params).await;
    }
}

#[cfg(test)]
mod tests {
    fn role_config(
        role: &str,
        request_kinds: std::vec::Vec<std::string::String>,
    ) -> kb_config::EndpointRoleConfig {
        return kb_config::EndpointRoleConfig {
            role: role.to_string(),
            enabled: true,
            request_kinds,
            priority: 1,
            requests_per_second: 10,
            burst_capacity: 10,
            max_concurrent_requests: 4,
            max_subscriptions: 16,
            pause_after_rate_limit_ms: 1500,
        };
    }

    fn endpoint(
        name: &str,
        role: &str,
        request_kinds: std::vec::Vec<std::string::String>,
    ) -> kb_config::WsEndpointConfig {
        return kb_config::WsEndpointConfig {
            name: name.to_string(),
            enabled: true,
            provider: "test".to_string(),
            cluster: "devnet".to_string(),
            url: format!("wss://{name}.invalid"),
            connect_timeout_ms: 100,
            request_timeout_ms: 100,
            unsubscribe_timeout_ms: 100,
            write_channel_capacity: 8,
            event_channel_capacity: 16,
            auto_reconnect: false,
            roles: std::vec![role_config(role, request_kinds)],
        };
    }

    fn client(endpoint: kb_config::WsEndpointConfig) -> crate::WsClient {
        match crate::WsClient::new(endpoint) {
            std::result::Result::Ok(client) => return client,
            std::result::Result::Err(error) => panic!("client creation failed: {error}"),
        }
    }

    #[test]
    fn new_rejects_empty_pool() {
        let result = crate::WsEndpointPool::new(std::vec::Vec::new());
        assert!(result.is_err());
    }

    #[test]
    fn snapshot_lists_every_client() {
        let pool = match crate::WsEndpointPool::new(std::vec![
            client(endpoint("a", "slot_notifications", std::vec!["slot_subscribe".to_string()])),
            client(endpoint("b", "slot_notifications", std::vec!["slot_subscribe".to_string()])),
        ]) {
            std::result::Result::Ok(pool) => pool,
            std::result::Result::Err(error) => panic!("pool creation failed: {error}"),
        };
        let snapshot = pool.snapshot();
        assert_eq!(snapshot.len(), 2);
        assert_eq!(snapshot[0].endpoint_name, "a");
        assert_eq!(snapshot[1].endpoint_name, "b");
    }

    #[test]
    fn select_client_round_robins_matching_clients() {
        let pool = match crate::WsEndpointPool::new(std::vec![
            client(endpoint("a", "slot_notifications", std::vec!["slot_subscribe".to_string()])),
            client(endpoint("b", "slot_notifications", std::vec!["slot_subscribe".to_string()])),
        ]) {
            std::result::Result::Ok(pool) => pool,
            std::result::Result::Err(error) => panic!("pool creation failed: {error}"),
        };
        let first =
            match pool.select_client_for_role_and_method("slot_notifications", "slotSubscribe") {
                std::result::Result::Ok(client) => client,
                std::result::Result::Err(error) => panic!("selection failed: {error}"),
            };
        let second =
            match pool.select_client_for_role_and_method("slot_notifications", "slotSubscribe") {
                std::result::Result::Ok(client) => client,
                std::result::Result::Err(error) => panic!("selection failed: {error}"),
            };
        assert_eq!(first.endpoint_name(), "a");
        assert_eq!(second.endpoint_name(), "b");
    }

    #[test]
    fn select_client_skips_unsupported_clients() {
        let pool = match crate::WsEndpointPool::new(std::vec![
            client(endpoint("a", "program_subscribe", std::vec!["program_subscribe".to_string()])),
            client(endpoint("b", "slot_notifications", std::vec!["slot_subscribe".to_string()])),
        ]) {
            std::result::Result::Ok(pool) => pool,
            std::result::Result::Err(error) => panic!("pool creation failed: {error}"),
        };
        let selected =
            match pool.select_client_for_role_and_method("slot_notifications", "slotSubscribe") {
                std::result::Result::Ok(client) => client,
                std::result::Result::Err(error) => panic!("selection failed: {error}"),
            };
        assert_eq!(selected.endpoint_name(), "b");
    }

    #[test]
    fn standard_subscription_selection_uses_the_subscribe_method() {
        let pool = match crate::WsEndpointPool::new(std::vec![client(endpoint(
            "a",
            "slot_notifications",
            std::vec!["slot_subscribe".to_string()],
        ))]) {
            std::result::Result::Ok(pool) => pool,
            std::result::Result::Err(error) => panic!("pool creation failed: {error}"),
        };
        let subscription = match crate::standard_ws_subscription("slotSubscribe") {
            std::option::Option::Some(value) => value,
            std::option::Option::None => panic!("slot subscription missing"),
        };
        let selected = match pool
            .select_client_for_standard_subscription("slot_notifications", subscription)
        {
            std::result::Result::Ok(client) => client,
            std::result::Result::Err(error) => panic!("selection failed: {error}"),
        };
        assert_eq!(selected.endpoint_name(), "a");
    }

    #[test]
    fn select_client_returns_error_for_missing_role() {
        let pool = match crate::WsEndpointPool::new(std::vec![client(endpoint(
            "a",
            "slot_notifications",
            std::vec!["slot_subscribe".to_string()]
        )),])
        {
            std::result::Result::Ok(pool) => pool,
            std::result::Result::Err(error) => panic!("pool creation failed: {error}"),
        };
        let selected =
            pool.select_client_for_role_and_kind("program_subscribe", "program_subscribe");
        assert!(selected.is_err());
    }
}