// Package ws is the WebSocket fan-out for SDP. Two flows:
//
//	agent            --(events)-->  /ws/agent          (one conn per agent)
//	dashboard client --(subscribe)-> /ws/deployments/{id}  (one conn per viewer)
//
// On agent connect we record the agent's nodeID. On dashboard connect we
// register a subscriber for one deployment. Events are best-effort fanned out;
// a slow client is dropped, not allowed to backpressure the agent.
package ws

import (
	"net/http"
	"sync"

	"github.com/gorilla/websocket"

	"github.com/sdp/protocol"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true }, // ponytail: internal tool, allow all
}

type Hub struct {
	mu sync.RWMutex

	// one channel per deployment; nil = no subscribers.
	subs map[string]map[chan protocol.Event]struct{}

	// one channel per connected agent; keyed by nodeID.
	agents map[string]chan<- []byte // outbound to agent (deploy requests, etc.)
}

func New() *Hub {
	return &Hub{
		subs:   make(map[string]map[chan protocol.Event]struct{}),
		agents: make(map[string]chan<- []byte),
	}
}

// Publish fans an event out to all subscribers of that deployment.
// Non-blocking; drops the event for a subscriber if its buffer is full.
func (h *Hub) Publish(e protocol.Event) {
	h.mu.RLock()
	subs := h.subs[e.DeploymentID]
	chans := make([]chan protocol.Event, 0, len(subs))
	for c := range subs {
		chans = append(chans, c)
	}
	h.mu.RUnlock()
	for _, c := range chans {
		select {
		case c <- e:
		default:
			// ponytail: drop slow subscriber; they'll reconnect.
		}
	}
}

// Subscribe registers a subscriber channel for one deployment.
// Returns an unsubscribe func the caller must invoke.
func (h *Hub) Subscribe(deploymentID string) (chan protocol.Event, func()) {
	ch := make(chan protocol.Event, 64)
	h.mu.Lock()
	if h.subs[deploymentID] == nil {
		h.subs[deploymentID] = make(map[chan protocol.Event]struct{})
	}
	h.subs[deploymentID][ch] = struct{}{}
	h.mu.Unlock()
	return ch, func() {
		h.mu.Lock()
		if subs, ok := h.subs[deploymentID]; ok {
			delete(subs, ch)
			if len(subs) == 0 {
				delete(h.subs, deploymentID)
			}
		}
		h.mu.Unlock()
		close(ch)
	}
}

// RegisterAgent stores the outbound channel for a node. Used when the control
// plane wants to send something back to an agent (deploy requests, stop, etc.).
func (h *Hub) RegisterAgent(nodeID string, out chan<- []byte) {
	h.mu.Lock()
	h.agents[nodeID] = out
	h.mu.Unlock()
}

func (h *Hub) UnregisterAgent(nodeID string) {
	h.mu.Lock()
	delete(h.agents, nodeID)
	h.mu.Unlock()
}

// SendToAgent best-effort sends a payload to a connected agent. Returns false
// if the agent isn't connected or its buffer is full.
func (h *Hub) SendToAgent(nodeID string, payload []byte) bool {
	h.mu.RLock()
	ch, ok := h.agents[nodeID]
	h.mu.RUnlock()
	if !ok {
		return false
	}
	select {
	case ch <- payload:
		return true
	default:
		return false
	}
}