package client

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"testing"

	"github.com/iceBear67/redapricot/client/wire"
)

// ---- packet builders (player/backend side) ----

func mkPacket(body []byte) []byte {
	return append(wire.AppendVarInt(nil, len(body)), body...)
}

func mkLoginStart(proto int, name string, uuid []byte, keyData bool) []byte {
	w := wire.NewWriter().VarInt(0x00).String(name)
	if proto >= protocol1_19 && proto < protocol1_19_3 {
		if keyData {
			w.U8(1).I64(1234567890)
			pub := bytes.Repeat([]byte{0xAA}, 33)
			sig := bytes.Repeat([]byte{0xBB}, 17)
			w.VarInt(len(pub)).Bytes(pub).VarInt(len(sig)).Bytes(sig)
		} else {
			w.U8(0)
		}
	}
	switch {
	case proto >= protocol1_20_2:
		w.Bytes(uuid)
	case proto >= protocol1_19_1:
		if uuid != nil {
			w.U8(1).Bytes(uuid)
		} else {
			w.U8(0)
		}
	}
	return mkPacket(w.Out())
}

func mkPluginRequest(msgID int, channel string, data []byte) []byte {
	body := wire.NewWriter().VarInt(loginS2CPluginRequest).VarInt(msgID).String(channel).Bytes(data).Out()
	return mkPacket(body)
}

// feedLogin drives a full player login prologue through ObserveC2S in n-byte
// chunks.
func feedLogin(v *velocityForwarder, proto int, intent int, loginStart []byte, chunk int) {
	stream := wire.BuildHandshake(proto, "mc.example.com", 25565, intent)
	if loginStart != nil {
		stream = append(stream, loginStart...)
	}
	for len(stream) > 0 {
		n := chunk
		if n > len(stream) {
			n = len(stream)
		}
		v.ObserveC2S(stream[:n])
		stream = stream[n:]
	}
}

// parseResponse validates the injected Login Plugin Response and returns the
// echoed message id and the signed forwarding payload.
func parseResponse(t *testing.T, secret string, inject []byte) (msgID int, payload *wire.Reader) {
	t.Helper()
	r := wire.NewReader(inject)
	plen, err := r.VarInt()
	if err != nil || plen != len(r.Remaining()) {
		t.Fatalf("bad response length prefix: %v (declared %d, have %d)", err, plen, len(r.Remaining()))
	}
	id, _ := r.VarInt()
	if id != loginC2SPluginResponse {
		t.Fatalf("response packet id = %#x, want 0x02", id)
	}
	msgID, _ = r.VarInt()
	ok, _ := r.U8()
	if ok != 1 {
		t.Fatalf("response not marked successful")
	}
	sig, err := r.Bytes(32)
	if err != nil {
		t.Fatalf("response missing signature: %v", err)
	}
	data := r.Remaining()
	mac := hmac.New(sha256.New, []byte(secret))
	mac.Write(data)
	if !hmac.Equal(sig, mac.Sum(nil)) {
		t.Fatalf("forwarding payload signature does not verify")
	}
	return msgID, wire.NewReader(data)
}

func assertPayload(t *testing.T, r *wire.Reader, version int, ip, name, uuidHex string) {
	t.Helper()
	gotVer, _ := r.VarInt()
	if gotVer != version {
		t.Fatalf("forwarding version = %d, want %d", gotVer, version)
	}
	gotIP, _ := r.String()
	if gotIP != ip {
		t.Fatalf("forwarded address = %q, want %q", gotIP, ip)
	}
	gotUUID, err := r.Bytes(16)
	if err != nil {
		t.Fatalf("payload missing uuid: %v", err)
	}
	if hex.EncodeToString(gotUUID) != uuidHex {
		t.Fatalf("forwarded uuid = %x, want %s", gotUUID, uuidHex)
	}
	gotName, _ := r.String()
	if gotName != name {
		t.Fatalf("forwarded username = %q, want %q", gotName, name)
	}
	props, err := r.VarInt()
	if err != nil || props != 0 {
		t.Fatalf("forwarded properties = %d (%v), want 0", props, err)
	}
	if len(r.Remaining()) != 0 {
		t.Fatalf("trailing bytes in forwarding payload: %x", r.Remaining())
	}
}

// ---- tests ----

const testSecret = "unit-secret"

