package prediction

import (
	"sync"
)

type Event struct {
	Status  string         `json:"status"`
	Logs    string         `json:"logs,omitempty"`
	Metrics map[string]any `json:"metrics,omitempty"`
	Output  any            `json:"output,omitempty"`
	Error   string         `json:"error,omitempty"`
	Done    bool           `json:"-"`
}

type Hub struct {
	mu   sync.RWMutex
	subs map[string]map[chan Event]struct{}
}

func NewHub() *Hub {
	return &Hub{subs: make(map[string]map[chan Event]struct{})}
}

func (h *Hub) Subscribe(predictionID string) chan Event {
	ch := make(chan Event, 16)
	h.mu.Lock()
	if h.subs[predictionID] == nil {
		h.subs[predictionID] = make(map[chan Event]struct{})
	}
	h.subs[predictionID][ch] = struct{}{}
	h.mu.Unlock()
	return ch
}

func (h *Hub) Unsubscribe(predictionID string, ch chan Event) {
	h.mu.Lock()
	if m := h.subs[predictionID]; m != nil {
		delete(m, ch)
		if len(m) == 0 {
			delete(h.subs, predictionID)
		}
	}
	h.mu.Unlock()
	close(ch)
}

func (h *Hub) Publish(predictionID string, ev Event) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for ch := range h.subs[predictionID] {
		select {
		case ch <- ev:
		default:
		}
	}
}