package btcwire

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"testing"

	"git.smesh.lol/nostrino/pkg/crypto/chainhash"
)

// frameHeader builds a 24-byte frame header with the given fields.
func frameHeader(magic uint32, command [CommandSize]byte, length uint32) []byte {
	hdr := make([]byte, MessageHeaderSize)
	binary.BigEndian.PutUint32(hdr[0:4], magic)
	copy(hdr[4:4+CommandSize], command[:])
	binary.LittleEndian.PutUint32(hdr[16:20], length)
	sum := Checksum(nil)
	copy(hdr[20:24], sum[:])
	return hdr
}

func TestChecksum(t *testing.T) {
	if got, want := Checksum(nil), [4]byte{0x5d, 0xf6, 0xe0, 0xe2}; got != want {
		t.Fatalf("Checksum(nil) = %x, want %x", got, want)
	}
	payload := []byte("nostrino compact filter client")
	sum := Checksum(payload)
	full := chainhash.DoubleHashB(payload)
	if !bytes.Equal(sum[:], full[:4]) {
		t.Fatalf("Checksum = %x, want first 4 bytes of %x", sum, full)
	}
}

// TestEncodeMessageLayout pins the four bytes of every network magic as they
// appear on the wire. These are the strings Bitcoin prints for each network
// read left to right: a node that receives any other order sees an unknown
// network and drops the connection before it parses the version payload.
func TestEncodeMessageLayout(t *testing.T) {
	wantMagic := []struct {
		name  string
		magic uint32
		bytes []byte
	}{
		{"mainnet", MainNet, []byte{0xf9, 0xbe, 0xb4, 0xd9}},
		{"testnet3", TestNet3, []byte{0x0b, 0x11, 0x09, 0x07}},
		{"regtest", RegTest, []byte{0xfa, 0xbf, 0xb5, 0xda}},
		{"signet", SigNet, []byte{0x0a, 0x03, 0xcf, 0x40}},
	}
	payload := []byte("hello, wire")
	for _, network := range wantMagic {
		t.Run(network.name, func(t *testing.T) {
			frame, err := EncodeMessage(network.magic, CmdPing, payload)
			if err != nil {
				t.Fatalf("EncodeMessage: %v", err)
			}
			var want []byte
			want = append(want, network.bytes...)
			// "ping" NUL padded to 12 bytes.
			want = append(want, []byte("ping")...)
			want = append(want, make([]byte, CommandSize-len("ping"))...)
			// Payload length, little endian.
			want = append(want, byte(len(payload)), 0, 0, 0)
			// Checksum.
			sum := Checksum(payload)
			want = append(want, sum[:]...)
			want = append(want, payload...)
			if !bytes.Equal(frame, want) {
				t.Fatalf("frame = %x, want %x", frame, want)
			}
		})
	}
}

// TestTheReversedMagicIsRejected: the pre-fix codec wrote each magic little
// endian, which reverses its four bytes on the wire. A round trip through the
// broken encoder and decoder agrees with itself, so it proves nothing; this
// test decodes the exact bytes the old code produced and requires a refusal.
func TestTheReversedMagicIsRejected(t *testing.T) {
	// A valid mainnet ping, as the fixed codec writes it.
	good, err := EncodeMessage(MainNet, CmdPing, nil)
	if err != nil {
		t.Fatalf("EncodeMessage: %v", err)
	}
	if got, want := good[0:4], []byte{0xf9, 0xbe, 0xb4, 0xd9}; !bytes.Equal(got, want) {
		t.Fatalf("mainnet magic on the wire = %x, want %x", got, want)
	}

	// The same frame with the magic byte-reversed, which is what the old
	// little-endian encoder put on the wire.
	reversed := append([]byte(nil), good...)
	copy(reversed[0:4], []byte{0xd9, 0xb4, 0xbe, 0xf9})

	if _, err := DecodeMessage(MainNet, bytes.NewReader(reversed)); !errors.Is(err, ErrBadMagic) {
		t.Fatalf("decoding the reversed mainnet magic: %v, want ErrBadMagic", err)
	}

	// The same bug on regtest: da b5 bf fa instead of fa bf b5 da.
	regtest, err := EncodeMessage(RegTest, CmdPing, nil)
	if err != nil {
		t.Fatalf("EncodeMessage: %v", err)
	}
	if got, want := regtest[0:4], []byte{0xfa, 0xbf, 0xb5, 0xda}; !bytes.Equal(got, want) {
		t.Fatalf("regtest magic on the wire = %x, want %x", got, want)
	}
	reversedRegtest := append([]byte(nil), regtest...)
	copy(reversedRegtest[0:4], []byte{0xda, 0xb5, 0xbf, 0xfa})
	if _, err := DecodeMessage(RegTest, bytes.NewReader(reversedRegtest)); !errors.Is(err, ErrBadMagic) {
		t.Fatalf("decoding the reversed regtest magic: %v, want ErrBadMagic", err)
	}
}