func TestVelocityInterceptModern(t *testing.T) {
	uuid, _ := hex.DecodeString("00112233445566778899aabbccddeeff")
	v := newVelocityForwarder(testSecret, "203.0.113.7")
	feedLogin(v, 767, intentLogin, mkLoginStart(767, "icybear", uuid, false), 1)

	// Backend query, requesting up to forwarding version 4, fed byte by byte:
	// nothing may reach the player, and the response appears with the last byte.
	req := mkPluginRequest(99, velocityChannel, []byte{0x04})
	var inject []byte
	for i, b := range req {
		fwd, inj := v.ProcessS2C([]byte{b})
		if len(fwd) != 0 {
			t.Fatalf("byte %d: request leaked to the player: %x", i, fwd)
		}
		if inj != nil {
			inject = inj
		}
	}
	if inject == nil {
		t.Fatalf("no response was injected")
	}
	msgID, payload := parseResponse(t, testSecret, inject)
	if msgID != 99 {
		t.Fatalf("echoed message id = %d, want 99", msgID)
	}
	assertPayload(t, payload, velocityVersionLazySession, "203.0.113.7", "icybear",
		"00112233445566778899aabbccddeeff")
	if !v.Passthrough() {
		t.Fatalf("interceptor should be passthrough after answering")
	}
	garbage := []byte{0xde, 0xad, 0xbe, 0xef}
	if fwd, inj := v.ProcessS2C(garbage); !bytes.Equal(fwd, garbage) || inj != nil {
		t.Fatalf("post-login bytes not passed through verbatim")
	}
}

func TestVelocityNegativeMessageID(t *testing.T) {
	// Paper picks the message id with ThreadLocalRandom.nextInt(): it is
	// negative half the time and must be echoed bit-exactly.
	uuid, _ := hex.DecodeString("00112233445566778899aabbccddeeff")
	v := newVelocityForwarder(testSecret, "198.51.100.1")
	feedLogin(v, 767, intentLogin, mkLoginStart(767, "neg", uuid, false), 64)
	_, inject := v.ProcessS2C(mkPluginRequest(-123456, velocityChannel, []byte{0x04}))
	if inject == nil {
		t.Fatalf("no response was injected")
	}
	msgID, _ := parseResponse(t, testSecret, inject)
	if !bytes.Equal(wire.AppendVarInt(nil, msgID), wire.AppendVarInt(nil, -123456)) {
		t.Fatalf("negative message id not echoed bit-exactly (got %d)", msgID)
	}
}

func TestVelocityOfflineUUIDAndV1(t *testing.T) {
	// 1.18.2 player: no UUID in Login Start -> Java's offline UUID; an old
	// backend requesting version 1 gets version 1.
	v := newVelocityForwarder(testSecret, "192.0.2.9")
	feedLogin(v, 758, intentLogin, mkLoginStart(758, "Notch", nil, false), 3)
	_, inject := v.ProcessS2C(mkPluginRequest(7, velocityChannel, []byte{0x01}))
	if inject == nil {
		t.Fatalf("no response was injected")
	}
	_, payload := parseResponse(t, testSecret, inject)
	// UUID.nameUUIDFromBytes("OfflinePlayer:Notch".getBytes(UTF_8)).
	assertPayload(t, payload, velocityVersionDefault, "192.0.2.9", "Notch",
		"b50ad385829d3141a2167e7d7539ba7f")
}

func TestVelocityVersionGating(t *testing.T) {
	// A modern backend (requests 4) behind a pre-1.19.3 player must get v1.
	v := newVelocityForwarder(testSecret, "192.0.2.9")
	feedLogin(v, 758, intentLogin, mkLoginStart(758, "Old", nil, false), 5)
	_, inject := v.ProcessS2C(mkPluginRequest(1, velocityChannel, []byte{0x04}))
	_, payload := parseResponse(t, testSecret, inject)
	ver, _ := payload.VarInt()
	if ver != velocityVersionDefault {
		t.Fatalf("version = %d, want 1 for a pre-1.19.3 player", ver)
	}
	// An empty request (very old backend) also means v1.
	v2 := newVelocityForwarder(testSecret, "192.0.2.9")
	feedLogin(v2, 767, intentLogin, mkLoginStart(767, "New", make([]byte, 16), false), 5)
	_, inject2 := v2.ProcessS2C(mkPluginRequest(1, velocityChannel, nil))
	_, payload2 := parseResponse(t, testSecret, inject2)
	ver2, _ := payload2.VarInt()
	if ver2 != velocityVersionDefault {
		t.Fatalf("version = %d, want 1 for an empty version request", ver2)
	}
}

func TestVelocityLoginStartVariants(t *testing.T) {
	uuid, _ := hex.DecodeString("ffeeddccbbaa99887766554433221100")
	cases := []struct {
		name    string
		proto   int
		start   []byte
		uuidHex string
	}{
		{"1.19 with key, no uuid", 759, mkLoginStart(759, "Notch", nil, true),
			"b50ad385829d3141a2167e7d7539ba7f"},
		{"1.19.1 with key and uuid", 760, mkLoginStart(760, "Notch", uuid, true),
			"ffeeddccbbaa99887766554433221100"},
		{"1.19.3 optional uuid present", 761, mkLoginStart(761, "Notch", uuid, false),
			"ffeeddccbbaa99887766554433221100"},
		{"1.19.3 optional uuid absent", 761, mkLoginStart(761, "Notch", nil, false),
			"b50ad385829d3141a2167e7d7539ba7f"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			v := newVelocityForwarder(testSecret, "192.0.2.1")
			feedLogin(v, tc.proto, intentLogin, tc.start, 2)
			_, inject := v.ProcessS2C(mkPluginRequest(3, velocityChannel, []byte{0x01}))
			if inject == nil {
				t.Fatalf("no response was injected")
			}
			_, payload := parseResponse(t, testSecret, inject)
			assertPayload(t, payload, velocityVersionDefault, "192.0.2.1", "Notch", tc.uuidHex)
		})
	}
}

