package mtprotoedge

import (
	"context"
	"errors"
	"reflect"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/iamxvbaba/td/bin"
	"github.com/iamxvbaba/td/clock"
	"github.com/iamxvbaba/td/proto"
	"github.com/iamxvbaba/td/tg"
	"go.uber.org/zap/zaptest"

	"github.com/iamxvbaba/td/tlprofile"
	appfiles "telesrv/internal/app/files"
	"telesrv/internal/rpc"
)

func TestLayerRPCAdmissionMaterializationConstants(t *testing.T) {
	constructors := make([]reflect.Type, 0, len(tg.TypesConstructorMap()))
	for _, constructor := range tg.TypesConstructorMap() {
		if typ := reflect.TypeOf(constructor()); typ != nil {
			constructors = append(constructors, typ)
		}
	}
	seen := make(map[reflect.Type]struct{})
	var (
		maxSize uintptr
		maxType reflect.Type
		visit   func(reflect.Type)
	)
	visit = func(typ reflect.Type) {
		if typ == nil {
			return
		}
		if _, ok := seen[typ]; ok {
			return
		}
		seen[typ] = struct{}{}
		switch typ.Kind() {
		case reflect.Pointer, reflect.Slice, reflect.Array:
			visit(typ.Elem())
		case reflect.Interface:
			for _, candidate := range constructors {
				// invokeWithLayer/query uses the deliberately broad bin.Object
				// interface. Exact admission restricts that slot to generated RPC
				// methods, so only request constructors are reachable there.
				broad := typ.PkgPath() != "github.com/iamxvbaba/td/tg"
				if candidate.Implements(typ) && (!broad || strings.HasSuffix(candidate.Elem().Name(), "Request")) {
					visit(candidate)
				}
			}
		case reflect.Struct:
			if typ.PkgPath() == "github.com/iamxvbaba/td/tg" && typ.Size() > maxSize {
				maxSize, maxType = typ.Size(), typ
			}
			for i := 0; i < typ.NumField(); i++ {
				visit(typ.Field(i).Type)
			}
		}
	}
	for _, typ := range constructors {
		if typ.Kind() == reflect.Pointer && strings.HasSuffix(typ.Elem().Name(), "Request") {
			visit(typ)
		}
	}
	if maxSize > layerRPCAdmissionStaticObjectBytes {
		t.Fatalf("request-reachable generated TL object %v is %d bytes, exceeds admission ceiling %d", maxType, maxSize, layerRPCAdmissionStaticObjectBytes)
	}
	if layerRPCAdmissionGraphSlack != layerRPCAdmissionStaticObjectBytes*inboundLayerDecodeLimits.MaxDepth {
		t.Fatalf("graph slack %d does not cover %d bytes across decode depth %d", layerRPCAdmissionGraphSlack, layerRPCAdmissionStaticObjectBytes, inboundLayerDecodeLimits.MaxDepth)
	}
	// Preserve enough room for the maximum upload payload plus any supported
	// transparent wrapper/client-info envelope on a default connection.
	if got := layerRPCAdmissionReservationSize(appfiles.MaxUploadPartBytes + (32 << 10)); got > maxInflightRPCBytes {
		t.Fatalf("largest legal upload request charge = %d, exceeds default connection budget %d", got, maxInflightRPCBytes)
	}
	maxInt := int(^uint(0) >> 1)
	if got := layerRPCAdmissionReservationSize(maxInt); got != maxInt {
		t.Fatalf("saturating charge = %d, want max int %d", got, maxInt)
	}
	if got := layerRPCAdmissionReservationSize(-1); got != layerRPCAdmissionGraphSlack {
		t.Fatalf("negative wire charge = %d, want fixed slack %d", got, layerRPCAdmissionGraphSlack)
	}
}

type countingLayerRPCAdmission struct {
	LayerRPCHandler
	decodeCalls atomic.Int32
}

type failingReplayLayerRPC struct {
	LayerRPCHandler
	err error
}