func TestMessageRoundTrip(t *testing.T) {
	magics := []uint32{MainNet, TestNet3, RegTest, SigNet}
	payloads := [][]byte{
		nil,
		{},
		{0x00},
		bytes.Repeat([]byte{0xab}, 253),
		bytes.Repeat([]byte{0xcd}, 1000),
	}
	for _, magic := range magics {
		for _, command := range allCommands {
			for _, payload := range payloads {
				name := fmt.Sprintf("%08x/%s/%d", magic, command, len(payload))
				t.Run(name, func(t *testing.T) {
					frame, err := EncodeMessage(magic, command, payload)
					if err != nil {
						t.Fatalf("EncodeMessage: %v", err)
					}
					if got, want := len(frame), MessageHeaderSize+len(payload); got != want {
						t.Fatalf("frame length = %d, want %d", got, want)
					}
					msg, err := DecodeMessage(magic, bytes.NewReader(frame))
					if err != nil {
						t.Fatalf("DecodeMessage: %v", err)
					}
					if msg.Command != command {
						t.Fatalf("command = %q, want %q", msg.Command, command)
					}
					if !bytes.Equal(msg.Payload, payload) {
						t.Fatalf("payload = %x, want %x", msg.Payload, payload)
					}
				})
			}
		}
	}
}

func TestDecodeMessageErrors(t *testing.T) {
	pingPayload := make([]byte, 8)
	goodPing, err := EncodeMessage(MainNet, CmdPing, pingPayload)
	if err != nil {
		t.Fatalf("EncodeMessage: %v", err)
	}
	var zeroCommand [CommandSize]byte
	var badCharCommand [CommandSize]byte
	copy(badCharCommand[:], "pi\x01ng")

	tests := []struct {
		name  string
		frame func() []byte
		want  error
	}{
		{
			name:  "empty input",
			frame: func() []byte { return nil },
			want:  io.EOF,
		},
		{
			name: "bad magic",
			frame: func() []byte {
				f := append([]byte(nil), goodPing...)
				binary.BigEndian.PutUint32(f[0:4], TestNet3)
				return f
			},
			want: ErrBadMagic,
		},
		{
			name: "bad checksum",
			frame: func() []byte {
				f, err := EncodeMessage(MainNet, CmdTx, []byte{1, 2, 3})
				if err != nil {
					t.Fatalf("EncodeMessage: %v", err)
				}
				f[len(f)-1] ^= 0xff
				return f
			},
			want: ErrBadChecksum,
		},
		{
			name:  "truncated header",
			frame: func() []byte { return goodPing[:10] },
			want:  io.ErrUnexpectedEOF,
		},
		{
			name:  "truncated payload",
			frame: func() []byte { return goodPing[:len(goodPing)-1] },
			want:  io.ErrUnexpectedEOF,
		},
		{
			name:  "missing payload",
			frame: func() []byte { return goodPing[:MessageHeaderSize] },
			want:  io.ErrUnexpectedEOF,
		},
		{
			name: "oversize payload",
			frame: func() []byte {
				var command [CommandSize]byte
				copy(command[:], CmdTx)
				return frameHeader(MainNet, command, MaxPayload+1)
			},
			want: ErrPayloadTooLarge,
		},
		{
			name: "all NUL command",
			frame: func() []byte {
				return frameHeader(MainNet, zeroCommand, 0)
			},
			want: ErrInvalidCommand,
		},
		{
			name: "non-printable command",
			frame: func() []byte {
				return frameHeader(MainNet, badCharCommand, 0)
			},
			want: ErrInvalidCommand,
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			_, err := DecodeMessage(MainNet, bytes.NewReader(tc.frame()))
			if !errors.Is(err, tc.want) {
				t.Fatalf("DecodeMessage error = %v, want %v", err, tc.want)
			}
		})
	}
}

