package hub

import (
	"encoding/json"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// WSMessage is sent to browsers (matches frontend ServerStatus shape).
type WSMessage struct {
	Type string                       `json:"type"`
	Data map[string]json.RawMessage `json:"data"`
}

type Hub struct {
	mu sync.RWMutex
	// serverId -> latest metrics JSON object (not wrapped)
	metrics map[string]json.RawMessage

	subsMu sync.Mutex
	// subscribers for broadcast
	conns map[*websocket.Conn]struct{}
}

func New() *Hub {
	return &Hub{
		metrics: make(map[string]json.RawMessage),
		conns:   make(map[*websocket.Conn]struct{}),
	}
}

func (h *Hub) SetMetrics(serverID string, metricsJSON []byte) {
	h.mu.Lock()
	if len(metricsJSON) == 0 {
		delete(h.metrics, serverID)
	} else {
		h.metrics[serverID] = json.RawMessage(append([]byte(nil), metricsJSON...))
	}
	h.mu.Unlock()
}

func (h *Hub) snapshotData() map[string]json.RawMessage {
	h.mu.RLock()
	defer h.mu.RUnlock()
	out := make(map[string]json.RawMessage, len(h.metrics))
	now := time.Now().UnixMilli()
	for id, raw := range h.metrics {
		// Wrap as ServerStatus for frontend compatibility
		wrap := map[string]any{
			"serverId":   id,
			"online":     true,
			"metrics":    json.RawMessage(raw),
			"lastUpdate": now,
		}
		b, err := json.Marshal(wrap)
		if err != nil {
			continue
		}
		out[id] = b
	}
	return out
}

func (h *Hub) BuildSnapshotMessage() ([]byte, error) {
	msg := WSMessage{
		Type: "snapshot",
		Data: h.snapshotData(),
	}
	return json.Marshal(msg)
}

// AddSubscriber registers for broadcast fan-out (call after sending the initial snapshot).
func (h *Hub) AddSubscriber(c *websocket.Conn) {
	h.subsMu.Lock()
	h.conns[c] = struct{}{}
	h.subsMu.Unlock()
}

func (h *Hub) Unregister(c *websocket.Conn) {
	h.subsMu.Lock()
	delete(h.conns, c)
	h.subsMu.Unlock()
}

// BroadcastSnapshot pushes full snapshot to all subscribers.
func (h *Hub) BroadcastSnapshot() {
	body, err := h.BuildSnapshotMessage()
	if err != nil {
		return
	}
	h.subsMu.Lock()
	defer h.subsMu.Unlock()
	for c := range h.conns {
		_ = c.SetWriteDeadline(time.Now().Add(10 * time.Second))
		if err := c.WriteMessage(websocket.TextMessage, body); err != nil {
			_ = c.Close()
			delete(h.conns, c)
		}
	}
}

// NotifyMetricsUpdate should be called after SetMetrics to push to browsers.
func (h *Hub) NotifyMetricsUpdate(serverID string, metricsJSON []byte) {
	h.SetMetrics(serverID, metricsJSON)
	h.BroadcastSnapshot()
}