bookhoard/internal/sync/websocket.go

package sync

import (
	"context"
	"log"
	"sync"
	"time"

	"github.com/google/uuid"
)

// Message types for WebSocket communication
const (
	MessageTypeProgressUpdate   = "progress_update"
	MessageTypeAnnotationUpdate = "annotation_update"
	MessageTypeConflict         = "conflict"
	MessageTypeSyncComplete     = "sync_complete"
	MessageTypeHeartbeat        = "heartbeat"
	MessageTypeInitial          = "initial_state"
)

// BroadcastMessage represents a message to broadcast to connected clients
type BroadcastMessage struct {
	Type      string        `json:"type"`
	Timestamp string        `json:"timestamp"`
	Data      interface{}   `json:"data"`
	Source    *SourceDevice `json:"source_device,omitempty"`
}

// SourceDevice represents the device that sent the update
type SourceDevice struct {
	ID   string `json:"id"`
	Name string `json:"name"`
	Type string `json:"type"`
}

// DeviceConnection represents an active WebSocket connection
type DeviceConnection struct {
	ID           string
	UserID       string
	DeviceType   string
	DeviceName   string
	Connected    time.Time
	LastPing     time.Time
	Send         chan BroadcastMessage
	Disconnected chan struct{}
}

// ConnectionManager manages WebSocket connections
type ConnectionManager struct {
	connections map[string]*DeviceConnection
	broadcast   chan BroadcastMessage
	mu          sync.RWMutex
}

// NewConnectionManager creates a new connection manager
func NewConnectionManager() *ConnectionManager {
	cm := &ConnectionManager{
		connections: make(map[string]*DeviceConnection),
		broadcast:   make(chan BroadcastMessage, 100),
	}

	// Start broadcast goroutine
	go cm.broadcastLoop()

	return cm
}

// broadcastLoop handles broadcasting messages to all connected clients
func (m *ConnectionManager) broadcastLoop() {
	for {
		select {
		case msg := <-m.broadcast:
			m.mu.RLock()
			for _, conn := range m.connections {
				select {
				case conn.Send <- msg:
				default:
					// Channel full, skip this connection
				}
			}
			m.mu.RUnlock()
		}
	}
}

// Broadcast sends a message to all connected clients
func (m *ConnectionManager) Broadcast(msg BroadcastMessage) {
	select {
	case m.broadcast <- msg:
	default:
		log.Println("Broadcast channel full, message dropped")
	}
}

// BroadcastToUser sends a message to all connections for a specific user
func (m *ConnectionManager) BroadcastToUser(userID string, msg BroadcastMessage) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	for _, conn := range m.connections {
		if conn.UserID == userID {
			select {
			case conn.Send <- msg:
			default:
				// Channel full, skip this connection
				log.Printf("WebSocket: Channel full for %s, skipping broadcast", conn.DeviceName)
			}
		}
	}
}

// BroadcastProgressUpdate broadcasts a progress update to all connected clients
func (m *ConnectionManager) BroadcastProgressUpdate(bookID uuid.UUID, percentage float64, source SourceDevice) {
	msg := BroadcastMessage{
		Type:      MessageTypeProgressUpdate,
		Timestamp: time.Now().Format(time.RFC3339),
		Data: map[string]interface{}{
			"book_id":    bookID.String(),
			"percentage": percentage,
		},
		Source: &source,
	}
	m.Broadcast(msg)
}

// BroadcastAnnotationUpdate broadcasts an annotation update to all connected clients
func (m *ConnectionManager) BroadcastAnnotationUpdate(bookID uuid.UUID, annotationType string, data interface{}, source SourceDevice) {
	msg := BroadcastMessage{
		Type:      MessageTypeAnnotationUpdate,
		Timestamp: time.Now().Format(time.RFC3339),
		Data: map[string]interface{}{
			"book_id":         bookID.String(),
			"annotation_type": annotationType,
			"data":            data,
		},
		Source: &source,
	}
	m.Broadcast(msg)
}

// BroadcastConflictNotification broadcasts a conflict notification to all connected clients
func (m *ConnectionManager) BroadcastConflictNotification(bookID [16]byte, notificationType string, conflictID string) {
	msg := BroadcastMessage{
		Type:      MessageTypeConflict,
		Timestamp: time.Now().Format(time.RFC3339),
		Data: map[string]interface{}{
			"book_id":           uuid.UUID(bookID).String(),
			"notification_type": notificationType,
			"conflict_id":       conflictID,
		},
	}
	m.Broadcast(msg)
}

// AddConnection adds a new WebSocket connection
func (m *ConnectionManager) AddConnection(conn *DeviceConnection) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.connections[conn.ID] = conn
	log.Printf("WebSocket: Connection added for device %s (user: %s, type: %s)", conn.DeviceName, conn.UserID, conn.DeviceType)
}

// RemoveConnection removes a WebSocket connection
func (m *ConnectionManager) RemoveConnection(connID string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if conn, exists := m.connections[connID]; exists {
		close(conn.Disconnected)
		delete(m.connections, connID)
		log.Printf("WebSocket: Connection removed for device %s", conn.DeviceName)
	}
}

// GetConnection retrieves a connection by ID
func (m *ConnectionManager) GetConnection(connID string) (*DeviceConnection, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	conn, exists := m.connections[connID]
	return conn, exists
}

// GetUserConnections returns all connections for a specific user
func (m *ConnectionManager) GetUserConnections(userID string) []*DeviceConnection {
	m.mu.RLock()
	defer m.mu.RUnlock()

	var conns []*DeviceConnection
	for _, conn := range m.connections {
		if conn.UserID == userID {
			conns = append(conns, conn)
		}
	}
	return conns
}

// GetConnectionCount returns the number of active connections
func (m *ConnectionManager) GetConnectionCount() int {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return len(m.connections)
}

// CleanupStaleConnections removes connections that haven't sent a ping in 2 minutes
func (m *ConnectionManager) CleanupStaleConnections() {
	m.mu.Lock()
	defer m.mu.Unlock()

	staleThreshold := time.Now().Add(-2 * time.Minute)
	for id, conn := range m.connections {
		if conn.LastPing.Before(staleThreshold) {
			log.Printf("WebSocket: Cleaning up stale connection for %s", conn.DeviceName)
			close(conn.Disconnected)
			delete(m.connections, id)
		}
	}
}

// GetConnectionStats returns statistics about connections
func (m *ConnectionManager) GetConnectionStats() map[string]int {
	m.mu.RLock()
	defer m.mu.RUnlock()

	stats := make(map[string]int)
	for _, conn := range m.connections {
		stats[conn.DeviceType]++
	}
	return stats
}

// StartCleanupTask starts a background task to cleanup stale connections
func (m *ConnectionManager) StartCleanupTask() context.CancelFunc {
	ticker := time.NewTicker(1 * time.Minute)
	ctx, cancel := context.WithCancel(context.Background())

	go func() {
		for {
			select {
			case <-ctx.Done():
				ticker.Stop()
				return
			case <-ticker.C:
				m.CleanupStaleConnections()
			}
		}
	}()

	return cancel
}