func TestDecodeUnknownCommand(t *testing.T) {
	// Commands the package does not implement must still frame correctly; the
	// caller decides what to do with them.
	payload := []byte{0xde, 0xad}
	frame, err := EncodeMessage(MainNet, "notacommand", payload)
	if err != nil {
		t.Fatalf("EncodeMessage: %v", err)
	}
	msg, err := DecodeMessage(MainNet, bytes.NewReader(frame))
	if err != nil {
		t.Fatalf("DecodeMessage: %v", err)
	}
	if msg.Command != "notacommand" {
		t.Fatalf("command = %q, want %q", msg.Command, "notacommand")
	}
	if !bytes.Equal(msg.Payload, payload) {
		t.Fatalf("payload = %x, want %x", msg.Payload, payload)
	}
}

func TestEncodeMessageErrors(t *testing.T) {
	tests := []struct {
		name    string
		command string
		payload []byte
		want    error
	}{
		{"empty command", "", nil, ErrInvalidCommand},
		{"command too long", "thirteenbytes", nil, ErrInvalidCommand},
		{"command with NUL", "pi\x00ng", nil, ErrInvalidCommand},
		{"command with DEL", "ping\x7f", nil, ErrInvalidCommand},
		{"command with high bit", "ping\xff", nil, ErrInvalidCommand},
		{"oversize payload", CmdTx, make([]byte, MaxPayload+1), ErrPayloadTooLarge},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			_, err := EncodeMessage(MainNet, tc.command, tc.payload)
			if !errors.Is(err, tc.want) {
				t.Fatalf("EncodeMessage error = %v, want %v", err, tc.want)
			}
		})
	}
}

func TestCommandConstants(t *testing.T) {
	seen := make(map[string]bool, len(allCommands))
	for _, command := range allCommands {
		if len(command) == 0 {
			t.Fatal("empty command constant")
		}
		if len(command) > CommandSize {
			t.Fatalf("command %q is %d bytes, maximum is %d", command, len(command), CommandSize)
		}
		for i := 0; i < len(command); i++ {
			if c := command[i]; c < 0x20 || c > 0x7e {
				t.Fatalf("command %q contains non-printable byte 0x%02x", command, c)
			}
		}
		if seen[command] {
			t.Fatalf("duplicate command %q", command)
		}
		seen[command] = true
		frame, err := EncodeMessage(MainNet, command, nil)
		if err != nil {
			t.Fatalf("EncodeMessage(%q): %v", command, err)
		}
		msg, err := DecodeMessage(MainNet, bytes.NewReader(frame))
		if err != nil {
			t.Fatalf("DecodeMessage(%q): %v", command, err)
		}
		if msg.Command != command {
			t.Fatalf("decoded command = %q, want %q", msg.Command, command)
		}
	}
}

func TestReader(t *testing.T) {
	var buf bytes.Buffer
	var wantCommands []string
	var wantPayloads [][]byte
	for i, command := range []string{CmdVersion, CmdVerAck, CmdPing, CmdPong} {
		payload := bytes.Repeat([]byte{byte(i)}, i)
		frame, err := EncodeMessage(MainNet, command, payload)
		if err != nil {
			t.Fatalf("EncodeMessage: %v", err)
		}
		buf.Write(frame)
		wantCommands = append(wantCommands, command)
		wantPayloads = append(wantPayloads, payload)
	}
	r := NewReader(&buf, MainNet)
	if r.Magic() != MainNet {
		t.Fatalf("Magic() = %08x, want %08x", r.Magic(), MainNet)
	}
	for i := 0; ; i++ {
		msg, err := r.Next()
		if errors.Is(err, io.EOF) {
			if i != len(wantCommands) {
				t.Fatalf("stream ended after %d messages, want %d", i, len(wantCommands))
			}
			break
		}
		if err != nil {
			t.Fatalf("Next: %v", err)
		}
		if i >= len(wantCommands) {
			t.Fatalf("unexpected extra message %q", msg.Command)
		}
		if msg.Command != wantCommands[i] {
			t.Fatalf("message %d command = %q, want %q", i, msg.Command, wantCommands[i])
		}
		if !bytes.Equal(msg.Payload, wantPayloads[i]) {
			t.Fatalf("message %d payload = %x, want %x", i, msg.Payload, wantPayloads[i])
		}
	}
}

