// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package rpc

import (
	"bytes"
	"errors"
	"io"
	"strings"
	"testing"
)

func TestFragmentHeaderLayout(t *testing.T) {
	got := AppendFragmentHeader(nil, 5, true)
	want := []byte{0x80, 0, 0, 5}
	if !bytes.Equal(got, want) {
		t.Fatalf("last fragment header: got %x, want %x", got, want)
	}
	got = AppendFragmentHeader(nil, 5, false)
	want = []byte{0, 0, 0, 5}
	if !bytes.Equal(got, want) {
		t.Fatalf("continuation header: got %x, want %x", got, want)
	}
}

func TestWriteReadRoundTrip(t *testing.T) {
	records := [][]byte{
		nil,
		[]byte("a"),
		[]byte("abc"),
		[]byte("abcd"),
		bytes.Repeat([]byte{0xc3}, 70000),
	}
	for _, want := range records {
		pr, pw := io.Pipe()
		go func() {
			err := WriteRecord(pw, want)
			if err != nil {
				pw.CloseWithError(err)
				return
			}
			pw.Close()
		}()
		got, err := ReadRecord(pr, 1<<20)
		if err != nil {
			t.Fatalf("record of %d bytes: %v", len(want), err)
		}
		if !bytes.Equal(got, want) {
			t.Fatalf("record of %d bytes came back as %d bytes", len(want), len(got))
		}
	}
}

func TestReadReassemblesFragments(t *testing.T) {
	// A record delivered as three fragments arrives as the same bytes.
	var stream bytes.Buffer
	stream.Write(AppendFragmentHeader(nil, 2, false))
	stream.WriteString("ab")
	stream.Write(AppendFragmentHeader(nil, 0, false))
	stream.Write(AppendFragmentHeader(nil, 3, true))
	stream.WriteString("cde")

	got, err := ReadRecord(&stream, 64)
	if err != nil {
		t.Fatalf("read: %v", err)
	}
	if string(got) != "abcde" {
		t.Fatalf("reassembled record: got %q, want %q", got, "abcde")
	}
}

func TestReadRecordTooLarge(t *testing.T) {
	var stream bytes.Buffer
	stream.Write(AppendFragmentHeader(nil, 100, true))
	stream.Write(bytes.Repeat([]byte{0}, 100))

	if _, err := ReadRecord(&stream, 64); !errors.Is(err, ErrRecordTooLarge) {
		t.Fatalf("a record of 100 bytes against a limit of 64 returned %v", err)
	}
	// A record split over fragments is bounded by the record total, not by
	// the single fragment length.
	stream.Reset()
	stream.Write(AppendFragmentHeader(nil, 50, false))
	stream.Write(bytes.Repeat([]byte{0}, 50))
	stream.Write(AppendFragmentHeader(nil, 50, true))
	stream.Write(bytes.Repeat([]byte{0}, 50))
	if _, err := ReadRecord(&stream, 64); !errors.Is(err, ErrRecordTooLarge) {
		t.Fatalf("a split record of 100 bytes against a limit of 64 returned %v", err)
	}
}

func TestReadRecordAtLimit(t *testing.T) {
	// A record of exactly the limit is legal; one byte more is not.
	// The boundary is strict, so a record of limit bytes must arrive.
	var stream bytes.Buffer
	stream.Write(AppendFragmentHeader(nil, 64, true))
	stream.Write(bytes.Repeat([]byte{0}, 64))
	if got, err := ReadRecord(&stream, 64); err != nil || len(got) != 64 {
		t.Fatalf("a record of exactly the limit: %d bytes, %v", len(got), err)
	}
	stream.Reset()
	stream.Write(AppendFragmentHeader(nil, 65, true))
	stream.Write(bytes.Repeat([]byte{0}, 65))
	if _, err := ReadRecord(&stream, 64); !errors.Is(err, ErrRecordTooLarge) {
		t.Fatalf("a record one byte over the limit returned %v", err)
	}
}

func TestReadCleanEOF(t *testing.T) {
	if _, err := ReadRecord(strings.NewReader(""), 64); !errors.Is(err, io.EOF) {
		t.Fatalf("an empty stream returned %v, want io.EOF", err)
	}
}

func TestReadTruncated(t *testing.T) {
	// A header that promises ten bytes and a body of three is a truncated
	// fragment, not a clean end of stream.
	var stream bytes.Buffer
	stream.Write(AppendFragmentHeader(nil, 10, true))
	stream.WriteString("abc")
	if _, err := ReadRecord(&stream, 64); !errors.Is(err, io.ErrUnexpectedEOF) {
		t.Fatalf("a truncated body returned %v", err)
	}
}

func TestReadTruncatedHeader(t *testing.T) {
	// A record already reassembling that loses its next header is also a
	// truncation, not a clean end of stream.
	var stream bytes.Buffer
	stream.Write(AppendFragmentHeader(nil, 1, false))
	stream.WriteString("a")
	stream.Write([]byte{0, 0})
	if _, err := ReadRecord(&stream, 64); !errors.Is(err, io.ErrUnexpectedEOF) {
		t.Fatalf("a truncated continuation header returned %v", err)
	}
}