package pushy

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

// Sink is a per-device outbound channel. Send returns false if the
// device's bounded write queue is full (a slow consumer) — the caller
// must NOT block, so one slow device can't head-of-line the whole node.
type Sink interface {
	Send(msg []byte) bool
}

// Node holds live device connections and routes messages to them.
type Node struct {
	ID       string
	Registry Registry
	TTL      time.Duration

	mu    sync.RWMutex
	sinks map[string]Sink
}

func NewNode(id string, r Registry) *Node {
	return &Node{ID: id, Registry: r, TTL: 60 * time.Second, sinks: map[string]Sink{}}
}

// Connect registers a device's sink on this node and records ownership.
func (n *Node) Connect(deviceID string, s Sink) {
	n.mu.Lock()
	n.sinks[deviceID] = s
	n.mu.Unlock()
	n.Registry.Register(deviceID, n.ID, n.TTL)
}

// Disconnect removes a device and clears registry ownership.
func (n *Node) Disconnect(deviceID string) {
	n.mu.Lock()
	delete(n.sinks, deviceID)
	n.mu.Unlock()
	n.Registry.Unregister(deviceID, n.ID)
}

// Deliver enqueues msg for a locally-connected device. Returns false if
// the device isn't on this node, or its queue is full.
func (n *Node) Deliver(deviceID string, msg []byte) bool {
	n.mu.RLock()
	s, ok := n.sinks[deviceID]
	n.mu.RUnlock()
	if !ok {
		return false
	}
	return s.Send(msg)
}

// Count returns the number of connected devices (connection density).
func (n *Node) Count() int {
	n.mu.RLock()
	defer n.mu.RUnlock()
	return len(n.sinks)
}

// Devices returns the connected device IDs (for drain).
func (n *Node) Devices() []string {
	n.mu.RLock()
	defer n.mu.RUnlock()
	out := make([]string, 0, len(n.sinks))
	for d := range n.sinks {
		out = append(out, d)
	}
	return out
}

// PushEvent is what a backend publishes: "deliver Payload to DeviceID".
type PushEvent struct {
	DeviceID string `json:"deviceId"`
	Payload  []byte `json:"payload"`
	MsgID    string `json:"msgId"` // for at-least-once dedupe
}

// NodeClient forwards a message to the (possibly remote) owning node.
type NodeClient interface {
	Forward(node string, ev PushEvent) bool
}

// LocalNodeClient maps node IDs to in-process *Node (the lab transport).
type LocalNodeClient map[string]*Node

func (c LocalNodeClient) Forward(node string, ev PushEvent) bool {
	n, ok := c[node]
	if !ok {
		return false
	}
	return n.Deliver(ev.DeviceID, ev.Payload)
}

// MessageProcessor consumes push events and routes each by registry
// lookup — decoupling publishers from delivery (Kafka in production).
type MessageProcessor struct {
	Registry Registry
	Nodes    NodeClient

	Delivered, Misses, ForwardFails int64
}

// Consume processes events off topic until ctx is cancelled.
func (mp *MessageProcessor) Consume(ctx context.Context, topic <-chan []byte) {
	for {
		select {
		case <-ctx.Done():
			return
		case raw, ok := <-topic:
			if !ok {
				return
			}
			mp.processOne(raw)
		}
	}
}

func (mp *MessageProcessor) processOne(raw []byte) {
	var ev PushEvent
	if json.Unmarshal(raw, &ev) != nil {
		return
	}
	node, ok := mp.Registry.Lookup(ev.DeviceID)
	if !ok {
		mp.Misses++ // device offline/moved: drop, or queue per policy
		return
	}
	if mp.Nodes.Forward(node, ev) {
		mp.Delivered++
	} else {
		mp.ForwardFails++ // node drained between lookup and forward; redeliver via reconnect
	}
}

// Publish enqueues an event (fire-and-forget for the backend).
func Publish(topic chan<- []byte, ev PushEvent) bool {
	raw, _ := json.Marshal(ev)
	select {
	case topic <- raw:
		return true
	default:
		return false // topic full: in prod Kafka absorbs this
	}
}