sonr/test/e2e/client/websocket.go

package client

import (
	"context"
	"encoding/json"
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/gorilla/websocket"
)

// WebSocketClient provides WebSocket client for CometBFT event subscription
type WebSocketClient struct {
	baseURL string
	conn    *websocket.Conn
}

// EventSubscription represents an event subscription
type EventSubscription struct {
	Query  string
	Events chan *SubscriptionEvent
	Errors chan error
	done   chan struct{}
}

// SubscriptionEvent represents an event received via subscription
type SubscriptionEvent struct {
	Query  string          `json:"query"`
	Data   EventResultData `json:"data"`
	Events []any           `json:"events,omitempty"`
}

// EventResultData represents the data part of a subscription event
type EventResultData struct {
	Type  string `json:"type"`
	Value any    `json:"value"`
}

// JSONRPCRequest represents a JSON-RPC request
type JSONRPCRequest struct {
	JSONRPC string `json:"jsonrpc"`
	Method  string `json:"method"`
	Params  any    `json:"params"`
	ID      int    `json:"id"`
}

// JSONRPCResponse represents a JSON-RPC response
type JSONRPCResponse struct {
	JSONRPC string        `json:"jsonrpc"`
	Result  any           `json:"result,omitempty"`
	Error   *JSONRPCError `json:"error,omitempty"`
	ID      int           `json:"id"`
}

// JSONRPCError represents a JSON-RPC error
type JSONRPCError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Data    string `json:"data,omitempty"`
}

// SubscribeParams represents subscription parameters
type SubscribeParams struct {
	Query string `json:"query"`
}

// NewWebSocketClient creates a new WebSocket client
func NewWebSocketClient(baseURL string) *WebSocketClient {
	return &WebSocketClient{
		baseURL: baseURL,
	}
}

// Connect establishes a WebSocket connection to CometBFT
func (ws *WebSocketClient) Connect(ctx context.Context) error {
	// Convert HTTP URL to WebSocket URL
	wsURL := strings.Replace(ws.baseURL, "http://", "ws://", 1)
	wsURL = strings.Replace(wsURL, "https://", "wss://", 1)
	wsURL += "/websocket"

	u, err := url.Parse(wsURL)
	if err != nil {
		return fmt.Errorf("invalid WebSocket URL: %w", err)
	}

	dialer := websocket.Dialer{
		HandshakeTimeout: 10 * time.Second,
	}

	conn, _, err := dialer.DialContext(ctx, u.String(), nil)
	if err != nil {
		return fmt.Errorf("failed to connect to WebSocket: %w", err)
	}

	ws.conn = conn
	return nil
}

// Close closes the WebSocket connection
func (ws *WebSocketClient) Close() error {
	if ws.conn != nil {
		return ws.conn.Close()
	}
	return nil
}

// Subscribe subscribes to events matching the given query
func (ws *WebSocketClient) Subscribe(ctx context.Context, query string) (*EventSubscription, error) {
	if ws.conn == nil {
		return nil, fmt.Errorf("WebSocket connection not established")
	}

	// Send subscription request
	req := JSONRPCRequest{
		JSONRPC: "2.0",
		Method:  "subscribe",
		Params: SubscribeParams{
			Query: query,
		},
		ID: 1,
	}

	if err := ws.conn.WriteJSON(req); err != nil {
		return nil, fmt.Errorf("failed to send subscription request: %w", err)
	}

	// Read subscription response
	var resp JSONRPCResponse
	if err := ws.conn.ReadJSON(&resp); err != nil {
		return nil, fmt.Errorf("failed to read subscription response: %w", err)
	}

	if resp.Error != nil {
		return nil, fmt.Errorf("subscription error: %s", resp.Error.Message)
	}

	// Create subscription
	subscription := &EventSubscription{
		Query:  query,
		Events: make(chan *SubscriptionEvent, 100),
		Errors: make(chan error, 10),
		done:   make(chan struct{}),
	}

	// Start listening for events
	go ws.listenForEvents(ctx, subscription)

	return subscription, nil
}

// SubscribeToNewBlocks subscribes to new block events
func (ws *WebSocketClient) SubscribeToNewBlocks(ctx context.Context) (*EventSubscription, error) {
	return ws.Subscribe(ctx, "tm.event = 'NewBlock'")
}

// SubscribeToNewBlockHeaders subscribes to new block header events
func (ws *WebSocketClient) SubscribeToNewBlockHeaders(ctx context.Context) (*EventSubscription, error) {
	return ws.Subscribe(ctx, "tm.event = 'NewBlockHeader'")
}

// SubscribeToTxEvents subscribes to transaction events
func (ws *WebSocketClient) SubscribeToTxEvents(ctx context.Context) (*EventSubscription, error) {
	return ws.Subscribe(ctx, "tm.event = 'Tx'")
}

