// file: kb-onchain-transport/src/ws_client.rs
// version: 8

//! Standard Solana WebSocket client helpers.

use futures_util::SinkExt; // rust-rules: trait-import
use futures_util::StreamExt; // rust-rules: trait-import

/// Snapshot of one pooled WebSocket endpoint.
#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub struct WsPoolClientSnapshot {
    /// Logical endpoint name.
    pub endpoint_name: std::string::String,
    /// Provider name.
    pub provider: std::string::String,
    /// Endpoint URL.
    pub endpoint_url: std::string::String,
    /// Supported roles.
    pub roles: std::vec::Vec<crate::EndpointRoleSnapshot>,
    /// Status string.
    pub status: std::string::String,
}

/// Standard Solana WebSocket client bound to one configured endpoint.
#[derive(Clone, Debug)]
pub struct WsClient {
    endpoint: kb_config::WsEndpointConfig,
    next_request_id: std::sync::Arc<std::sync::atomic::AtomicU64>,
}

impl crate::WsClient {
    /// Creates a new WebSocket client bound to one endpoint.
    pub fn new(endpoint: kb_config::WsEndpointConfig) -> kb_core::Result<Self> {
        if !endpoint.enabled {
            tracing::error!(target: crate::TRACING_TARGET, action = "create_ws_client", endpoint_name = %endpoint.name, provider = %endpoint.provider, error_code = "ws_endpoint_disabled", "cannot create WebSocket client for disabled endpoint");
            return std::result::Result::Err(kb_core::Error::config(format!(
                "ws endpoint '{}' is disabled",
                endpoint.name
            )));
        }
        tracing::debug!(target: crate::TRACING_TARGET, action = "create_ws_client", endpoint_name = %endpoint.name, provider = %endpoint.provider, role_count = endpoint.roles.len(), "WebSocket client created");
        return std::result::Result::Ok(Self {
            endpoint,
            next_request_id: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(1)),
        });
    }

    /// Returns the endpoint name.
    pub fn endpoint_name(&self) -> &str {
        return self.endpoint.name.as_str();
    }

    /// Returns the provider name.
    pub fn provider(&self) -> &str {
        return self.endpoint.provider.as_str();
    }

    /// Returns the endpoint URL.
    pub fn endpoint_url(&self) -> &str {
        return self.endpoint.url.as_str();
    }

    /// Returns the endpoint configuration.
    pub fn endpoint_config(&self) -> &kb_config::WsEndpointConfig {
        return &self.endpoint;
    }

    /// Returns true when this endpoint supports the required role and request kind.
    pub fn can_handle(&self, required_role: &str, request_kind: &str) -> bool {
        if !self.endpoint.enabled {
            return false;
        }
        for role in &self.endpoint.roles {
            if crate::role_matches(role, required_role, request_kind) {
                return true;
            }
        }
        return false;
    }

    /// Returns a serializable endpoint snapshot.
    pub fn snapshot(&self) -> crate::WsPoolClientSnapshot {
        let mut roles = std::vec::Vec::new();
        for role in &self.endpoint.roles {
            roles.push(crate::EndpointRoleSnapshot::from_config(role));
        }
        return crate::WsPoolClientSnapshot {
            endpoint_name: self.endpoint.name.clone(),
            provider: self.endpoint.provider.clone(),
            endpoint_url: self.endpoint.url.clone(),
            roles,
            status: "idle".to_string(),
        };
    }

    /// Builds a JSON-RPC request with a generated id.
    pub fn build_json_rpc_request(
        &self,
        method: std::string::String,
        params: std::vec::Vec<serde_json::Value>,
    ) -> crate::JsonRpcRequest {
        let request_id = self.next_request_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        return crate::JsonRpcRequest::new_with_u64_id(request_id, method, params);
    }

    /// Builds a subscribe request for one explicitly registered standard subscription.
    pub fn build_standard_subscribe_request(
        &self,
        subscription: &crate::StandardWsSubscriptionSpec,
        params: std::vec::Vec<serde_json::Value>,
    ) -> crate::JsonRpcRequest {
        return self.build_json_rpc_request(subscription.subscribe_method.to_string(), params);
    }

    /// Builds an unsubscribe request for one explicitly registered standard subscription.
    pub fn build_standard_unsubscribe_request(
        &self,
        subscription: &crate::StandardWsSubscriptionSpec,
        subscription_id: u64,
    ) -> crate::JsonRpcRequest {
        return self.build_json_rpc_request(
            subscription.unsubscribe_method.to_string(),
            std::vec![serde_json::Value::from(subscription_id)],
        );
    }

    /// Connects, sends one JSON-RPC request, waits for one response and closes.
    pub async fn execute_json_rpc_once(
        &self,
        method: std::string::String,
        params: std::vec::Vec<serde_json::Value>,
    ) -> kb_core::Result<crate::JsonRpcResponse> {
        let parameter_count = params.len();
        let request = self.build_json_rpc_request(method.clone(), params);
        tracing::debug!(target: crate::TRACING_TARGET, action = "execute_ws_json_rpc_once", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, parameter_count, "start one-shot WebSocket JSON-RPC request");
        let request_text = match request.to_json_string() {
            std::result::Result::Ok(text) => text,
            std::result::Result::Err(error) => {
                tracing::error!(target: crate::TRACING_TARGET, action = "serialize_ws_json_rpc_request", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error = %error, "WebSocket JSON-RPC request serialization failed");
                return std::result::Result::Err(error);
            },
        };
        let connect_timeout = std::time::Duration::from_millis(self.endpoint.connect_timeout_ms);
        let connect_future = tokio_tungstenite::connect_async(self.endpoint.url.as_str());
        let connect_timeout_result = tokio::time::timeout(connect_timeout, connect_future).await;
        let connect_result = match connect_timeout_result {
            std::result::Result::Ok(result) => result,
            std::result::Result::Err(_) => {
                tracing::error!(target: crate::TRACING_TARGET, action = "connect_ws_endpoint", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, timeout_ms = self.endpoint.connect_timeout_ms, error_code = "ws_connect_timeout", "WebSocket endpoint connection timed out");
                return std::result::Result::Err(kb_core::Error::ws(format!(
                    "websocket connect timed out for endpoint '{}'",
                    self.endpoint.name
                )));
            },
        };
        let (mut stream, _response) = match connect_result {
            std::result::Result::Ok(pair) => pair,
            std::result::Result::Err(error) => {
                tracing::error!(target: crate::TRACING_TARGET, action = "connect_ws_endpoint", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error = %error, "WebSocket endpoint connection failed");
                return std::result::Result::Err(kb_core::Error::ws(format!(
                    "cannot connect websocket endpoint '{}': {error}",
                    self.endpoint.name
                )));
            },
        };
        let send_result = stream
            .send(tokio_tungstenite::tungstenite::Message::Text(request_text.into()))
            .await;
        if let std::result::Result::Err(error) = send_result {
            tracing::error!(target: crate::TRACING_TARGET, action = "send_ws_json_rpc_request", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error = %error, "WebSocket JSON-RPC request send failed");
            return std::result::Result::Err(kb_core::Error::ws(format!(
                "cannot send websocket request '{}' to endpoint '{}': {error}",
                method, self.endpoint.name
            )));
        }
        let response_timeout = std::time::Duration::from_millis(self.endpoint.request_timeout_ms);
        let next_timeout_result = tokio::time::timeout(response_timeout, stream.next()).await;
        let next_result = match next_timeout_result {
            std::result::Result::Ok(result) => result,
            std::result::Result::Err(_) => {
                tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, timeout_ms = self.endpoint.request_timeout_ms, error_code = "ws_response_timeout", "WebSocket JSON-RPC response timed out");
                return std::result::Result::Err(kb_core::Error::ws(format!(
                    "websocket response timed out for endpoint '{}'",
                    self.endpoint.name
                )));
            },
        };
        let message = match next_result {
            std::option::Option::Some(result) => match result {
                std::result::Result::Ok(message) => message,
                std::result::Result::Err(error) => {
                    tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error = %error, "WebSocket JSON-RPC response read failed");
                    return std::result::Result::Err(kb_core::Error::ws(format!(
                        "websocket read failed for endpoint '{}': {error}",
                        self.endpoint.name
                    )));
                },
            },
            std::option::Option::None => {
                tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error_code = "ws_closed_before_response", "WebSocket endpoint closed before response");
                return std::result::Result::Err(kb_core::Error::ws(format!(
                    "websocket endpoint '{}' closed before response",
                    self.endpoint.name
                )));
            },
        };
        let close_result = stream
            .send(tokio_tungstenite::tungstenite::Message::Close(std::option::Option::None))
            .await;
        if let std::result::Result::Err(error) = close_result {
            tracing::debug!(target: crate::TRACING_TARGET, action = "close_ws_endpoint", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, error = %error, "WebSocket close send failed");
        }
        return match message {
            tokio_tungstenite::tungstenite::Message::Text(text) => {
                let parse_result = crate::parse_json_rpc_text(text.as_str());
                match parse_result {
                    std::result::Result::Ok(response) => {
                        if let crate::JsonRpcResponse::Error(error_response) = &response {
                            tracing::error!(target: crate::TRACING_TARGET, action = "execute_ws_json_rpc_once", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, response_kind = response.kind_name(), rpc_error_code = error_response.error.code, rpc_error_message = %error_response.error.message, "WebSocket JSON-RPC endpoint returned an RPC error");
                        } else {
                            tracing::debug!(target: crate::TRACING_TARGET, action = "execute_ws_json_rpc_once", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, response_kind = response.kind_name(), "one-shot WebSocket JSON-RPC request completed");
                        }
                        std::result::Result::Ok(response)
                    },
                    std::result::Result::Err(error) => {
                        tracing::error!(target: crate::TRACING_TARGET, action = "parse_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, response_byte_length = text.len(), error = %error, "WebSocket JSON-RPC response parsing failed");
                        std::result::Result::Err(error)
                    },
                }
            },
            tokio_tungstenite::tungstenite::Message::Binary(binary) => {
                tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, message_kind = "binary", response_byte_length = binary.len(), "WebSocket endpoint returned binary data before JSON response");
                std::result::Result::Err(kb_core::Error::ws(format!(
                    "websocket endpoint '{}' returned binary message with {} bytes",
                    self.endpoint.name,
                    binary.len()
                )))
            },
            tokio_tungstenite::tungstenite::Message::Ping(_) => {
                tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, message_kind = "ping", "WebSocket endpoint returned ping before JSON response");
                std::result::Result::Err(kb_core::Error::ws(format!(
                    "websocket endpoint '{}' returned ping before json response",
                    self.endpoint.name
                )))
            },
            tokio_tungstenite::tungstenite::Message::Pong(_) => {
                tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, message_kind = "pong", "WebSocket endpoint returned pong before JSON response");
                std::result::Result::Err(kb_core::Error::ws(format!(
                    "websocket endpoint '{}' returned pong before json response",
                    self.endpoint.name
                )))
            },
            tokio_tungstenite::tungstenite::Message::Close(_) => {
                tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, message_kind = "close", "WebSocket endpoint closed before JSON response");
                std::result::Result::Err(kb_core::Error::ws(format!(
                    "websocket endpoint '{}' closed before json response",
                    self.endpoint.name
                )))
            },
            tokio_tungstenite::tungstenite::Message::Frame(_) => {
                tracing::error!(target: crate::TRACING_TARGET, action = "read_ws_json_rpc_response", endpoint_name = %self.endpoint.name, provider = %self.endpoint.provider, request_id = ?request.id, method = %method, message_kind = "frame", "WebSocket endpoint returned raw frame before JSON response");
                std::result::Result::Err(kb_core::Error::ws(format!(
                    "websocket endpoint '{}' returned raw frame before json response",
                    self.endpoint.name
                )))
            },
        };
    }
}

