package agent

import (
	"context"
	"encoding/binary"
	"fmt"
	"log/slog"
	"testing"
	"time"

	"github.com/go-zeromq/zmq4"
	"github.com/vmihailenco/msgpack/v5"

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

func kvAgent() *Agent {
	return &Agent{log: slog.New(slog.DiscardHandler), kv: kvevents.NewMirror(),
		kvSubs: map[string]*kvevents.Subscriber{}}
}

// TestAnEnginesEventsReachTheControlPlaneStream runs the agent's half end to
// end: a listener for the shard, a publisher standing in for vLLM connecting
// to the port the engine would be given, and the stream to the control plane
// carrying the result.
func TestAnEnginesEventsReachTheControlPlaneStream(t *testing.T) {
	a := kvAgent()
	spec := engines.ShardSpec{ReplicaID: "r1", ShardID: "r1-0", NNodes: 1,
		Profile: engines.Profile{Engine: "vllm", KVEvents: true}}
	a.listenKV(&spec, 0)
	if spec.KVEventsPort == 0 {
		t.Fatal("the shard should have been given a port to publish to")
	}
	defer a.stopKV(spec.ShardID, spec.ReplicaID)

	pub := zmq4.NewPub(context.Background())
	defer func() { _ = pub.Close() }()
	if err := pub.Dial(fmt.Sprintf("tcp://127.0.0.1:%d", spec.KVEventsPort)); err != nil {
		t.Fatal(err)
	}
	tokens := make([]int, 32)
	for i := range tokens {
		tokens[i] = 500 + i
	}
	hs := kvevents.BlockHashes(tokens, 16, 0)
	payload, _ := msgpack.Marshal([]any{1.0, []any{map[string]any{
		"type": "BlockStored", "block_hashes": []any{hs[0], hs[1]}, "parent_block_hash": nil,
		"token_ids": tokens, "block_size": 16, "lora_id": nil, "medium": "GPU",
	}}})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	out := make(chan *llmgwv1.AgentMessage, 16)
	go a.forwardKV(ctx, out)

	deadline := time.After(5 * time.Second)
	var seq uint64
	for {
		// Sent until it arrives: a subscription reaches the publisher after
		// the connection, and what is published before it is dropped.
		var s [8]byte
		binary.BigEndian.PutUint64(s[:], seq)
		seq++
		_ = pub.Send(zmq4.NewMsgFrom([]byte("kv-events"), s[:], payload))
		select {
		case m := <-out:
			kv := m.GetKvEvents()
			if kv == nil || kv.GetReplicaId() != "r1" || kv.GetBlockSize() != 16 || len(kv.GetStored()) != 2 {
				t.Fatalf("message %v", m)
			}
			return
		case <-deadline:
			t.Fatal("nothing reached the stream")
		case <-time.After(50 * time.Millisecond):
		}
	}
}

func TestAShardWithoutKVEventsGetsNoListener(t *testing.T) {
	a := kvAgent()
	spec := engines.ShardSpec{ReplicaID: "r1", ShardID: "r1-0", NNodes: 1,
		Profile: engines.Profile{Engine: "vllm"}}
	a.listenKV(&spec, 0)
	if spec.KVEventsPort != 0 || len(a.kvSubs) != 0 {
		t.Fatal("kv_events is opt-in")
	}
}

func TestARepeatedStartReusesTheShardsListener(t *testing.T) {
	a := kvAgent()
	spec := engines.ShardSpec{ReplicaID: "r1", ShardID: "r1-0", NNodes: 1,
		Profile: engines.Profile{Engine: "vllm", KVEvents: true}}
	a.listenKV(&spec, 0)
	defer a.stopKV(spec.ShardID, spec.ReplicaID)
	again := spec
	again.KVEventsPort = 0
	a.listenKV(&again, 0)
	if again.KVEventsPort != spec.KVEventsPort || len(a.kvSubs) != 1 {
		t.Fatal("a start command repeated must not open a second listener")
	}
}