package mxgateway

import (
	"context"
	"errors"
	"io"
	"net"
	"testing"
	"time"

	pb "gitea.dohertylan.com/dohertj2/mxaccessgw/clients/go/internal/generated"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"google.golang.org/grpc/test/bufconn"
)

const bufSize = 1024 * 1024

func TestDialAttachesAuthMetadataToUnaryCalls(t *testing.T) {
	fake := &fakeGatewayServer{
		openReply: &pb.OpenSessionReply{
			SessionId:              "session-1",
			GatewayProtocolVersion: GatewayProtocolVersion,
			WorkerProtocolVersion:  WorkerProtocolVersion,
			ProtocolStatus: &pb.ProtocolStatus{
				Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
			},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()

	_, err := client.OpenSession(context.Background(), OpenSessionOptions{ClientSessionName: "fixture"})
	if err != nil {
		t.Fatalf("OpenSession() error = %v", err)
	}

	if got := fake.openAuth; got != "Bearer test-api-key" {
		t.Fatalf("authorization metadata = %q, want %q", got, "Bearer test-api-key")
	}
}

func TestStreamEventsAttachesAuthMetadataAndClosesOnCancellation(t *testing.T) {
	fake := &fakeGatewayServer{
		streamStarted: make(chan struct{}),
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")
	ctx, cancel := context.WithCancel(context.Background())

	events, err := session.Events(ctx)
	if err != nil {
		t.Fatalf("Events() error = %v", err)
	}
	<-fake.streamStarted

	first := <-events
	if first.Err != nil {
		t.Fatalf("first event error = %v", first.Err)
	}
	if first.Event.GetWorkerSequence() != 1 {
		t.Fatalf("worker sequence = %d, want 1", first.Event.GetWorkerSequence())
	}
	if got := fake.streamAuth; got != "Bearer test-api-key" {
		t.Fatalf("stream authorization metadata = %q, want %q", got, "Bearer test-api-key")
	}

	cancel()
	select {
	case _, ok := <-events:
		if ok {
			t.Fatal("events channel produced an extra item after cancellation")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("events channel did not close after cancellation")
	}
}

func TestEventSubscriptionCloseStopsStream(t *testing.T) {
	fake := &fakeGatewayServer{
		streamStarted: make(chan struct{}),
		streamDone:    make(chan struct{}),
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	subscription, err := session.SubscribeEvents(context.Background())
	if err != nil {
		t.Fatalf("SubscribeEvents() error = %v", err)
	}
	<-fake.streamStarted
	first := <-subscription.Events()
	if first.Err != nil {
		t.Fatalf("first event error = %v", first.Err)
	}

	subscription.Close()

	select {
	case <-fake.streamDone:
	case <-time.After(2 * time.Second):
		t.Fatal("event stream did not stop after subscription close")
	}
	select {
	case _, ok := <-subscription.Events():
		if ok {
			t.Fatal("subscription channel remained open after close")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("subscription channel did not close")
	}
}

func TestEventsAfterCancelsStreamWhenCompatibilityChannelIsAbandoned(t *testing.T) {
	fake := &fakeGatewayServer{
		streamStarted:    make(chan struct{}),
		streamDone:       make(chan struct{}),
		streamEventCount: 64,
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	events, err := session.EventsAfter(context.Background(), 0)
	if err != nil {
		t.Fatalf("EventsAfter() error = %v", err)
	}
	<-fake.streamStarted

	select {
	case <-fake.streamDone:
	case <-time.After(2 * time.Second):
		t.Fatal("compatibility event stream did not stop after result channel filled")
	}

	for {
		select {
		case _, ok := <-events:
			if !ok {
				return
			}
		case <-time.After(2 * time.Second):
			t.Fatal("compatibility event channel did not close")
		}
	}
}

func TestEventsSlowConsumerYieldsErrSlowConsumerBeforeClose(t *testing.T) {
	fake := &fakeGatewayServer{
		streamStarted:    make(chan struct{}),
		streamDone:       make(chan struct{}),
		streamEventCount: 64,
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	events, err := session.EventsAfter(context.Background(), 0)
	if err != nil {
		t.Fatalf("EventsAfter() error = %v", err)
	}
	<-fake.streamStarted

	// Do not drain the channel so the buffer overflows. The stream stops once
	// the slow-consumer cancel fires on the producer side.
	select {
	case <-fake.streamDone:
	case <-time.After(2 * time.Second):
		t.Fatal("event stream did not stop after buffer overflow")
	}

	var last EventResult
	gotResult := false
	for {
		select {
		case res, ok := <-events:
			if !ok {
				if !gotResult {
					t.Fatal("events channel closed without yielding any result")
				}
				if !errors.Is(last.Err, ErrSlowConsumer) {
					t.Fatalf("final event result err = %v, want one wrapping ErrSlowConsumer", last.Err)
				}
				var gwErr *GatewayError
				if !errors.As(last.Err, &gwErr) {
					t.Fatalf("final event result err is %T, want *GatewayError", last.Err)
				}
				if gwErr.Op != "stream events" {
					t.Fatalf("final event result err Op = %q, want %q", gwErr.Op, "stream events")
				}
				return
			}
			last = res
			gotResult = true
		case <-time.After(2 * time.Second):
			t.Fatal("events channel did not close after slow-consumer termination")
		}
	}
}

func TestEventsFullBufferTerminalErrorKeepsRootCause(t *testing.T) {
	fake := &fakeGatewayServer{
		streamStarted:     make(chan struct{}),
		streamDone:        make(chan struct{}),
		streamEventCount:  eventBufferSize,
		streamTerminalErr: status.Error(codes.Internal, "boom"),
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	events, err := session.EventsAfter(context.Background(), 0)
	if err != nil {
		t.Fatalf("EventsAfter() error = %v", err)
	}
	<-fake.streamStarted

	// Do not drain until the stream has fully ended: the server sends exactly
	// eventBufferSize events (filling the data slots) and then returns a genuine
	// terminal gRPC error. The client must report that error as itself, using the
	// reserved slot, rather than mislabeling it as ErrSlowConsumer.
	select {
	case <-fake.streamDone:
	case <-time.After(2 * time.Second):
		t.Fatal("event stream did not stop after terminal error")
	}
	// streamDone fires when the server returns; the client's producer goroutine
	// still needs a moment to drain the gRPC stream, fill all data slots, and
	// enqueue the terminal result. Let it settle before draining so the buffer is
	// genuinely full when the terminal error is processed (which is what makes the
	// mislabel bug observable).
	time.Sleep(250 * time.Millisecond)

	var last EventResult
	gotResult := false
	for {
		select {
		case res, ok := <-events:
			if !ok {
				if !gotResult {
					t.Fatal("events channel closed without yielding any result")
				}
				var gwErr *GatewayError
				if !errors.As(last.Err, &gwErr) {
					t.Fatalf("final event result err is %T, want *GatewayError", last.Err)
				}
				if code := status.Code(last.Err); code != codes.Internal {
					t.Fatalf("final event result gRPC code = %s, want %s", code, codes.Internal)
				}
				if errors.Is(last.Err, ErrSlowConsumer) {
					t.Fatalf("final event result err = %v, must not be mislabeled as ErrSlowConsumer", last.Err)
				}
				return
			}
			last = res
			gotResult = true
		case <-time.After(2 * time.Second):
			t.Fatal("events channel did not close after terminal error")
		}
	}
}

// TestSubscribeEventsFullBufferDeliversTerminalError is the CLI-44 regression for
// the never-drop Subscribe path. SubscribeEvents/SubscribeEventsAfter use
// cancelWhenResultBufferFull=false, so ordinary sends are blocking and uncapped and
// can fill every slot in the results channel — including the reserved terminal slot.
// A genuine terminal Recv error must still be delivered as the final result, never
// silently dropped. The server sends eventBufferSize+eventBufferReservedSlots events
// (filling every slot) and then returns a genuine gRPC error; with an unconditional
// non-blocking terminal send the error is dropped, so this fails red until the send
// path blocks for the never-drop mode.
func TestSubscribeEventsFullBufferDeliversTerminalError(t *testing.T) {
	fake := &fakeGatewayServer{
		streamStarted:     make(chan struct{}),
		streamDone:        make(chan struct{}),
		streamEventCount:  eventBufferSize + eventBufferReservedSlots,
		streamTerminalErr: status.Error(codes.Internal, "boom"),
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	subscription, err := session.SubscribeEvents(context.Background())
	if err != nil {
		t.Fatalf("SubscribeEvents() error = %v", err)
	}
	defer subscription.Close()
	<-fake.streamStarted

	// Wait for the server to finish sending every event and return the terminal
	// error, so the producer goroutine has filled every buffered slot before the
	// terminal result is processed. That is what makes the dropped-terminal bug
	// observable: with the buffer full, an unconditional non-blocking send discards
	// the terminal error.
	select {
	case <-fake.streamDone:
	case <-time.After(2 * time.Second):
		t.Fatal("event stream did not stop after terminal error")
	}
	time.Sleep(250 * time.Millisecond)

	// Drain fully. Every data event, then the terminal gRPC error as the final
	// result, must arrive; the channel must not close without yielding it.
	events := subscription.Events()
	var last EventResult
	gotResult := false
	for {
		select {
		case res, ok := <-events:
			if !ok {
				if !gotResult {
					t.Fatal("events channel closed without yielding any result")
				}
				var gwErr *GatewayError
				if !errors.As(last.Err, &gwErr) {
					t.Fatalf("final event result err is %T (%v), want the terminal *GatewayError; it was dropped", last.Err, last.Err)
				}
				if code := status.Code(last.Err); code != codes.Internal {
					t.Fatalf("final event result gRPC code = %s, want %s", code, codes.Internal)
				}
				return
			}
			last = res
			gotResult = true
		case <-time.After(2 * time.Second):
			t.Fatal("events channel did not close after terminal error")
		}
	}
}

func TestEventsSurfacesReplayGapSentinelAsTypedSignal(t *testing.T) {
	fake := &fakeGatewayServer{
		streamStarted: make(chan struct{}),
		streamReplayGap: &pb.ReplayGap{
			RequestedAfterSequence:  5,
			OldestAvailableSequence: 42,
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	events, err := session.EventsAfter(ctx, 5)
	if err != nil {
		t.Fatalf("EventsAfter() error = %v", err)
	}
	<-fake.streamStarted

	// First result must be the typed replay-gap signal, not a normal event.
	first := <-events
	if first.Err != nil {
		t.Fatalf("first result error = %v", first.Err)
	}
	if !first.IsReplayGap() {
		t.Fatalf("first result IsReplayGap() = false, want true")
	}
	if first.Event != nil {
		t.Fatalf("replay-gap result carried a non-nil Event %+v; want the sentinel to clear Event", first.Event)
	}
	if got := first.ReplayGap.GetRequestedAfterSequence(); got != 5 {
		t.Fatalf("ReplayGap.RequestedAfterSequence = %d, want 5", got)
	}
	if got := first.ReplayGap.GetOldestAvailableSequence(); got != 42 {
		t.Fatalf("ReplayGap.OldestAvailableSequence = %d, want 42", got)
	}

	// Normal events after the sentinel are unaffected: Event set, ReplayGap nil.
	second := <-events
	if second.Err != nil {
		t.Fatalf("second result error = %v", second.Err)
	}
	if second.IsReplayGap() {
		t.Fatalf("second result IsReplayGap() = true, want false for a normal event")
	}
	if second.Event == nil {
		t.Fatal("second result Event = nil, want a normal event")
	}
	if got := second.Event.GetWorkerSequence(); got != 1 {
		t.Fatalf("normal event worker sequence = %d, want 1", got)
	}
	if second.Event.GetFamily() != pb.MxEventFamily_MX_EVENT_FAMILY_ON_DATA_CHANGE {
		t.Fatalf("normal event family = %s, want ON_DATA_CHANGE", second.Event.GetFamily())
	}
}

func TestSessionHelpersBuildCommandsAndExposeRawReply(t *testing.T) {
	fake := &fakeGatewayServer{
		invokeReply: &pb.MxCommandReply{
			SessionId: "session-1",
			Kind:      pb.MxCommandKind_MX_COMMAND_KIND_ADD_ITEM2,
			ProtocolStatus: &pb.ProtocolStatus{
				Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
			},
			Payload: &pb.MxCommandReply_AddItem2{
				AddItem2: &pb.AddItem2Reply{ItemHandle: 42},
			},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	itemHandle, err := session.AddItem2(context.Background(), 12, "Area001.Pump001.Speed", "runtime")
	if err != nil {
		t.Fatalf("AddItem2() error = %v", err)
	}

	if itemHandle != 42 {
		t.Fatalf("item handle = %d, want 42", itemHandle)
	}
	req := fake.invokeRequest
	if req.GetSessionId() != "session-1" {
		t.Fatalf("session id = %q, want session-1", req.GetSessionId())
	}
	if req.GetClientCorrelationId() == "" {
		t.Fatal("client correlation id is empty")
	}
	if req.GetCommand().GetKind() != pb.MxCommandKind_MX_COMMAND_KIND_ADD_ITEM2 {
		t.Fatalf("command kind = %s", req.GetCommand().GetKind())
	}
	if req.GetCommand().GetAddItem2().GetItemContext() != "runtime" {
		t.Fatalf("item context = %q, want runtime", req.GetCommand().GetAddItem2().GetItemContext())
	}
}

func TestSubscribeBulkBuildsOneBulkCommandAndReturnsResults(t *testing.T) {
	fake := &fakeGatewayServer{
		invokeReply: &pb.MxCommandReply{
			SessionId: "session-1",
			Kind:      pb.MxCommandKind_MX_COMMAND_KIND_SUBSCRIBE_BULK,
			ProtocolStatus: &pb.ProtocolStatus{
				Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
			},
			Payload: &pb.MxCommandReply_SubscribeBulk{
				SubscribeBulk: &pb.BulkSubscribeReply{
					Results: []*pb.SubscribeResult{
						{
							ServerHandle:  12,
							TagAddress:    "Area001.Pump001.Speed",
							ItemHandle:    34,
							WasSuccessful: true,
						},
					},
				},
			},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	results, err := session.SubscribeBulk(context.Background(), 12, []string{"Area001.Pump001.Speed"})
	if err != nil {
		t.Fatalf("SubscribeBulk() error = %v", err)
	}

	if len(results) != 1 || results[0].GetItemHandle() != 34 {
		t.Fatalf("results = %#v, want item handle 34", results)
	}
	req := fake.invokeRequest
	if req.GetCommand().GetKind() != pb.MxCommandKind_MX_COMMAND_KIND_SUBSCRIBE_BULK {
		t.Fatalf("command kind = %s", req.GetCommand().GetKind())
	}
	if got := req.GetCommand().GetSubscribeBulk().GetTagAddresses(); len(got) != 1 || got[0] != "Area001.Pump001.Speed" {
		t.Fatalf("tag addresses = %#v", got)
	}
}

func TestWriteBulkBuildsOneBulkCommandAndReturnsPerEntryResults(t *testing.T) {
	fake := &fakeGatewayServer{
		invokeReply: &pb.MxCommandReply{
			SessionId: "session-1",
			Kind:      pb.MxCommandKind_MX_COMMAND_KIND_WRITE_BULK,
			ProtocolStatus: &pb.ProtocolStatus{
				Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
			},
			Payload: &pb.MxCommandReply_WriteBulk{
				WriteBulk: &pb.BulkWriteReply{
					Results: []*pb.BulkWriteResult{
						{ItemHandle: 10, WasSuccessful: true},
						{ItemHandle: 11, WasSuccessful: true},
					},
				},
			},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	entries := []*WriteBulkEntry{
		{ItemHandle: 10, Value: Int32Value(7), UserId: 100},
		{ItemHandle: 11, Value: Int32Value(8), UserId: 100},
	}
	results, err := session.WriteBulk(context.Background(), 12, entries)
	if err != nil {
		t.Fatalf("WriteBulk() error = %v", err)
	}
	if len(results) != 2 {
		t.Fatalf("results len = %d, want 2", len(results))
	}
	req := fake.invokeRequest
	if req.GetCommand().GetKind() != pb.MxCommandKind_MX_COMMAND_KIND_WRITE_BULK {
		t.Fatalf("command kind = %s", req.GetCommand().GetKind())
	}
	if got := req.GetCommand().GetWriteBulk().GetEntries(); len(got) != 2 {
		t.Fatalf("entry count = %d, want 2", len(got))
	}
}

func TestWriteBulkRejectsNilEntries(t *testing.T) {
	fake := &fakeGatewayServer{}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	if _, err := session.WriteBulk(context.Background(), 12, nil); err == nil {
		t.Fatal("WriteBulk(nil) returned no error")
	}
	if _, err := session.Write2Bulk(context.Background(), 12, nil); err == nil {
		t.Fatal("Write2Bulk(nil) returned no error")
	}
	if _, err := session.WriteSecuredBulk(context.Background(), 12, nil); err == nil {
		t.Fatal("WriteSecuredBulk(nil) returned no error")
	}
	if _, err := session.WriteSecured2Bulk(context.Background(), 12, nil); err == nil {
		t.Fatal("WriteSecured2Bulk(nil) returned no error")
	}
	if _, err := session.ReadBulk(context.Background(), 12, nil, 0); err == nil {
		t.Fatal("ReadBulk(nil) returned no error")
	}
}

func TestBulkMethodsShortCircuitOnEmptySliceWithoutRoundTrip(t *testing.T) {
	fake := &fakeGatewayServer{
		invokeReply: &pb.MxCommandReply{
			ProtocolStatus: &pb.ProtocolStatus{
				Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
			},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	results, err := session.WriteBulk(context.Background(), 12, []*WriteBulkEntry{})
	if err != nil {
		t.Fatalf("WriteBulk(empty) error = %v", err)
	}
	if len(results) != 0 {
		t.Fatalf("WriteBulk(empty) results len = %d, want 0", len(results))
	}
	if fake.invokeRequest != nil {
		t.Fatal("WriteBulk(empty) sent a round trip; expected short-circuit")
	}

	results2, err := session.Write2Bulk(context.Background(), 12, []*Write2BulkEntry{})
	if err != nil {
		t.Fatalf("Write2Bulk(empty) error = %v", err)
	}
	if len(results2) != 0 {
		t.Fatalf("Write2Bulk(empty) results len = %d, want 0", len(results2))
	}
	if fake.invokeRequest != nil {
		t.Fatal("Write2Bulk(empty) sent a round trip; expected short-circuit")
	}

	results3, err := session.WriteSecuredBulk(context.Background(), 12, []*WriteSecuredBulkEntry{})
	if err != nil {
		t.Fatalf("WriteSecuredBulk(empty) error = %v", err)
	}
	if len(results3) != 0 {
		t.Fatalf("WriteSecuredBulk(empty) results len = %d, want 0", len(results3))
	}
	if fake.invokeRequest != nil {
		t.Fatal("WriteSecuredBulk(empty) sent a round trip; expected short-circuit")
	}

	results4, err := session.WriteSecured2Bulk(context.Background(), 12, []*WriteSecured2BulkEntry{})
	if err != nil {
		t.Fatalf("WriteSecured2Bulk(empty) error = %v", err)
	}
	if len(results4) != 0 {
		t.Fatalf("WriteSecured2Bulk(empty) results len = %d, want 0", len(results4))
	}
	if fake.invokeRequest != nil {
		t.Fatal("WriteSecured2Bulk(empty) sent a round trip; expected short-circuit")
	}

	readResults, err := session.ReadBulk(context.Background(), 12, []string{}, 0)
	if err != nil {
		t.Fatalf("ReadBulk(empty) error = %v", err)
	}
	if len(readResults) != 0 {
		t.Fatalf("ReadBulk(empty) results len = %d, want 0", len(readResults))
	}
	if fake.invokeRequest != nil {
		t.Fatal("ReadBulk(empty) sent a round trip; expected short-circuit")
	}
}

func TestWrite2BuildsCommandWithTimestampAndReturnsNoError(t *testing.T) {
	fake := &fakeGatewayServer{
		invokeReply: &pb.MxCommandReply{
			SessionId: "session-1",
			Kind:      pb.MxCommandKind_MX_COMMAND_KIND_WRITE2,
			ProtocolStatus: &pb.ProtocolStatus{
				Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
			},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	val := Int32Value(99)
	ts := Int32Value(77)
	err := session.Write2(context.Background(), 12, 34, val, ts, 100)
	if err != nil {
		t.Fatalf("Write2() error = %v", err)
	}

	req := fake.invokeRequest
	if req.GetCommand().GetKind() != pb.MxCommandKind_MX_COMMAND_KIND_WRITE2 {
		t.Fatalf("command kind = %s, want WRITE2", req.GetCommand().GetKind())
	}
	w2 := req.GetCommand().GetWrite2()
	if w2.GetServerHandle() != 12 {
		t.Fatalf("server handle = %d, want 12", w2.GetServerHandle())
	}
	if w2.GetItemHandle() != 34 {
		t.Fatalf("item handle = %d, want 34", w2.GetItemHandle())
	}
	if w2.GetValue().GetInt32Value() != 99 {
		t.Fatalf("value int32 = %d, want 99", w2.GetValue().GetInt32Value())
	}
	if w2.GetTimestampValue().GetInt32Value() != 77 {
		t.Fatalf("timestamp value int32 = %d, want 77", w2.GetTimestampValue().GetInt32Value())
	}
	if w2.GetUserId() != 100 {
		t.Fatalf("user id = %d, want 100", w2.GetUserId())
	}
}

func TestWrite2RawReturnsRawReply(t *testing.T) {
	fake := &fakeGatewayServer{
		invokeReply: &pb.MxCommandReply{
			SessionId: "session-1",
			Kind:      pb.MxCommandKind_MX_COMMAND_KIND_WRITE2,
			ProtocolStatus: &pb.ProtocolStatus{
				Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
			},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	reply, err := session.Write2Raw(context.Background(), 12, 34, Int32Value(1), Int32Value(0), 0)
	if err != nil {
		t.Fatalf("Write2Raw() error = %v", err)
	}
	if reply == nil {
		t.Fatal("Write2Raw() returned nil reply")
	}
	if reply.GetKind() != pb.MxCommandKind_MX_COMMAND_KIND_WRITE2 {
		t.Fatalf("reply kind = %s, want WRITE2", reply.GetKind())
	}
}

func TestWrite2RejectsNilValue(t *testing.T) {
	fake := &fakeGatewayServer{}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	if err := session.Write2(context.Background(), 12, 34, nil, Int32Value(0), 0); err == nil {
		t.Fatal("Write2(nil value) returned no error")
	}
	if err := session.Write2(context.Background(), 12, 34, Int32Value(1), nil, 0); err == nil {
		t.Fatal("Write2(nil timestampValue) returned no error")
	}
}

func TestReadBulkForwardsTimeoutAndUnpacksCachedFlag(t *testing.T) {
	fake := &fakeGatewayServer{
		invokeReply: &pb.MxCommandReply{
			SessionId: "session-1",
			Kind:      pb.MxCommandKind_MX_COMMAND_KIND_READ_BULK,
			ProtocolStatus: &pb.ProtocolStatus{
				Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
			},
			Payload: &pb.MxCommandReply_ReadBulk{
				ReadBulk: &pb.BulkReadReply{
					Results: []*pb.BulkReadResult{
						{TagAddress: "Tank01.Level", WasSuccessful: true, WasCached: true},
						{TagAddress: "Tank02.Level", WasSuccessful: true, WasCached: false},
					},
				},
			},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	results, err := session.ReadBulk(context.Background(), 12, []string{"Tank01.Level", "Tank02.Level"}, 250*time.Millisecond)
	if err != nil {
		t.Fatalf("ReadBulk() error = %v", err)
	}
	if len(results) != 2 {
		t.Fatalf("results len = %d, want 2", len(results))
	}
	if !results[0].GetWasCached() || results[1].GetWasCached() {
		t.Fatalf("WasCached flags = [%v %v], want [true false]", results[0].GetWasCached(), results[1].GetWasCached())
	}
	req := fake.invokeRequest
	if req.GetCommand().GetKind() != pb.MxCommandKind_MX_COMMAND_KIND_READ_BULK {
		t.Fatalf("command kind = %s", req.GetCommand().GetKind())
	}
	if got := req.GetCommand().GetReadBulk().GetTimeoutMs(); got != 250 {
		t.Fatalf("timeout ms = %d, want 250", got)
	}
}

func TestReadBulkSaturatesTimeoutAboveMaxUint32(t *testing.T) {
	fake := &fakeGatewayServer{
		invokeReply: &pb.MxCommandReply{
			SessionId: "session-1",
			Kind:      pb.MxCommandKind_MX_COMMAND_KIND_READ_BULK,
			ProtocolStatus: &pb.ProtocolStatus{
				Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
			},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	// 100 days in milliseconds exceeds MaxUint32 (~49.7 days).
	hugeTimeout := 100 * 24 * time.Hour
	_, err := session.ReadBulk(context.Background(), 12, []string{"Tank01.Level"}, hugeTimeout)
	if err != nil {
		t.Fatalf("ReadBulk() error = %v", err)
	}
	got := fake.invokeRequest.GetCommand().GetReadBulk().GetTimeoutMs()
	if got != ^uint32(0) {
		t.Fatalf("timeout ms = %d, want %d (MaxUint32)", got, ^uint32(0))
	}
}

func TestInvokeReturnsTypedMxAccessErrorWithRawReply(t *testing.T) {
	hresult := int32(-2147467259)
	fake := &fakeGatewayServer{
		invokeReply: &pb.MxCommandReply{
			SessionId:         "session-1",
			Kind:              pb.MxCommandKind_MX_COMMAND_KIND_ADVISE,
			Hresult:           &hresult,
			DiagnosticMessage: "native failure",
			ProtocolStatus:    &pb.ProtocolStatus{Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_MXACCESS_FAILURE, Message: "MXAccess failed"},
			Statuses:          []*pb.MxStatusProxy{{Success: 0, DiagnosticText: "failed"}},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	err := session.Advise(context.Background(), 12, 34)

	var mxErr *MxAccessError
	if !errors.As(err, &mxErr) {
		t.Fatalf("error %T does not support errors.As(*MxAccessError)", err)
	}
	if mxErr.Reply.GetHresult() != hresult {
		t.Fatalf("raw reply HRESULT = %d, want %d", mxErr.Reply.GetHresult(), hresult)
	}
	var commandErr *CommandError
	if !errors.As(err, &commandErr) {
		t.Fatalf("error %T does not support errors.As(*CommandError)", err)
	}
	if commandErr.Reply.GetDiagnosticMessage() != "native failure" {
		t.Fatalf("raw diagnostic = %q", commandErr.Reply.GetDiagnosticMessage())
	}
}

func newBufconnClient(t *testing.T, fake *fakeGatewayServer) (*Client, func()) {
	t.Helper()

	listener := bufconn.Listen(bufSize)
	server := grpc.NewServer()
	pb.RegisterMxAccessGatewayServer(server, fake)
	go func() {
		if err := server.Serve(listener); err != nil && !errors.Is(err, grpc.ErrServerStopped) {
			t.Errorf("bufconn server failed: %v", err)
		}
	}()

	dialer := func(ctx context.Context, _ string) (net.Conn, error) {
		return listener.DialContext(ctx)
	}
	client, err := Dial(context.Background(), Options{
		Endpoint:  "bufnet",
		APIKey:    "test-api-key",
		Plaintext: true,
		DialOptions: []grpc.DialOption{
			grpc.WithContextDialer(dialer),
		},
	})
	if err != nil {
		t.Fatalf("Dial() error = %v", err)
	}

	return client, func() {
		client.Close()
		server.Stop()
		listener.Close()
	}
}

type fakeGatewayServer struct {
	pb.UnimplementedMxAccessGatewayServer

	openReply         *pb.OpenSessionReply
	openAuth          string
	streamAuth        string
	streamStarted     chan struct{}
	streamDone        chan struct{}
	streamEventCount  int
	streamReplayGap   *pb.ReplayGap
	streamTerminalErr error
	invokeReply       *pb.MxCommandReply
	invokeRequest     *pb.MxCommandRequest
}

func (s *fakeGatewayServer) OpenSession(ctx context.Context, req *pb.OpenSessionRequest) (*pb.OpenSessionReply, error) {
	s.openAuth = authorizationFromContext(ctx)
	if s.openReply != nil {
		return s.openReply, nil
	}
	return &pb.OpenSessionReply{
		SessionId: "session-1",
		ProtocolStatus: &pb.ProtocolStatus{
			Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
		},
	}, nil
}

func (s *fakeGatewayServer) CloseSession(ctx context.Context, req *pb.CloseSessionRequest) (*pb.CloseSessionReply, error) {
	return &pb.CloseSessionReply{
		SessionId: req.GetSessionId(),
		ProtocolStatus: &pb.ProtocolStatus{
			Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
		},
	}, nil
}

func (s *fakeGatewayServer) Invoke(ctx context.Context, req *pb.MxCommandRequest) (*pb.MxCommandReply, error) {
	s.invokeRequest = req
	if s.invokeReply != nil {
		return s.invokeReply, nil
	}
	return &pb.MxCommandReply{
		SessionId: req.GetSessionId(),
		Kind:      req.GetCommand().GetKind(),
		ProtocolStatus: &pb.ProtocolStatus{
			Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
		},
	}, nil
}

func (s *fakeGatewayServer) StreamEvents(req *pb.StreamEventsRequest, stream grpc.ServerStreamingServer[pb.MxEvent]) error {
	s.streamAuth = authorizationFromContext(stream.Context())
	if s.streamDone != nil {
		defer close(s.streamDone)
	}
	if s.streamStarted != nil {
		close(s.streamStarted)
	}
	if s.streamReplayGap != nil {
		// Emit the reconnect-replay gap sentinel at the head of the resumed
		// stream: family UNSPECIFIED, body unset, only replay_gap populated.
		if err := stream.Send(&pb.MxEvent{
			SessionId: req.GetSessionId(),
			ReplayGap: s.streamReplayGap,
		}); err != nil {
			return err
		}
	}
	eventCount := s.streamEventCount
	if eventCount == 0 {
		eventCount = 1
	}
	for sequence := 1; sequence <= eventCount; sequence++ {
		if err := stream.Send(&pb.MxEvent{
			SessionId:      req.GetSessionId(),
			Family:         pb.MxEventFamily_MX_EVENT_FAMILY_ON_DATA_CHANGE,
			WorkerSequence: uint64(sequence),
		}); err != nil {
			return err
		}
	}
	if s.streamTerminalErr != nil {
		// Return a genuine terminal stream error immediately after sending the
		// events, without waiting on the client to cancel. This exercises the
		// Recv-error path while the client's result buffer is still full.
		return s.streamTerminalErr
	}
	<-stream.Context().Done()
	return io.EOF
}

func authorizationFromContext(ctx context.Context) string {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return ""
	}
	values := md.Get(authorizationHeader)
	if len(values) == 0 {
		return ""
	}
	return values[0]
}

// ---------------------------------------------------------------------------
// WriteArrayElements / buildSparseArrayValue unit tests
// ---------------------------------------------------------------------------

func TestBuildSparseArrayValueSetsSparseOneof(t *testing.T) {
	elements := map[uint32]*MxValue{
		2: Int32Value(99),
		0: Int32Value(10),
	}
	v := buildSparseArrayValue(DataTypeInteger, 5, elements)

	sa, ok := v.Kind.(*pb.MxValue_SparseArrayValue)
	if !ok {
		t.Fatalf("Kind is %T, want *pb.MxValue_SparseArrayValue", v.Kind)
	}
	got := sa.SparseArrayValue
	if got.GetElementDataType() != DataTypeInteger {
		t.Errorf("ElementDataType = %v, want DataTypeInteger", got.GetElementDataType())
	}
	if got.GetTotalLength() != 5 {
		t.Errorf("TotalLength = %d, want 5", got.GetTotalLength())
	}
	if len(got.GetElements()) != 2 {
		t.Fatalf("len(Elements) = %d, want 2", len(got.GetElements()))
	}
	// Elements must be sorted by index (ascending).
	if got.GetElements()[0].GetIndex() != 0 {
		t.Errorf("Elements[0].Index = %d, want 0", got.GetElements()[0].GetIndex())
	}
	if got.GetElements()[0].GetValue().GetInt32Value() != 10 {
		t.Errorf("Elements[0].Value = %v, want 10", got.GetElements()[0].GetValue())
	}
	if got.GetElements()[1].GetIndex() != 2 {
		t.Errorf("Elements[1].Index = %d, want 2", got.GetElements()[1].GetIndex())
	}
	if got.GetElements()[1].GetValue().GetInt32Value() != 99 {
		t.Errorf("Elements[1].Value = %v, want 99", got.GetElements()[1].GetValue())
	}
}

func TestBuildSparseArrayValueEmptyMapProducesEmptyElements(t *testing.T) {
	v := buildSparseArrayValue(DataTypeBoolean, 4, map[uint32]*MxValue{})

	sa, ok := v.Kind.(*pb.MxValue_SparseArrayValue)
	if !ok {
		t.Fatalf("Kind is %T, want *pb.MxValue_SparseArrayValue", v.Kind)
	}
	if len(sa.SparseArrayValue.GetElements()) != 0 {
		t.Errorf("len(Elements) = %d, want 0", len(sa.SparseArrayValue.GetElements()))
	}
	if sa.SparseArrayValue.GetTotalLength() != 4 {
		t.Errorf("TotalLength = %d, want 4", sa.SparseArrayValue.GetTotalLength())
	}
}

func TestWriteArrayElementsSendsWriteCommandWithSparseOneof(t *testing.T) {
	fake := &fakeGatewayServer{
		invokeReply: &pb.MxCommandReply{
			SessionId: "session-1",
			Kind:      pb.MxCommandKind_MX_COMMAND_KIND_WRITE,
			ProtocolStatus: &pb.ProtocolStatus{
				Code: pb.ProtocolStatusCode_PROTOCOL_STATUS_CODE_OK,
			},
		},
	}
	client, cleanup := newBufconnClient(t, fake)
	defer cleanup()
	session := NewSessionForID(client, "session-1")

	err := session.WriteArrayElements(
		context.Background(),
		1, 2,
		DataTypeFloat,
		10,
		map[uint32]*MxValue{3: FloatValue(1.5)},
		42,
	)
	if err != nil {
		t.Fatalf("WriteArrayElements() error = %v", err)
	}

	cmd := fake.invokeRequest.GetCommand()
	if cmd.GetKind() != pb.MxCommandKind_MX_COMMAND_KIND_WRITE {
		t.Fatalf("command kind = %s, want WRITE", cmd.GetKind())
	}
	w := cmd.GetWrite()
	if w.GetServerHandle() != 1 {
		t.Fatalf("server handle = %d, want 1", w.GetServerHandle())
	}
	if w.GetItemHandle() != 2 {
		t.Fatalf("item handle = %d, want 2", w.GetItemHandle())
	}
	if w.GetUserId() == 0 {
		t.Fatal("user id = 0, want non-zero")
	}
	if w.GetUserId() != 42 {
		t.Fatalf("user id = %d, want 42", w.GetUserId())
	}
	val := w.GetValue()
	sa, ok := val.Kind.(*pb.MxValue_SparseArrayValue)
	if !ok {
		t.Fatalf("value kind is %T, want *pb.MxValue_SparseArrayValue", val.Kind)
	}
	if sa.SparseArrayValue.GetTotalLength() != 10 {
		t.Errorf("TotalLength = %d, want 10", sa.SparseArrayValue.GetTotalLength())
	}
	if sa.SparseArrayValue.GetElementDataType() != DataTypeFloat {
		t.Errorf("ElementDataType = %v, want DataTypeFloat", sa.SparseArrayValue.GetElementDataType())
	}
	if len(sa.SparseArrayValue.GetElements()) != 1 {
		t.Fatalf("len(Elements) = %d, want 1", len(sa.SparseArrayValue.GetElements()))
	}
	elem := sa.SparseArrayValue.GetElements()[0]
	if elem.GetIndex() != 3 {
		t.Errorf("element index = %d, want 3", elem.GetIndex())
	}
	if elem.GetValue().GetFloatValue() != 1.5 {
		t.Errorf("element float value = %v, want 1.5", elem.GetValue().GetFloatValue())
	}
}