func (h *failingReplayLayerRPC) PrepareAdmittedReplay(
	context.Context,
	[8]byte,
	int64,
	int64,
	uint64,
	tlprofile.Admission,
) (func() error, error) {
	return nil, h.err
}

func (h *countingLayerRPCAdmission) AdmitLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
	h.decodeCalls.Add(1)
	return h.LayerRPCHandler.AdmitLayer(profile, b, limits)
}

func (h *countingLayerRPCAdmission) AdmitLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
	h.decodeCalls.Add(1)
	if admitter, ok := h.LayerRPCHandler.(LayerRPCOptionsAdmitter); ok {
		return admitter.AdmitLayerWithOptions(profile, b, options)
	}
	return h.LayerRPCHandler.AdmitLayer(profile, b, options.Limits)
}

func (h *countingLayerRPCAdmission) AdmitUnprofiled(b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
	h.decodeCalls.Add(1)
	return h.LayerRPCHandler.AdmitUnprofiled(b, limits)
}

func (h *countingLayerRPCAdmission) AdmitUnprofiledWithOptions(b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
	h.decodeCalls.Add(1)
	if admitter, ok := h.LayerRPCHandler.(LayerRPCOptionsAdmitter); ok {
		return admitter.AdmitUnprofiledWithOptions(b, options)
	}
	return h.LayerRPCHandler.AdmitUnprofiled(b, options.Limits)
}

func TestLayerRPCAdmissionCapacityRejectsBeforeDecoder(t *testing.T) {
	for _, test := range []struct {
		name       string
		fillGlobal bool
	}{
		{name: "connection_queue"},
		{name: "global_task", fillGlobal: true},
	} {
		t.Run(test.name, func(t *testing.T) {
			router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
			counting := &countingLayerRPCAdmission{LayerRPCHandler: router}
			s := New(Options{DC: 2, LayerRPC: counting})
			scheduler := newInboundRPCScheduler(1, 1, 1<<30)
			s.rpcScheduler = scheduler

			target := &Conn{authKeyID: [8]byte{8, 1}, sessionID: 81, metrics: NopMetrics{}}
			target.startInboundRPCScheduler(scheduler, 1, 1, time.Second)
			holder := target
			if test.fillGlobal {
				holder = &Conn{authKeyID: [8]byte{8, 2}, sessionID: 82, metrics: NopMetrics{}}
				holder.startInboundRPCScheduler(scheduler, 1, 1, time.Second)
			}
			occupied, err := holder.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "occupied", size: 1}})
			if err != nil {
				t.Fatal(err)
			}
			defer occupied.abort()

			body := exactLayerRPCBody(t, &tg.InvokeWithLayerRequest{Layer: 225, Query: &tg.HelpGetConfigRequest{}})
			plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
			defer plan.close()
			if err := s.prepareInboundLayerRPCBatch(context.Background(), target, plan); err != nil {
				t.Fatal(err)
			}
			if got := counting.decodeCalls.Load(); got != 0 {
				t.Fatalf("typed decoder entered %d times after capacity rejection", got)
			}
			if plan.items[0].kind != inboundItemCapacityError || plan.rpcReservation != nil || len(plan.rpcTasks) != 0 {
				t.Fatalf("rejected plan = kind:%d reservation:%v tasks:%d", plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks))
			}
		})
	}
}

