dumpsterChat/cmd/tui/ws.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/url"
	"strings"
	"sync"
	"time"

	tea "github.com/charmbracelet/bubbletea"
	"github.com/gorilla/websocket"
)

// WSEvent represents a raw websocket event from the server.
type WSEvent struct {
	Type string          `json:"type"`
	Data json.RawMessage `json:"data,omitempty"`
}

// WebSocket connection manager.
type WSConn struct {
	conn   *websocket.Conn
	mu     sync.Mutex
	done   chan struct{}
	closed bool
}

// ConnectWS dials the websocket endpoint and authenticates via an initial
// message frame (not a URL query parameter, to avoid leaking the token in logs).
func ConnectWS(serverURL, token string) (*WSConn, error) {
	u, err := url.Parse(serverURL)
	if err != nil {
		return nil, err
	}
	// Switch to ws/wss scheme
	if u.Scheme == "https" {
		u.Scheme = "wss"
	} else {
		u.Scheme = "ws"
	}
	u.Path = "/ws"

	conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		return nil, err
	}

	// Send auth token as the first message frame.
	authMsg, _ := json.Marshal(map[string]string{"token": token})
	if err := conn.WriteMessage(websocket.TextMessage, authMsg); err != nil {
		conn.Close()
		return nil, fmt.Errorf("send auth: %w", err)
	}

	// Wait for the server's ready response.
	conn.SetReadDeadline(time.Now().Add(10 * time.Second))
	_, raw, err := conn.ReadMessage()
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("auth timeout: %w", err)
	}
	var resp struct {
		Type  string `json:"type"`
		Error string `json:"error"`
	}
	if err := json.Unmarshal(raw, &resp); err != nil || resp.Type != "ready" {
		conn.Close()
		msg := "unexpected response"
		if resp.Error != "" {
			msg = resp.Error
		}
		return nil, fmt.Errorf("auth failed: %s", msg)
	}
	conn.SetReadDeadline(time.Time{})

	ws := &WSConn{
		conn: conn,
		done: make(chan struct{}),
	}
	return ws, nil
}

// Listen starts reading messages and sends them as tea.Msgs to the program.
// It returns a tea.Cmd that listens for the next message.
func (ws *WSConn) Listen() tea.Cmd {
	return func() tea.Msg {
		_, raw, err := ws.conn.ReadMessage()
		if err != nil {
			return wsErrMsg{err: err}
		}

		var event WSEvent
		if err := json.Unmarshal(raw, &event); err != nil {
			return wsErrMsg{err: err}
		}

		return wsEventMsg(event)
	}
}

// Send sends a JSON event to the server.
func (ws *WSConn) Send(eventType string, data interface{}) error {
	ws.mu.Lock()
	defer ws.mu.Unlock()

	msg := map[string]interface{}{
		"type": eventType,
		"data": data,
	}
	payload, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	return ws.conn.WriteMessage(websocket.TextMessage, payload)
}

// Close closes the websocket connection.
func (ws *WSConn) Close() {
	ws.mu.Lock()
	defer ws.mu.Unlock()
	if !ws.closed {
		ws.closed = true
		ws.conn.WriteMessage(websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		ws.conn.Close()
		close(ws.done)
	}
}

// StartPing sends periodic pings to keep the connection alive.
func (ws *WSConn) StartPing() {
	go func() {
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				ws.mu.Lock()
				if !ws.closed {
					_ = ws.conn.WriteMessage(websocket.PingMessage, nil)
				}
				ws.mu.Unlock()
			case <-ws.done:
				return
			}
		}
	}()
}

// ---- Message types for bubbletea ----

type wsEventMsg WSEvent

type wsErrMsg struct {
	err error
}

// Parse helpers for ws event data

type wsMessageData struct {
	ID          string  `json:"id"`
	ChannelID   string  `json:"channel_id"`
	AuthorID    string  `json:"author_id"`
	AuthorName  string  `json:"author_username"`
	DisplayName *string `json:"author_display_name"`
	Content     string  `json:"content"`
	EditedAt    *string `json:"edited_at"`
	CreatedAt   string  `json:"created_at"`
}

type wsPresenceData struct {
	UserID   string `json:"user_id"`
	Username string `json:"username"`
	Status   string `json:"status"`
}

type wsTypingData struct {
	ChannelID string `json:"channel_id"`
	UserID    string `json:"user_id"`
	Username  string `json:"username"`
}

type wsVoiceData struct {
	RoomName string `json:"room_name"`
	UserID   string `json:"user_id"`
	Username string `json:"username"`
}

// parseWSEventMessage extracts a Message from event data.
func parseWSEventMessage(data json.RawMessage) (*Message, error) {
	var m Message
	if err := json.Unmarshal(data, &m); err != nil {
		return nil, err
	}
	return &m, nil
}

// parseWSDeleteData extracts the message ID from a delete event.
func parseWSDeleteData(data json.RawMessage) (id string, channelID string) {
	var d struct {
		ID        string `json:"id"`
		ChannelID string `json:"channel_id"`
	}
	if err := json.Unmarshal(data, &d); err != nil {
		log.Printf("ws: failed to parse delete data: %v", err)
		return
	}
	return d.ID, d.ChannelID
}

// parseWSTyping extracts typing event data.
func parseWSTyping(data json.RawMessage) (username string, channelID string) {
	var d struct {
		ChannelID string `json:"channel_id"`
		Username  string `json:"username"`
	}
	if err := json.Unmarshal(data, &d); err != nil {
		return
	}
	return d.Username, d.ChannelID
}

// trimURLScheme strips http(s):// for display
func trimURLScheme(s string) string {
	s = strings.TrimPrefix(s, "https://")
	s = strings.TrimPrefix(s, "http://")
	return s
}