bookhoard/internal/handlers/websocket.go

package handlers

import (
	"bookhoard/internal/database"
	"bookhoard/internal/middleware"
	"bookhoard/internal/sync"
	"context"
	"log"
	"net/http"
	"time"

	jwt "github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/labstack/echo/v5"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type WSHandler struct {
	db                   *database.Queries
	connManager          *sync.ConnectionManager
	jwtSecret            string
	deviceAuthMiddleware *middleware.DeviceAuthMiddleware
}

func NewWSHandler(db *database.Queries, connManager *sync.ConnectionManager, jwtSecret string, deviceAuth *middleware.DeviceAuthMiddleware) *WSHandler {
	return &WSHandler{
		db:                   db,
		connManager:          connManager,
		jwtSecret:            jwtSecret,
		deviceAuthMiddleware: deviceAuth,
	}
}

type WSMessage struct {
	Type      string                 `json:"type"`
	Timestamp string                 `json:"timestamp"`
	Data      map[string]interface{} `json:"data"`
}

type ClientInfo struct {
	UserID     string
	DeviceID   string
	DeviceType string
	DeviceName string
	IsDevice   bool
}

func (h *WSHandler) HandleWebSocket(c *echo.Context) error {
	token := c.QueryParam("token")
	if token == "" {
		return c.JSON(http.StatusUnauthorized, map[string]string{"error": "token required"})
	}

	clientInfo, err := h.authenticateClient(token, c.Request().Header.Get("Authorization"))
	if err != nil {
		return c.JSON(http.StatusUnauthorized, map[string]string{"error": err.Error()})
	}

	ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		log.Printf("WebSocket upgrade failed: %v", err)
		return err
	}

	conn := h.createConnection(clientInfo)
	h.connManager.AddConnection(conn)

	go h.readPump(conn, ws, clientInfo)
	go h.writePump(conn, ws, clientInfo)

	return nil
}

func (h *WSHandler) authenticateClient(token string, authHeader string) (*ClientInfo, error) {
	if authHeader != "" && len(authHeader) > 7 && authHeader[:7] == "Bearer " {
		deviceToken := authHeader[7:]
		device, err := h.deviceAuthMiddleware.ValidateDeviceToken(deviceToken)
		if err == nil {
			userID := uuid.UUID(device.UserID.Bytes).String()
			return &ClientInfo{
				UserID:     userID,
				DeviceID:   uuid.UUID(device.ID.Bytes).String(),
				DeviceType: device.DeviceType,
				DeviceName: device.DeviceName,
				IsDevice:   true,
			}, nil
		}
	}

	parsedToken, err := jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
		return []byte(h.jwtSecret), nil
	})
	if err == nil && parsedToken.Valid {
		claims := parsedToken.Claims.(jwt.MapClaims)
		userID := claims["user_id"].(string)
		return &ClientInfo{
			UserID:     userID,
			DeviceID:   "web-" + userID,
			DeviceType: "web",
			DeviceName: "Web Client",
			IsDevice:   false,
		}, nil
	}

	return nil, echo.NewHTTPError(http.StatusUnauthorized, "invalid token")
}

func (h *WSHandler) createConnection(clientInfo *ClientInfo) *sync.DeviceConnection {
	return &sync.DeviceConnection{
		ID:           clientInfo.DeviceID,
		UserID:       clientInfo.UserID,
		DeviceType:   clientInfo.DeviceType,
		DeviceName:   clientInfo.DeviceName,
		Connected:    time.Now(),
		LastPing:     time.Now(),
		Send:         make(chan sync.BroadcastMessage, 100),
		Disconnected: make(chan struct{}),
	}
}

func (h *WSHandler) readPump(conn *sync.DeviceConnection, ws *websocket.Conn, clientInfo *ClientInfo) {
	defer func() {
		h.connManager.RemoveConnection(conn.ID)
		ws.Close()
	}()

	ws.SetReadDeadline(time.Now().Add(90 * time.Second))
	ws.SetPongHandler(func(string) error {
		ws.SetReadDeadline(time.Now().Add(90 * time.Second))
		conn.LastPing = time.Now()
		return nil
	})

	for {
		_, message, err := ws.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("WebSocket error for %s: %v", conn.DeviceName, err)
			}
			break
		}

		h.handleMessage(message, conn, clientInfo)
	}
}

func (h *WSHandler) writePump(conn *sync.DeviceConnection, ws *websocket.Conn, clientInfo *ClientInfo) {
	ticker := time.NewTicker(30 * time.Second)
	defer func() {
		ticker.Stop()
		ws.Close()
	}()

	initialState := h.getInitialState(clientInfo.UserID)
	conn.Send <- sync.BroadcastMessage{
		Type:      sync.MessageTypeInitial,
		Timestamp: time.Now().Format(time.RFC3339),
		Data:      initialState,
	}

	for {
		select {
		case msg, ok := <-conn.Send:
			if !ok {
				return
			}
			ws.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := ws.WriteJSON(msg); err != nil {
				log.Printf("WebSocket write error for %s: %v", conn.DeviceName, err)
				return
			}
		case <-conn.Disconnected:
			return
		case <-ticker.C:
			ws.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := ws.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func (h *WSHandler) handleMessage(message []byte, conn *sync.DeviceConnection, clientInfo *ClientInfo) {
	log.Printf("WebSocket message from %s: %s", conn.DeviceName, string(message))
}

func (h *WSHandler) getInitialState(userID string) map[string]interface{} {
	userUUID, err := uuid.Parse(userID)
	if err != nil {
		return map[string]interface{}{"error": "invalid user ID"}
	}

	pgUserID := pgtype.UUID{Bytes: userUUID, Valid: true}

	mediaItems, err := h.db.GetUserMediaItemsForSync(context.Background(), pgUserID)
	if err != nil {
		return map[string]interface{}{"error": "failed to fetch initial state"}
	}

	progressMap := make(map[string]map[string]interface{})
	for _, item := range mediaItems {
		progress, err := h.db.GetUniversalProgress(context.Background(), database.GetUniversalProgressParams{
			MediaItemID: item.ID,
			UserID:      pgUserID,
		})
		if err == nil {
			bookID := uuid.UUID(item.ID.Bytes).String()
			progressMap[bookID] = map[string]interface{}{
				"percentage":   progress.Percentage.Float64,
				"current_page": progress.CurrentPage.Int32,
				"total_pages":  progress.TotalPages.Int32,
				"last_read":    progress.LastReadAt.Time,
			}
		}
	}

	return map[string]interface{}{
		"progress": progressMap,
		"devices":  h.connManager.GetConnectionStats(),
	}
}