func TestLayerRPCAdmissionExpandsTDLibNestedGZIPUnderTransferredBudget(t *testing.T) {
	router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
	s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
	c := &Conn{authKeyID: [8]byte{8, 21}, sessionID: 821, metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
	defer func() {
		c.closeInboundRPCScheduler()
		s.rpcScheduler.stop(time.Second)
	}()

	body, expandedBytes := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetConfigRequest{})
	plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
	defer plan.close()
	if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
		t.Fatal(err)
	}
	if plan.items[0].kind != inboundItemRPC || plan.rpcReservation == nil || len(plan.rpcTasks) != 1 {
		t.Fatalf("admitted nested gzip plan = kind:%d reservation:%v tasks:%d", plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks))
	}
	wantCharge := int64(layerRPCAdmissionReservationSize(len(body) + expandedBytes))
	if got := c.inflightRPCBytes.Load(); got != wantCharge {
		t.Fatalf("nested gzip connection charge = %d, want %d", got, wantCharge)
	}
	if got := plan.rpcReservation.entries[0].size; int64(got) != wantCharge {
		t.Fatalf("nested gzip reservation charge = %d, want %d", got, wantCharge)
	}
	if got := plan.gzipExpandedBytes; got != expandedBytes {
		t.Fatalf("nested gzip cumulative expansion = %d, want %d", got, expandedBytes)
	}
	if got := s.frameBudget.usedBytes(); got != 0 {
		t.Fatalf("nested gzip temporary frame budget retained after materialization: %d", got)
	}
}

func TestLayerRPCAdmissionNestedGZIPGrowFailureRejectsWholeBatch(t *testing.T) {
	router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
	s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
	first, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetConfigRequest{})
	second, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetNearestDCRequest{})
	initialCharge := int64(layerRPCAdmissionReservationSize(len(first)) + layerRPCAdmissionReservationSize(len(second)))
	s.rpcScheduler = newInboundRPCScheduler(1, 4, initialCharge)
	c := &Conn{authKeyID: [8]byte{8, 22}, sessionID: 822, metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
	defer func() {
		c.closeInboundRPCScheduler()
		s.rpcScheduler.stop(time.Second)
	}()

	plan := &inboundPlan{items: []inboundItem{
		{kind: inboundItemRPC, msgID: 100, body: first},
		{kind: inboundItemRPC, msgID: 104, body: second},
	}}
	defer plan.close()
	if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
		t.Fatal(err)
	}
	for index := range plan.items {
		if plan.items[index].kind != inboundItemCapacityError {
			t.Fatalf("item %d kind = %d, want capacity error", index, plan.items[index].kind)
		}
	}
	if plan.rpcReservation != nil || len(plan.rpcTasks) != 0 {
		t.Fatalf("capacity plan retained reservation/tasks = %v/%d", plan.rpcReservation != nil, len(plan.rpcTasks))
	}
	if got := c.inflightRPCBytes.Load(); got != 0 {
		t.Fatalf("grow failure leaked connection charge %d", got)
	}
	if tasks, bytes := s.rpcScheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
		t.Fatalf("grow failure leaked global budget %d/%d", tasks, bytes)
	}
	if got := s.frameBudget.usedBytes(); got != 0 {
		t.Fatalf("grow failure leaked temporary frame budget %d", got)
	}
}

func TestLayerRPCAdmissionNestedGZIPSiblingsShareFrameExpansionLimit(t *testing.T) {
	router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
	s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
	c := &Conn{authKeyID: [8]byte{8, 23}, sessionID: 823, metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
	defer func() {
		c.closeInboundRPCScheduler()
		s.rpcScheduler.stop(time.Second)
	}()

	first, expandedBytes := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetConfigRequest{})
	second, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetNearestDCRequest{})
	plan := &inboundPlan{
		gzipExpandedBytes: maxDispatchExpandedBytes - expandedBytes,
		items: []inboundItem{
			{kind: inboundItemRPC, msgID: 100, body: first},
			{kind: inboundItemRPC, msgID: 104, body: second},
		},
	}
	defer plan.close()
	if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
		t.Fatal(err)
	}
	for index := range plan.items {
		if plan.items[index].kind != inboundItemCapacityError {
			t.Fatalf("item %d kind = %d, want capacity error", index, plan.items[index].kind)
		}
	}
	if got := plan.gzipExpandedBytes; got != maxDispatchExpandedBytes {
		t.Fatalf("shared cumulative expansion = %d, want %d", got, maxDispatchExpandedBytes)
	}
	if got := c.inflightRPCBytes.Load(); got != 0 {
		t.Fatalf("shared-limit rejection leaked connection charge %d", got)
	}
	if got := s.frameBudget.usedBytes(); got != 0 {
		t.Fatalf("shared-limit rejection leaked temporary frame budget %d", got)
	}
}