// SubscribeToDIDEvents subscribes to DID module events
func (ws *WebSocketClient) SubscribeToDIDEvents(ctx context.Context) (*EventSubscription, error) {
	return ws.Subscribe(ctx, "did.v1.EventDIDCreated EXISTS OR did.v1.EventDIDUpdated EXISTS OR did.v1.EventDIDDeactivated EXISTS")
}

// SubscribeToDWNEvents subscribes to DWN module events
func (ws *WebSocketClient) SubscribeToDWNEvents(ctx context.Context) (*EventSubscription, error) {
	return ws.Subscribe(ctx, "dwn.v1.EventRecordWritten EXISTS OR dwn.v1.EventRecordDeleted EXISTS")
}

// SubscribeToCustomEvents subscribes to custom events with specific attributes
func (ws *WebSocketClient) SubscribeToCustomEvents(ctx context.Context, eventType, attributeKey, attributeValue string) (*EventSubscription, error) {
	query := fmt.Sprintf("%s EXISTS", eventType)
	if attributeKey != "" && attributeValue != "" {
		query += fmt.Sprintf(" AND %s.%s = '%s'", eventType, attributeKey, attributeValue)
	}
	return ws.Subscribe(ctx, query)
}

// listenForEvents listens for incoming events on the WebSocket connection
func (ws *WebSocketClient) listenForEvents(ctx context.Context, subscription *EventSubscription) {
	defer close(subscription.Events)
	defer close(subscription.Errors)

	for {
		select {
		case <-ctx.Done():
			return
		case <-subscription.done:
			return
		default:
			// Set read deadline
			if err := ws.conn.SetReadDeadline(time.Now().Add(30 * time.Second)); err != nil {
				subscription.Errors <- fmt.Errorf("failed to set read deadline: %w", err)
				return
			}

			var message json.RawMessage
			if err := ws.conn.ReadJSON(&message); err != nil {
				if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
					return
				}
				subscription.Errors <- fmt.Errorf("failed to read WebSocket message: %w", err)
				continue
			}

			// Try to parse as JSON-RPC response first
			var resp JSONRPCResponse
			if err := json.Unmarshal(message, &resp); err == nil && resp.Result != nil {
				// This is likely an event notification
				var event SubscriptionEvent
				if eventBytes, err := json.Marshal(resp.Result); err == nil {
					if err := json.Unmarshal(eventBytes, &event); err == nil {
						event.Query = subscription.Query
						select {
						case subscription.Events <- &event:
						case <-ctx.Done():
							return
						case <-subscription.done:
							return
						}
					}
				}
			}
		}
	}
}

// Unsubscribe unsubscribes from the event subscription
func (ws *WebSocketClient) Unsubscribe(ctx context.Context, subscription *EventSubscription) error {
	if ws.conn == nil {
		return fmt.Errorf("WebSocket connection not established")
	}

	// Send unsubscribe request
	req := JSONRPCRequest{
		JSONRPC: "2.0",
		Method:  "unsubscribe",
		Params: SubscribeParams{
			Query: subscription.Query,
		},
		ID: 2,
	}

	if err := ws.conn.WriteJSON(req); err != nil {
		return fmt.Errorf("failed to send unsubscribe request: %w", err)
	}

	// Signal the listening goroutine to stop
	close(subscription.done)

	return nil
}

// Close closes the event subscription
func (sub *EventSubscription) Close() {
	if sub.done != nil {
		select {
		case <-sub.done:
			// Already closed
		default:
			close(sub.done)
		}
	}
}

// WaitForEvent waits for a specific event with timeout
func (sub *EventSubscription) WaitForEvent(ctx context.Context, timeout time.Duration, eventFilter func(*SubscriptionEvent) bool) (*SubscriptionEvent, error) {
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	for {
		select {
		case <-ctx.Done():
			return nil, fmt.Errorf("timeout waiting for event")
		case err := <-sub.Errors:
			return nil, fmt.Errorf("subscription error: %w", err)
		case event := <-sub.Events:
			if event == nil {
				return nil, fmt.Errorf("event channel closed")
			}
			if eventFilter == nil || eventFilter(event) {
				return event, nil
			}
		}
	}
}

// WaitForEventByType waits for an event of a specific type
func (sub *EventSubscription) WaitForEventByType(ctx context.Context, timeout time.Duration, eventType string) (*SubscriptionEvent, error) {
	return sub.WaitForEvent(ctx, timeout, func(event *SubscriptionEvent) bool {
		// This is a simplified check - in practice, you'd parse the event data more carefully
		eventStr := fmt.Sprintf("%v", event.Data.Value)
		return strings.Contains(eventStr, eventType)
	})
}

// GetAllEvents returns all events received so far (non-blocking)
func (sub *EventSubscription) GetAllEvents() []*SubscriptionEvent {
	var events []*SubscriptionEvent

	for {
		select {
		case event := <-sub.Events:
			if event == nil {
				return events
			}
			events = append(events, event)
		default:
			return events
		}
	}
}