package tcp

import (
	"reflect"
	"testing"

	"github.com/apernet/OpenGFW/analyzer"
)

func TestMinecraftParsing_HandshakeLogin(t *testing.T) {
	handshake := buildMinecraftHandshake(t, 767, "211.136.162.178", 25565, minecraftNextStateLogin)
	want := analyzer.PropMap{
		"yes":             true,
		"candidate":       false,
		"protocol":        767,
		"server_addr":     "211.136.162.178",
		"server_port":     uint16(25565),
		"next_state":      minecraftNextStateLogin,
		"next_state_name": "login",
		"packet_length":   len(handshake) - 1,
	}

	u, done := newMinecraftStream(nil).Feed(false, false, false, 0, handshake)
	if !done {
		t.Fatal("stream not done after full handshake")
	}
	if !reflect.DeepEqual(u.M, want) {
		t.Errorf("parsed = %#v, want %#v", u.M, want)
	}
}

func TestMinecraftParsing_FragmentedHandshake(t *testing.T) {
	handshake := buildMinecraftHandshake(t, 767, "mc.example.com", 25565, minecraftNextStateStatus)
	s := newMinecraftStream(nil)

	u, done := s.Feed(false, false, false, 0, handshake[:2])
	if done {
		t.Fatal("stream done on partial handshake")
	}
	if got := u.M.Get("candidate"); got != true {
		t.Fatalf("candidate = %v, want true", got)
	}

	u, done = s.Feed(false, false, false, 0, handshake[2:])
	if !done {
		t.Fatal("stream not done after remaining handshake")
	}
	if got := u.M.Get("yes"); got != true {
		t.Fatalf("yes = %v, want true", got)
	}
	if got := u.M.Get("server_addr"); got != "mc.example.com" {
		t.Fatalf("server_addr = %v, want mc.example.com", got)
	}
}

func TestMinecraftParsing_InvalidAfterCandidate(t *testing.T) {
	s := newMinecraftStream(nil)
	u, done := s.Feed(false, false, false, 0, []byte{0x04})
	if done {
		t.Fatal("stream done on partial packet")
	}
	if got := u.M.Get("candidate"); got != true {
		t.Fatalf("candidate = %v, want true", got)
	}

	u, done = s.Feed(false, false, false, 0, []byte{0x01, 0x02, 0x03, 0x04})
	if !done {
		t.Fatal("stream not done on invalid packet")
	}
	if got := u.M.Get("candidate"); got != false {
		t.Fatalf("candidate = %v, want false", got)
	}
	if got := u.M.Get("yes"); got != false {
		t.Fatalf("yes = %v, want false", got)
	}
}

func TestMinecraftParsing_RejectsInvalidUTF8Address(t *testing.T) {
	packet := []byte{0x00}
	packet = append(packet, encodeMinecraftVarInt(767)...)
	packet = append(packet, 0x01, 0xff)
	packet = append(packet, 0x63, 0xdd)
	packet = append(packet, encodeMinecraftVarInt(minecraftNextStateLogin)...)
	data := append(encodeMinecraftVarInt(int32(len(packet))), packet...)

	u, done := newMinecraftStream(nil).Feed(false, false, false, 0, data)
	if !done {
		t.Fatal("stream not done on invalid UTF-8 address")
	}
	if got := u.M.Get("yes"); got != false {
		t.Fatalf("yes = %v, want false", got)
	}
}

func buildMinecraftHandshake(t *testing.T, protocol int32, addr string, port uint16, nextState int32) []byte {
	t.Helper()
	packet := []byte{0x00}
	packet = append(packet, encodeMinecraftVarInt(protocol)...)
	packet = append(packet, encodeMinecraftVarInt(int32(len(addr)))...)
	packet = append(packet, []byte(addr)...)
	packet = append(packet, byte(port>>8), byte(port))
	packet = append(packet, encodeMinecraftVarInt(nextState)...)
	return append(encodeMinecraftVarInt(int32(len(packet))), packet...)
}

func encodeMinecraftVarInt(value int32) []byte {
	var out []byte
	for {
		b := byte(value & 0x7f)
		value = int32(uint32(value) >> 7)
		if value != 0 {
			b |= 0x80
		}
		out = append(out, b)
		if value == 0 {
			return out
		}
	}
}