#[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(enabled: bool) -> kb_config::WsEndpointConfig {
        return kb_config::WsEndpointConfig {
            name: "ws_a".to_string(),
            enabled,
            provider: "test".to_string(),
            cluster: "devnet".to_string(),
            url: "wss://example.invalid".to_string(),
            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("slot_notifications", std::vec!["slot_subscribe".to_string()]),
                role_config("program_subscribe", std::vec!["program_subscribe".to_string()]),
                role_config("ws_any", std::vec!["*".to_string()]),
            ],
        };
    }

    #[test]
    fn new_rejects_disabled_endpoint() {
        let result = crate::WsClient::new(endpoint(false));
        assert!(result.is_err());
    }

    #[test]
    fn can_handle_matches_exact_role_and_kind() {
        let client = match crate::WsClient::new(endpoint(true)) {
            std::result::Result::Ok(client) => client,
            std::result::Result::Err(error) => panic!("client creation failed: {error}"),
        };
        assert!(client.can_handle("slot_notifications", "slot_subscribe"));
        assert!(!client.can_handle("slot_notifications", "root_subscribe"));
    }

    #[test]
    fn can_handle_matches_wildcard_kind() {
        let client = match crate::WsClient::new(endpoint(true)) {
            std::result::Result::Ok(client) => client,
            std::result::Result::Err(error) => panic!("client creation failed: {error}"),
        };
        assert!(client.can_handle("ws_any", "logs_subscribe_mentions"));
    }

    #[test]
    fn snapshot_preserves_endpoint_metadata() {
        let client = match crate::WsClient::new(endpoint(true)) {
            std::result::Result::Ok(client) => client,
            std::result::Result::Err(error) => panic!("client creation failed: {error}"),
        };
        let snapshot = client.snapshot();
        assert_eq!(snapshot.endpoint_name, "ws_a");
        assert_eq!(snapshot.provider, "test");
        assert_eq!(snapshot.endpoint_url, "wss://example.invalid");
        assert_eq!(snapshot.roles.len(), 3);
    }

    #[test]
    fn standard_subscription_builders_preserve_exact_methods_and_subscription_id() {
        let client = match crate::WsClient::new(endpoint(true)) {
            std::result::Result::Ok(client) => client,
            std::result::Result::Err(error) => panic!("client 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 subscribe = client.build_standard_subscribe_request(
            subscription,
            std::vec![serde_json::json!({"commitment": "confirmed"})],
        );
        let unsubscribe = client.build_standard_unsubscribe_request(subscription, 42);
        assert_eq!(subscribe.method, "slotSubscribe");
        assert_eq!(subscribe.params.len(), 1);
        assert_eq!(unsubscribe.method, "slotUnsubscribe");
        assert_eq!(unsubscribe.params, std::vec![serde_json::Value::from(42)]);
        assert_eq!(subscribe.id, serde_json::Value::from(1));
        assert_eq!(unsubscribe.id, serde_json::Value::from(2));
    }

    #[test]
    fn build_json_rpc_request_increments_ids() {
        let client = match crate::WsClient::new(endpoint(true)) {
            std::result::Result::Ok(client) => client,
            std::result::Result::Err(error) => panic!("client creation failed: {error}"),
        };
        let first =
            client.build_json_rpc_request("slotSubscribe".to_string(), std::vec::Vec::new());
        let second =
            client.build_json_rpc_request("rootSubscribe".to_string(), std::vec::Vec::new());
        assert_eq!(first.id, serde_json::Value::from(1));
        assert_eq!(second.id, serde_json::Value::from(2));
    }
}