func TestLayerRPCAdmissionNestedGZIPReDecodeReusesMaterializationCharge(t *testing.T) {
	handler := newAdmissionOnlyLayerRPC()
	s := New(Options{DC: 2, LayerRPC: handler, Logger: zaptest.NewLogger(t)})
	s.rpcResults = newRPCResultCacheWithFlightLimit(time.Now, 8)
	scheduler := newInboundRPCScheduler(1, 4, 1<<30)
	s.rpcScheduler = scheduler
	authKeyID := [8]byte{8, 24}
	const sessionID = int64(824)
	c225 := &Conn{authKeyID: authKeyID, sessionID: sessionID, metrics: NopMetrics{}}
	c227 := &Conn{authKeyID: authKeyID, sessionID: sessionID, metrics: NopMetrics{}}
	c225.startInboundRPCScheduler(scheduler, 1, 2, time.Second)
	c227.startInboundRPCScheduler(scheduler, 1, 2, time.Second)
	defer func() {
		c225.closeInboundRPCScheduler()
		c227.closeInboundRPCScheduler()
		scheduler.stop(time.Second)
	}()

	terminal := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.MessagesGetHistoryRequest{
		Peer: &tg.InputPeerSelf{}, Limit: 1,
	})
	body := exactLayerRPCBody(t, &tg.InvokeWithoutUpdatesRequest{Query: &proto.GZIP{Data: terminal}})
	initialCharge := layerRPCAdmissionReservationSize(len(body))
	reservation225, err := c225.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "messages.getHistory", size: initialCharge}})
	if err != nil {
		t.Fatal(err)
	}
	defer reservation225.abort()
	reservation227, err := c227.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "messages.getHistory", size: initialCharge}})
	if err != nil {
		t.Fatal(err)
	}
	defer reservation227.abort()

	plan225 := &inboundPlan{}
	budget225 := &layerRPCGZIPExpansionBudget{
		server: s, plan: plan225, reservation: reservation225,
		baseSourceBytes: len(body), chargedSourceBytes: len(body),
	}
	options225 := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits, ExpandGZIP: budget225.expand}
	item225 := inboundItem{msgID: 100, body: body}
	item225.admitted, item225.method, err = s.decodeInboundLayerRPCWithOptions(
		LayerProfileSnapshot{Profile: tlprofile.Profile225, Origin: LayerProfileInherited}, body, options225,
	)
	if err != nil {
		t.Fatal(err)
	}

	plan227 := &inboundPlan{}
	budget227 := &layerRPCGZIPExpansionBudget{
		server: s, plan: plan227, reservation: reservation227,
		baseSourceBytes: len(body), chargedSourceBytes: len(body),
	}
	options227 := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits, ExpandGZIP: budget227.expand}
	item227 := inboundItem{msgID: 100, body: body}
	item227.admitted, item227.method, err = s.decodeInboundLayerRPCWithOptions(
		LayerProfileSnapshot{Profile: tlprofile.Profile227, Origin: LayerProfileInherited}, body, options227,
	)
	if err != nil {
		t.Fatal(err)
	}
	if item225.admitted.Prepared().Identity() == item227.admitted.Prepared().Identity() {
		t.Fatal("test request identity is invariant; need authoritative-profile re-decode")
	}

	winner, err := s.acquireAdmittedLayerRPC(c225, &item225, nil, options225, budget225)
	if err != nil || winner.state != rpcResultAcquireOwner || winner.owner == nil {
		t.Fatalf("winner = state:%d err:%v", winner.state, err)
	}
	defer winner.owner.Abort()
	loser, err := s.acquireAdmittedLayerRPC(c227, &item227, nil, options227, budget227)
	if err != nil || loser.state != rpcResultAcquirePending {
		t.Fatalf("loser = state:%d err:%v", loser.state, err)
	}
	if got := item227.admitted.Call().Profile(); got != tlprofile.Profile225 {
		t.Fatalf("loser re-admitted profile = %d, want 225", got)
	}
	wantCharge := layerRPCAdmissionReservationSize(len(body) + len(terminal))
	if got := reservation227.entries[0].size; got != wantCharge {
		t.Fatalf("re-decode reservation charge = %d, want single-graph maximum %d", got, wantCharge)
	}
	if got := plan227.gzipExpandedBytes; got != 2*len(terminal) {
		t.Fatalf("re-decode cumulative work = %d, want %d", got, 2*len(terminal))
	}
	if got := s.frameBudget.usedBytes(); got != 0 {
		t.Fatalf("re-decode leaked temporary frame budget %d", got)
	}
}

