package mtprotoedge

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"testing"
	"time"

	"github.com/iamxvbaba/td/bin"
	"github.com/iamxvbaba/td/crypto"
	"github.com/iamxvbaba/td/proto"
)

func TestRPCResultCloneReservationIsOneShotUnderReleaseRace(t *testing.T) {
	const iterations = 256
	for i := 0; i < iterations; i++ {
		encoded := encodedRPCResultForPriorityTest(int64(i+1), 0)
		budget := newOutboundTrackedBudget(int64(len(encoded.body)))
		if !budget.reserve(len(encoded.body)) {
			t.Fatal("reserve body")
		}
		reserved := &outboundBodyReservation{budget: budget, bytes: len(encoded.body)}
		start := make(chan struct{})
		taken := make(chan *outboundOp, 1)
		go func() {
			<-start
			op, _ := reserved.take(encoded, fallbackOutboundOpPool)
			taken <- op
		}()
		released := make(chan struct{})
		go func() {
			<-start
			reserved.release()
			close(released)
		}()
		close(start)
		op := <-taken
		<-released
		op.releaseReservation(budget)
		reserved.release()
		if got := budget.snapshot(); got != 0 {
			t.Fatalf("iteration %d retained bytes = %d, want 0", i, got)
		}
	}
}

func TestRPCResultReservationReleaseWinsAdmissionRollback(t *testing.T) {
	encoded := encodedRPCResultForPriorityTest(6999, 0)
	budget := newOutboundTrackedBudget(int64(len(encoded.body)))
	if !budget.reserve(len(encoded.body)) {
		t.Fatal("reserve body")
	}
	reserved := &outboundBodyReservation{budget: budget, bytes: len(encoded.body)}
	op, err := reserved.take(encoded, fallbackOutboundOpPool)
	if err != nil {
		t.Fatalf("take reservation: %v", err)
	}
	// Model a watchdog that retires the queued owner just before queue admission
	// rolls the op back. The rollback must observe the release request and return
	// the raw op charge instead of resurrecting an owner nobody will release.
	reserved.release()
	if !reserved.reclaim(op) {
		t.Fatal("reclaim actor reservation")
	}
	if got := budget.snapshot(); got != 0 {
		t.Fatalf("retained bytes after watchdog/rollback race = %d, want 0", got)
	}
	reserved.release()
	op.releaseReservation(budget)
}

func TestCachedRPCResultReplayUsesPreReservedBodyWithoutDoubleCharge(t *testing.T) {
	encoded := encodedRPCResultForPriorityTest(7001, 32<<10)
	budget := newOutboundTrackedBudget(int64(len(encoded.body)))
	tr := newGatedRecordingTransport()
	c := newOutboundTestConn(t, tr, budget)
	s := New(Options{WriteTimeout: time.Second})

	done := make(chan error, 1)
	go func() {
		done <- s.sendReplayedRPCResult(context.Background(), c, encoded)
	}()
	select {
	case <-tr.started:
	case <-time.After(time.Second):
		t.Fatal("cached replay did not reach the blocked physical write")
	}
	if got, want := budget.snapshot(), int64(len(encoded.body)); got != want {
		t.Fatalf("blocked replay retained bytes = %d, want exactly one body %d", got, want)
	}
	tr.once.Do(func() { close(tr.release) })
	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("cached replay: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("cached replay did not finish")
	}
	// rpc_result is reliable and remains charged once as resend state until ACK/close.
	if got, want := budget.snapshot(), int64(len(encoded.body)); got != want {
		t.Fatalf("pending replay retained bytes = %d, want %d", got, want)
	}
	c.Close()
	if got := budget.snapshot(); got != 0 {
		t.Fatalf("replay retained bytes after close = %d, want 0", got)
	}
}

func TestQueuedRPCRewrapClonesAreBoundedBeforeAllocation(t *testing.T) {
	source := encodedRPCResultForPriorityTest(7101, 64<<10)
	perBody := len(source.body)
	budget := newOutboundTrackedBudget(int64(2 * perBody))
	c := &Conn{metrics: NopMetrics{}, outboundTrackedBudget: budget}

	first, firstReserved, err := c.cloneRPCResultForRequestReserved(source, 7102, false)
	if err != nil || first == nil || firstReserved == nil {
		t.Fatalf("first queued clone = %p reservation=%p err=%v", first, firstReserved, err)
	}
	firstAlias := &rpcRewrapAlias{bodyReservation: firstReserved}
	second, secondReserved, err := c.cloneRPCResultForRequestReserved(source, 7103, false)
	if err != nil || second == nil || secondReserved == nil {
		t.Fatalf("second queued clone = %p reservation=%p err=%v", second, secondReserved, err)
	}
	secondAlias := &rpcRewrapAlias{bodyReservation: secondReserved}
	if got, want := budget.snapshot(), int64(2*perBody); got != want {
		t.Fatalf("two queued aliases retained bytes = %d, want %d", got, want)
	}
	third, thirdReserved, err := c.cloneRPCResultForRequestReserved(source, 7104, false)
	if !errors.Is(err, ErrOutboundTrackedBudget) || third != nil || thirdReserved != nil {
		t.Fatalf("third queued clone = %p reservation=%p err=%v, want pre-allocation budget rejection", third, thirdReserved, err)
	}
	if got, want := budget.snapshot(), int64(2*perBody); got != want {
		t.Fatalf("budget changed after rejected clone = %d, want %d", got, want)
	}
	firstAlias.finishReplayRestoreWithoutDelivery()
	secondAlias.finishReplayRestoreWithoutDelivery()
	if got := budget.snapshot(); got != 0 {
		t.Fatalf("queued alias bytes after terminal release = %d, want 0", got)
	}
}