func TestReaderCleanEOF(t *testing.T) {
	r := NewReader(bytes.NewReader(nil), MainNet)
	if _, err := r.Next(); !errors.Is(err, io.EOF) {
		t.Fatalf("Next on empty stream = %v, want io.EOF", err)
	}
}

func TestReaderTruncatedStream(t *testing.T) {
	frame, err := EncodeMessage(MainNet, CmdPing, make([]byte, 8))
	if err != nil {
		t.Fatalf("EncodeMessage: %v", err)
	}
	stream := append(append([]byte(nil), frame...), frame[:5]...)
	r := NewReader(bytes.NewReader(stream), MainNet)
	if _, err := r.Next(); err != nil {
		t.Fatalf("first Next: %v", err)
	}
	_, err = r.Next()
	if !errors.Is(err, io.ErrUnexpectedEOF) {
		t.Fatalf("second Next = %v, want io.ErrUnexpectedEOF", err)
	}
}

func TestReaderHeaderOnlyFrame(t *testing.T) {
	// A frame header that declares a payload, followed by end of stream, is
	// truncation, not a clean end of the message stream.
	frame, err := EncodeMessage(MainNet, CmdPing, make([]byte, 8))
	if err != nil {
		t.Fatalf("EncodeMessage: %v", err)
	}
	r := NewReader(bytes.NewReader(frame[:MessageHeaderSize]), MainNet)
	if _, err := r.Next(); !errors.Is(err, io.ErrUnexpectedEOF) {
		t.Fatalf("Next = %v, want io.ErrUnexpectedEOF", err)
	}
}

func TestEncodeMessagePayloadError(t *testing.T) {
	headers := &Headers{Headers: make([]BlockHeader, MaxBlockHeadersPerMsg+1)}
	if _, err := EncodeMessagePayload(MainNet, CmdHeaders, headers); !errors.Is(err, ErrTooManyItems) {
		t.Fatalf("EncodeMessagePayload error = %v, want ErrTooManyItems", err)
	}
}

func TestMessageDecodeError(t *testing.T) {
	frame, err := EncodeMessage(MainNet, CmdVerAck, []byte{0x01})
	if err != nil {
		t.Fatalf("EncodeMessage: %v", err)
	}
	msg, err := DecodeMessage(MainNet, bytes.NewReader(frame))
	if err != nil {
		t.Fatalf("DecodeMessage: %v", err)
	}
	if err := msg.Decode(&VerAck{}); !errors.Is(err, ErrInvalidPayload) {
		t.Fatalf("Decode error = %v, want ErrInvalidPayload", err)
	}
}

// TestThePayloadLimitIsEnforcedBeforeAllocation: a frame that claims more than the cap is
// refused from its header alone, so a hostile peer cannot make a block read allocate 32 MiB.
func TestThePayloadLimitIsEnforcedBeforeAllocation(t *testing.T) {
	const magic = RegTest

	header := make([]byte, MessageHeaderSize)
	putMagic(header[0:4], magic)
	copy(header[4:4+CommandSize], CmdBlock)
	binary.LittleEndian.PutUint32(header[16:20], 9<<20)

	if _, err := decodeMessageLimit(magic, bytes.NewReader(header), 8<<20); !errors.Is(err, ErrPayloadTooLarge) {
		t.Fatalf("a 9 MiB frame under an 8 MiB cap: %v", err)
	}

	// The reader carries the same limit, which is what the peer uses.
	reader := NewReader(bytes.NewReader(header), magic)
	reader.SetLimit(8 << 20)

	if _, err := reader.Next(); !errors.Is(err, ErrPayloadTooLarge) {
		t.Fatalf("the reader: %v", err)
	}
}

// TestTheReaderLimitResets: a block read must not lower the cap every other read lives under.
func TestTheReaderLimitResets(t *testing.T) {
	reader := NewReader(bytes.NewReader(nil), 1)
	reader.SetLimit(8 << 20)
	reader.SetLimit(MaxPayload)

	if reader.limit != MaxPayload {
		t.Fatalf("the limit is %d", reader.limit)
	}
}