func TestLayerRPCAdmissionTransfersOriginalReservationToFreshOwner(t *testing.T) {
	router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
	s := New(Options{DC: 2, LayerRPC: router})
	c := &Conn{authKeyID: [8]byte{8, 3}, sessionID: 83, metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
	if err := c.FreezeLayerProfile(tlprofile.Profile225); err != nil {
		t.Fatal(err)
	}

	bad := make([]byte, bin.Word)
	bad[0], bad[1], bad[2], bad[3] = 0x04, 0x03, 0x02, 0x01
	fresh := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
	plan := &inboundPlan{items: []inboundItem{
		{kind: inboundItemRPC, msgID: 100, body: bad},
		{kind: inboundItemRPC, msgID: 104, body: fresh},
	}}
	defer plan.close()
	if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
		t.Fatal(err)
	}
	if plan.items[0].kind != inboundItemRPCAdmissionError || len(plan.rpcTasks) != 1 || plan.rpcReservation == nil {
		t.Fatalf("classified plan = bad:%d tasks:%d reservation:%v", plan.items[0].kind, len(plan.rpcTasks), plan.rpcReservation != nil)
	}
	wantCharge := int64(layerRPCAdmissionReservationSize(len(fresh)))
	if got := c.inflightRPCBytes.Load(); got != wantCharge {
		t.Fatalf("connection retained bytes = %d, want %d", got, wantCharge)
	}
	s.rpcScheduler.budgetMu.Lock()
	globalTasks, globalBytes := s.rpcScheduler.tasks, s.rpcScheduler.bytes
	s.rpcScheduler.budgetMu.Unlock()
	if globalTasks != 1 || globalBytes != wantCharge {
		t.Fatalf("global retained budget = %d/%d, want 1/%d", globalTasks, globalBytes, wantCharge)
	}
	plan.close()
	if got := c.inflightRPCBytes.Load(); got != 0 {
		t.Fatalf("plan abort leaked %d connection bytes", got)
	}
	s.rpcScheduler.budgetMu.Lock()
	globalTasks, globalBytes = s.rpcScheduler.tasks, s.rpcScheduler.bytes
	s.rpcScheduler.budgetMu.Unlock()
	if globalTasks != 0 || globalBytes != 0 {
		t.Fatalf("plan abort leaked global budget %d/%d", globalTasks, globalBytes)
	}
}

func tdlibNestedGZIPBody(t *testing.T, profile tlprofile.Profile, terminal bin.Object) ([]byte, int) {
	t.Helper()
	terminalWire := exactOutboundLayerRPCBody(t, profile, terminal)
	request := &tg.InvokeWithLayerRequest{
		Layer: int(profile),
		Query: &tg.InitConnectionRequest{
			APIID:          1,
			DeviceModel:    "android",
			SystemVersion:  "test",
			AppVersion:     "1.0",
			SystemLangCode: "en",
			LangPack:       "",
			LangCode:       "en",
			Query:          &proto.GZIP{Data: terminalWire},
		},
	}
	return exactLayerRPCBody(t, request), len(terminalWire)
}