func TestVelocityStatusPassthrough(t *testing.T) {
	v := newVelocityForwarder(testSecret, "192.0.2.1")
	feedLogin(v, 767, 1 /* status */, nil, 100)
	if !v.Passthrough() {
		t.Fatalf("status intent should turn the stream transparent")
	}
	data := []byte("not a minecraft packet at all")
	if fwd, inj := v.ProcessS2C(data); !bytes.Equal(fwd, data) || inj != nil {
		t.Fatalf("status traffic must pass through untouched")
	}
}

func TestVelocityOtherChannelForwarded(t *testing.T) {
	uuid := make([]byte, 16)
	v := newVelocityForwarder(testSecret, "192.0.2.1")
	feedLogin(v, 767, intentLogin, mkLoginStart(767, "modded", uuid, false), 50)

	// A modded-handshake query and the velocity query coalesced in one chunk:
	// the first must reach the player, the second must not.
	other := mkPluginRequest(1, "fml:loginwrapper", []byte{0x00, 0x01})
	velo := mkPluginRequest(2, velocityChannel, []byte{0x04})
	fwd, inject := v.ProcessS2C(append(append([]byte{}, other...), velo...))
	if !bytes.Equal(fwd, other) {
		t.Fatalf("non-velocity query not forwarded verbatim:\n got %x\nwant %x", fwd, other)
	}
	if inject == nil {
		t.Fatalf("velocity query in the same chunk was not answered")
	}
}

func TestVelocityNoQueryLoginSuccess(t *testing.T) {
	// Backend without velocity forwarding: the first login packet ends
	// interception and everything flows verbatim.
	uuid := make([]byte, 16)
	v := newVelocityForwarder(testSecret, "192.0.2.1")
	feedLogin(v, 767, intentLogin, mkLoginStart(767, "plain", uuid, false), 50)

	success := mkPacket(wire.NewWriter().VarInt(loginS2CSuccess).Bytes(uuid).String("plain").VarInt(0).Out())
	tail := []byte("compressed gibberish after login")
	fwd, inject := v.ProcessS2C(append(append([]byte{}, success...), tail...))
	if inject != nil {
		t.Fatalf("nothing should be injected without a velocity query")
	}
	want := append(append([]byte{}, success...), tail...)
	if !bytes.Equal(fwd, want) {
		t.Fatalf("login success not flushed verbatim")
	}
	if !v.Passthrough() {
		t.Fatalf("interceptor should be passthrough after Login Success")
	}
}

func TestVelocityQueryBeforeLoginStartFailsOpen(t *testing.T) {
	// A query arriving before the Login Start was observed cannot be answered:
	// it must reach the player unmodified (who will then be kicked by the
	// backend with its own message).
	v := newVelocityForwarder(testSecret, "192.0.2.1")
	feedLogin(v, 767, intentLogin, nil, 50) // handshake only
	req := mkPluginRequest(5, velocityChannel, []byte{0x04})
	fwd, inject := v.ProcessS2C(req)
	if inject != nil {
		t.Fatalf("must not answer without a Login Start")
	}
	if !bytes.Equal(fwd, req) {
		t.Fatalf("unanswerable query not passed through")
	}
	if !v.Passthrough() {
		t.Fatalf("interceptor should fail open")
	}
}

func TestVelocityOversizedC2SFailsOpen(t *testing.T) {
	v := newVelocityForwarder(testSecret, "192.0.2.1")
	// A declared c2s packet length beyond the sniff cap aborts interception.
	v.ObserveC2S(wire.AppendVarInt(nil, maxC2SSniff+1))
	if !v.Passthrough() {
		t.Fatalf("oversized login packet should turn the stream transparent")
	}
}

func TestVelocityC2SAbortFlushesBufferedS2C(t *testing.T) {
	v := newVelocityForwarder(testSecret, "192.0.2.1")
	feedLogin(v, 767, intentLogin, nil, 50) // handshake only; login pending
	req := mkPluginRequest(5, velocityChannel, []byte{0x04})
	half := len(req) / 2
	if fwd, _ := v.ProcessS2C(req[:half]); len(fwd) != 0 {
		t.Fatalf("partial packet must stay buffered")
	}
	// The c2s side now aborts (e.g. unparseable player bytes) while s2c bytes
	// sit buffered: they must not be lost.
	v.ObserveC2S(wire.AppendVarInt(nil, maxC2SSniff+1))
	fwd, inject := v.ProcessS2C(req[half:])
	if inject != nil {
		t.Fatalf("aborted interceptor must not inject")
	}
	if !bytes.Equal(fwd, req) {
		t.Fatalf("buffered s2c bytes lost on abort:\n got %x\nwant %x", fwd, req)
	}
}