func TestOutboundActorRetargetRequiresSecondBodyReservation(t *testing.T) {
	const (
		oldReqID = int64(7201)
		newReqID = int64(7202)
	)
	encoded := encodedRPCResultForPriorityTest(oldReqID, 32<<10)
	encoded.delivery = newRPCResultDelivery(oldReqID)
	encoded.markQueued()
	if !encoded.tryRetarget(newReqID) {
		t.Fatal("retarget prepared rpc_result")
	}
	budget := newOutboundTrackedBudget(int64(len(encoded.body)))
	if !budget.reserve(len(encoded.body)) {
		t.Fatal("reserve original queued body")
	}
	tr := &failAfterTransport{}
	c := newOutboundTestConn(t, tr, budget)
	state := newOutboundState(budget)
	var terminalErr error
	var terminalBytes int64
	op := outboundOp{
		kind:              outboundSend,
		ctx:               context.Background(),
		msgType:           proto.MessageServerResponse,
		encoded:           encoded,
		reservedBytes:     len(encoded.body),
		reservationBudget: budget,
		enqueuedAt:        time.Now(),
		terminal: func(err error) {
			terminalErr = err
			terminalBytes = budget.snapshot()
		},
	}
	err := c.handleOutboundSend(state, &op)
	op.finish(outboundResult{err: err})
	if !errors.Is(terminalErr, ErrOutboundTrackedBudget) {
		t.Fatalf("retarget terminal error = %v, want %v", terminalErr, ErrOutboundTrackedBudget)
	}
	if terminalBytes != 0 {
		t.Fatalf("bytes visible to terminal = %d, want producer reservation released", terminalBytes)
	}
	if got := tr.sends.Load(); got != 0 {
		t.Fatalf("retarget under one-body budget wrote %d frames, want 0", got)
	}
	if got := int64(binary.LittleEndian.Uint64(encoded.body[4:12])); got != oldReqID {
		t.Fatalf("source req_msg_id mutated to %d before second-body admission", got)
	}
	if got := budget.snapshot(); got != 0 {
		t.Fatalf("retarget bytes after terminal = %d, want 0", got)
	}
}

func TestOutboundActorRetargetTransfersOnlyReplacementToPending(t *testing.T) {
	const (
		oldReqID = int64(7301)
		newReqID = int64(7302)
	)
	encoded := encodedRPCResultForPriorityTest(oldReqID, 32<<10)
	encoded.delivery = newRPCResultDelivery(oldReqID)
	encoded.markQueued()
	if !encoded.tryRetarget(newReqID) {
		t.Fatal("retarget prepared rpc_result")
	}
	perBody := len(encoded.body)
	budget := newOutboundTrackedBudget(int64(2 * perBody))
	if !budget.reserve(perBody) {
		t.Fatal("reserve original queued body")
	}
	tr := &failAfterTransport{}
	c := newOutboundTestConn(t, tr, budget)
	state := newOutboundState(budget)
	var terminalErr error
	var terminalBytes int64
	op := outboundOp{
		kind:              outboundSend,
		ctx:               context.Background(),
		msgType:           proto.MessageServerResponse,
		encoded:           encoded,
		reservedBytes:     perBody,
		reservationBudget: budget,
		enqueuedAt:        time.Now(),
		terminal: func(err error) {
			terminalErr = err
			terminalBytes = budget.snapshot()
		},
	}
	err := c.handleOutboundSend(state, &op)
	op.finish(outboundResult{err: err})
	if terminalErr != nil {
		t.Fatalf("retarget terminal error: %v", terminalErr)
	}
	if terminalBytes != int64(perBody) {
		t.Fatalf("bytes visible to terminal = %d, want only pending replacement %d", terminalBytes, perBody)
	}
	if got := budget.snapshot(); got != int64(perBody) {
		t.Fatalf("bytes after terminal = %d, want one pending replacement %d", got, perBody)
	}
	data, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
	if err != nil {
		t.Fatalf("decrypt retargeted frame: %v", err)
	}
	var result proto.Result
	if err := result.Decode(&bin.Buffer{Buf: append([]byte(nil), data.Data()...)}); err != nil {
		t.Fatalf("decode retargeted rpc_result: %v", err)
	}
	if result.RequestMessageID != newReqID {
		t.Fatalf("wire req_msg_id = %d, want %d", result.RequestMessageID, newReqID)
	}
	state.releaseAll()
	if got := budget.snapshot(); got != 0 {
		t.Fatalf("retarget bytes after pending release = %d, want 0", got)
	}
}