func TestLayerRPCAdmissionPendingReplayReleasesProvisionalEntry(t *testing.T) {
	router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
	s := New(Options{DC: 2, LayerRPC: router})
	c := &Conn{authKeyID: [8]byte{8, 4}, sessionID: 84, metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
	if err := c.FreezeLayerProfile(tlprofile.Profile225); err != nil {
		t.Fatal(err)
	}

	pendingBody := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
	identityBuffer := &bin.Buffer{Buf: append([]byte(nil), pendingBody...)}
	pendingRequest, err := router.AdmitLayer(tlprofile.Profile225, identityBuffer, tlprofile.Limits{})
	if err != nil {
		t.Fatal(err)
	}
	pending, err := s.rpcResults.AcquireIdentified(c.authKeyID, c.sessionID, 100, pendingRequest.Prepared().Identity())
	if err != nil || pending.owner == nil {
		t.Fatalf("pending owner = %v, %v", pending.owner, err)
	}

	freshBody := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.HelpGetNearestDCRequest{})
	plan := &inboundPlan{items: []inboundItem{
		{kind: inboundItemRPC, msgID: 100, body: pendingBody},
		{kind: inboundItemRPC, msgID: 104, body: freshBody},
	}}
	defer plan.close()
	if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
		t.Fatal(err)
	}
	if plan.items[0].kind != inboundItemRewrappedRPC || len(plan.rewrapAliases) != 1 || len(plan.rpcTasks) != 1 {
		t.Fatalf("pending/fresh classification = kind:%d aliases:%d tasks:%d", plan.items[0].kind, len(plan.rewrapAliases), len(plan.rpcTasks))
	}
	wantCharge := int64(layerRPCAdmissionReservationSize(len(freshBody)))
	if got := c.inflightRPCBytes.Load(); got != wantCharge {
		t.Fatalf("pending replay retained bytes = %d, want only fresh %d", got, wantCharge)
	}
	plan.close()
	if !pending.owner.Abort() {
		t.Fatal("plan cleanup aborted the pre-existing pending replay owner")
	}
}

func TestLayerRPCAdmissionCompletedReplayReleasesWholeProvisionalBatch(t *testing.T) {
	router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
	s := New(Options{DC: 2, LayerRPC: router})
	c := &Conn{authKeyID: [8]byte{8, 8}, sessionID: 88, metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
	if err := c.FreezeLayerProfile(tlprofile.Profile225); err != nil {
		t.Fatal(err)
	}
	body := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
	identityBuffer := &bin.Buffer{Buf: append([]byte(nil), body...)}
	request, err := router.AdmitLayer(tlprofile.Profile225, identityBuffer, tlprofile.Limits{})
	if err != nil {
		t.Fatal(err)
	}
	claim, err := s.rpcResults.AcquireIdentified(c.authKeyID, c.sessionID, 100, request.Prepared().Identity())
	if err != nil || claim.owner == nil {
		t.Fatalf("completed replay owner = %v, %v", claim.owner, err)
	}
	if !claim.owner.CompleteExecution(true) {
		t.Fatal("complete replay business outcome failed")
	}
	s.rpcResults.Put(c.authKeyID, c.sessionID, 100, &encodedOutboundMessage{body: []byte{1, 2, 3, 4}})

	plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
	defer plan.close()
	if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
		t.Fatal(err)
	}
	if plan.items[0].kind != inboundItemReplayRPC || plan.rpcReservation != nil || len(plan.rpcTasks) != 0 {
		t.Fatalf("completed replay plan = kind:%d reservation:%v tasks:%d", plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks))
	}
	if got := c.inflightRPCBytes.Load(); got != 0 || c.rpcReserved != 0 {
		t.Fatalf("completed replay leaked connection budget bytes:%d tasks:%d", got, c.rpcReserved)
	}
	s.rpcScheduler.budgetMu.Lock()
	globalTasks, globalBytes := s.rpcScheduler.tasks, s.rpcScheduler.bytes
	s.rpcScheduler.budgetMu.Unlock()
	if globalTasks != 0 || globalBytes != 0 {
		t.Fatalf("completed replay leaked global budget %d/%d", globalTasks, globalBytes)
	}
}

