package service

import (
	"encoding/json"
	"sync"
	"system-altrak/pkg/utils"

	"github.com/gofiber/websocket/v2"
	"go.uber.org/zap"
)

type WsMessage struct {
	Type    string      `json:"type"`
	Message string      `json:"message"`
	Data    interface{} `json:"data"`
}

type WsHub struct {
	clients    map[*websocket.Conn]bool
	Broadcast  chan WsMessage
	Register   chan *websocket.Conn
	Unregister chan *websocket.Conn
	mu         sync.Mutex
}

func NewWsHub() *WsHub {
	return &WsHub{
		clients:    make(map[*websocket.Conn]bool),
		Broadcast:  make(chan WsMessage),
		Register:   make(chan *websocket.Conn),
		Unregister: make(chan *websocket.Conn),
	}
}

func (h *WsHub) Run() {
	for {
		select {
		case client := <-h.Register:
			h.mu.Lock()
			h.clients[client] = true
			activeClients := len(h.clients)
			h.mu.Unlock()
			utils.Info("WS client connected", zap.Int("active_clients", activeClients))

		case client := <-h.Unregister:
			h.mu.Lock()
			if _, ok := h.clients[client]; ok {
				delete(h.clients, client)
				client.Close()
			}
			activeClients := len(h.clients)
			h.mu.Unlock()
			utils.Info("WS client disconnected", zap.Int("active_clients", activeClients))

		case message := <-h.Broadcast:
			h.mu.Lock()
			msgBytes, err := json.Marshal(message)
			if err != nil {
				h.mu.Unlock()
				utils.Error("WS message marshal failed",
					zap.Error(err),
					zap.String("type", message.Type),
				)
				continue
			}

			for client := range h.clients {
				if err := client.WriteMessage(websocket.TextMessage, msgBytes); err != nil {
					utils.Error("WS write failed", zap.Error(err))
					client.Close()
					delete(h.clients, client)
				}
			}
			h.mu.Unlock()
		}
	}
}