// Filename: internal/handlers/websocket_handler.go
package handlers

import (
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/sirupsen/logrus"
)

type connWrapper struct {
	conn *websocket.Conn
	mu   sync.Mutex
}

type WebSocketHandler struct {
	logger   *logrus.Logger
	clients  sync.Map
	upgrader websocket.Upgrader
}

func NewWebSocketHandler(logger *logrus.Logger) *WebSocketHandler {
	return &WebSocketHandler{
		logger: logger,
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool { return true },
		},
	}
}

func (h *WebSocketHandler) HandleSystemLogs(c *gin.Context) {
	conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		h.logger.WithError(err).Error("Failed to upgrade websocket")
		return
	}
	defer conn.Close()

	clientID := time.Now().UnixNano()
	h.clients.Store(clientID, &connWrapper{conn: conn})
	defer h.clients.Delete(clientID)

	for {
		if _, _, err := conn.ReadMessage(); err != nil {
			break
		}
	}
}

func (h *WebSocketHandler) BroadcastLog(entry *logrus.Entry) {
	msg := map[string]interface{}{
		"timestamp": entry.Time.Format(time.RFC3339),
		"level":     entry.Level.String(),
		"message":   entry.Message,
		"fields":    entry.Data,
	}

	h.clients.Range(func(key, value interface{}) bool {
		wrapper := value.(*connWrapper)
		wrapper.mu.Lock()
		wrapper.conn.WriteJSON(msg)
		wrapper.mu.Unlock()
		return true
	})
}