func TestLayerRPCAdmissionReplayPreparationErrorIsNotSilentlyDelivered(t *testing.T) {
	router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
	prepareErr := errors.New("invalid replay wrapper metadata")
	s := New(Options{DC: 2, LayerRPC: &failingReplayLayerRPC{
		LayerRPCHandler: router,
		err:             prepareErr,
	}})
	c := &Conn{authKeyID: [8]byte{8, 9}, sessionID: 89, metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
	if err := c.FreezeLayerProfile(tlprofile.Profile225); err != nil {
		t.Fatal(err)
	}
	body := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
	identityBuffer := &bin.Buffer{Buf: append([]byte(nil), body...)}
	request, err := router.AdmitLayer(tlprofile.Profile225, identityBuffer, tlprofile.Limits{})
	if err != nil {
		t.Fatal(err)
	}
	claim, err := s.rpcResults.AcquireIdentified(c.authKeyID, c.sessionID, 100, request.Prepared().Identity())
	if err != nil || claim.owner == nil {
		t.Fatalf("completed replay owner = %v, %v", claim.owner, err)
	}
	if !claim.owner.CompleteExecution(true) {
		t.Fatal("complete replay business outcome failed")
	}
	s.rpcResults.Put(c.authKeyID, c.sessionID, 100, &encodedOutboundMessage{body: []byte{1, 2, 3, 4}})

	plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
	if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); !errors.Is(err, prepareErr) {
		plan.close()
		t.Fatalf("replay preparation error = %v, want %v", err, prepareErr)
	}
	if plan.items[0].kind == inboundItemReplayRPC {
		plan.close()
		t.Fatal("invalid replay metadata was converted into a deliverable cached result")
	}
	plan.close()
	if got := c.inflightRPCBytes.Load(); got != 0 || c.rpcReserved != 0 {
		t.Fatalf("failed replay preparation leaked connection budget bytes:%d tasks:%d", got, c.rpcReserved)
	}
	s.rpcScheduler.budgetMu.Lock()
	globalTasks, globalBytes := s.rpcScheduler.tasks, s.rpcScheduler.bytes
	s.rpcScheduler.budgetMu.Unlock()
	if globalTasks != 0 || globalBytes != 0 {
		t.Fatalf("failed replay preparation leaked global budget %d/%d", globalTasks, globalBytes)
	}
}

func TestLayerRPCAdmissionTransferredBatchClosesWithoutLeak(t *testing.T) {
	router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
	s := New(Options{DC: 2, LayerRPC: router})
	c := &Conn{authKeyID: [8]byte{8, 5}, sessionID: 85, metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
	if err := c.FreezeLayerProfile(tlprofile.Profile225); err != nil {
		t.Fatal(err)
	}
	plan := &inboundPlan{items: []inboundItem{{
		kind: inboundItemRPC, msgID: 100,
		body: exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{}),
	}}}
	defer plan.close()
	if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
		t.Fatal(err)
	}
	c.beginCloseInboundRPCScheduler()
	if err := plan.commitRPCBatch(); err != ErrConnClosed {
		t.Fatalf("commit after connection close = %v, want ErrConnClosed", err)
	}
	plan.close()
	if !c.waitInboundShutdown(time.Second) {
		t.Fatal("connection close did not converge after transferred reservation failed commit")
	}
	if got := c.inflightRPCBytes.Load(); got != 0 {
		t.Fatalf("connection close leaked %d admission bytes", got)
	}
	s.rpcScheduler.budgetMu.Lock()
	globalTasks, globalBytes := s.rpcScheduler.tasks, s.rpcScheduler.bytes
	s.rpcScheduler.budgetMu.Unlock()
	if globalTasks != 0 || globalBytes != 0 {
		t.Fatalf("connection close leaked global budget %d/%d", globalTasks, globalBytes)
	}
}

