package agent

import (
	"context"
	"time"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
)

// kvFlushInterval is how often KV cache changes go up the stream. Often
// enough that a prefix cached by one request is known before the next few
// arrive; seldom enough that a busy engine's churn is sent as its net effect.
const kvFlushInterval = 250 * time.Millisecond

// listenKV starts receiving a shard's KV cache events, when its profile asks
// for them and this rank publishes any, and records the port in the spec for
// the engine's arguments. port is the one to reuse, or 0 for any.
//
// A failure is logged and the engine started without events: the router then
// routes the replica by tier 1, which is what it did before, rather than the
// replica not starting at all for want of an optimisation.
func (a *Agent) listenKV(spec *engines.ShardSpec, port int) {
	if !spec.Profile.KVEvents || spec.Profile.Engine != "vllm" || !spec.Leader() {
		return
	}
	a.mu.Lock()
	if sub, ok := a.kvSubs[spec.ShardID]; ok {
		a.mu.Unlock()
		// A repeated start command: the shard already has its listener.
		spec.KVEventsPort = sub.Port()
		return
	}
	a.mu.Unlock()

	sub, err := kvevents.Listen(spec.ReplicaID, port, a.kv, a.log)
	if err != nil {
		a.log.Warn("cannot receive kv cache events; the router will route this replica without them",
			"shard", spec.ShardID, "replica", spec.ReplicaID, "err", err)
		return
	}
	a.mu.Lock()
	a.kvSubs[spec.ShardID] = sub
	a.mu.Unlock()
	spec.KVEventsPort = sub.Port()
}

// stopKV stops a shard's listener and tells the control plane its replica
// holds nothing any more.
func (a *Agent) stopKV(shardID, replicaID string) {
	a.mu.Lock()
	sub, ok := a.kvSubs[shardID]
	delete(a.kvSubs, shardID)
	a.mu.Unlock()
	if !ok {
		return
	}
	sub.Close()
	a.kv.Remove(replicaID)
}

// forwardKV sends KV cache changes up one session's stream until it ends,
// beginning with a full copy of every replica: the control plane at the other
// end may have just started, or just taken over, and know nothing.
func (a *Agent) forwardKV(ctx context.Context, out chan<- *llmgwv1.AgentMessage) {
	s := a.kv.Open()
	defer s.Close()
	t := time.NewTicker(kvFlushInterval)
	defer t.Stop()
	for {
		for _, u := range s.Drain() {
			select {
			case out <- agentsvc.KVEventsMessage(u):
			case <-ctx.Done():
				return
			}
		}
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}