func TestLayerRPCAdmissionTransferredBatchCommitsConservativeCharge(t *testing.T) {
	router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
	s := New(Options{DC: 2, LayerRPC: router})
	c := &Conn{authKeyID: [8]byte{8, 6}, sessionID: 86, metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
	if err := c.FreezeLayerProfile(tlprofile.Profile225); err != nil {
		t.Fatal(err)
	}
	body := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
	plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
	defer plan.close()
	if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
		t.Fatal(err)
	}
	if err := plan.commitRPCBatch(); err != nil {
		t.Fatal(err)
	}
	wantCharge := layerRPCAdmissionReservationSize(len(body))
	c.rpcMu.Lock()
	queued := len(c.rpcQueue)
	gotCharge := 0
	if queued == 1 {
		gotCharge = c.rpcQueue[0].size
	}
	c.rpcMu.Unlock()
	if queued != 1 || gotCharge != wantCharge {
		t.Fatalf("committed queue = len:%d charge:%d, want 1/%d", queued, gotCharge, wantCharge)
	}
	c.beginCloseInboundRPCScheduler()
	if got := c.inflightRPCBytes.Load(); got != 0 {
		t.Fatalf("queued exact task close leaked %d bytes", got)
	}
	s.rpcScheduler.budgetMu.Lock()
	globalTasks, globalBytes := s.rpcScheduler.tasks, s.rpcScheduler.bytes
	s.rpcScheduler.budgetMu.Unlock()
	if globalTasks != 0 || globalBytes != 0 {
		t.Fatalf("queued exact task close leaked global budget %d/%d", globalTasks, globalBytes)
	}
}

func TestLayerRPCAdmissionLocalDuplicateConsumesNoProvisionalEntry(t *testing.T) {
	router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
	s := New(Options{DC: 2, LayerRPC: router})
	c := &Conn{authKeyID: [8]byte{8, 7}, sessionID: 87, metrics: NopMetrics{}}
	c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
	if err := c.FreezeLayerProfile(tlprofile.Profile225); err != nil {
		t.Fatal(err)
	}
	body := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
	plan := &inboundPlan{items: []inboundItem{
		{kind: inboundItemDuplicate, msgID: 96, body: body},
		{kind: inboundItemRPC, msgID: 100, body: body},
	}}
	defer plan.close()
	if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
		t.Fatal(err)
	}
	if plan.items[0].kind != inboundItemDuplicate || len(plan.rpcTasks) != 1 || c.rpcReserved != 1 {
		t.Fatalf("duplicate/fresh admission = duplicate:%d tasks:%d reserved:%d", plan.items[0].kind, len(plan.rpcTasks), c.rpcReserved)
	}
}

func TestTakeInboundRPCFIFOAdvancesSliceHead(t *testing.T) {
	scheduler := newInboundRPCScheduler(1, 8, 1<<20)
	c := &Conn{metrics: NopMetrics{}}
	c.startInboundRPCScheduler(scheduler, 1, 8, time.Second)
	c.rpcQueue = []inboundRPC{{method: "first"}, {method: "second"}, {method: "third"}}
	c.rpcReady = true
	oldSecond := &c.rpcQueue[1]
	task, ok, _ := c.takeInboundRPC()
	if !ok || task.method != "first" {
		t.Fatalf("take = (%q,%v), want first", task.method, ok)
	}
	if len(c.rpcQueue) != 2 || &c.rpcQueue[0] != oldSecond {
		t.Fatal("FIFO take copied the queue instead of advancing its slice head")
	}
	c.finishInboundRPC(task)
	c.beginCloseInboundRPCScheduler()
}

func BenchmarkLayerRPCAdmissionReservationSize(b *testing.B) {
	for _, size := range []int{64, appfiles.MaxUploadPartBytes + 24} {
		b.Run(time.Duration(size).String(), func(b *testing.B) {
			b.ReportAllocs()
			var charge int
			for i := 0; i < b.N; i++ {
				charge = layerRPCAdmissionReservationSize(size)
			}
			if charge == 0 {
				b.Fatal("zero admission charge")
			}
		})
	}
}