package e2e

import (
	"bytes"
	"fmt"
	"io"
	"math/rand"
	"net"
	"sync"
	"testing"
	"time"

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

// resumePlayer opens a player connection and returns it, having sent only the
// handshake. The caller keeps it open across the outage — which is the whole
// point: before stream resumption this socket was closed by the hub the instant
// its worker conn died.
func resumePlayer(t *testing.T, hubAddr, host string) net.Conn {
	t.Helper()
	pc, err := net.DialTimeout("tcp", hubAddr, 5*time.Second)
	if err != nil {
		t.Fatalf("player dial: %v", err)
	}
	t.Cleanup(func() { _ = pc.Close() })
	if _, err := pc.Write(wire.BuildHandshake(767, host, 25565, 2)); err != nil {
		t.Fatalf("player handshake: %v", err)
	}
	return pc
}

// echoExchange streams payload through an echoing destination and verifies that
// what comes back is byte-for-byte identical, calling disrupt once `at` bytes
// have made the round trip.
//
// Comparing the whole stream rather than sampling is deliberate: a resumption
// bug does not corrupt bytes, it duplicates or skips a range, and only an exact
// comparison of the full sequence catches an off-by-one in the offsets.
func echoExchange(t *testing.T, pc net.Conn, payload []byte, at int, disrupt func()) {
	t.Helper()
	const chunk = 16 << 10

	var wg sync.WaitGroup
	wg.Add(1)
	writeErr := make(chan error, 1)
	go func() {
		defer wg.Done()
		for off := 0; off < len(payload); off += chunk {
			end := min(off+chunk, len(payload))
			_ = pc.SetWriteDeadline(time.Now().Add(60 * time.Second))
			if _, err := pc.Write(payload[off:end]); err != nil {
				writeErr <- fmt.Errorf("write at %d: %w", off, err)
				return
			}
		}
		writeErr <- nil
	}()

	got := make([]byte, len(payload))
	read, fired := 0, false
	for read < len(got) {
		_ = pc.SetReadDeadline(time.Now().Add(60 * time.Second))
		n, err := pc.Read(got[read:])
		read += n
		if !fired && read >= at {
			fired = true
			disrupt()
		}
		if err != nil {
			t.Fatalf("player read failed after %d/%d bytes: %v", read, len(got), err)
		}
	}
	wg.Wait()
	if err := <-writeErr; err != nil {
		t.Fatalf("player write: %v", err)
	}

	if !bytes.Equal(got, payload) {
		// Report the first divergence: its offset says whether the stream gained
		// or lost bytes, which is the difference between a retransmit that
		// replayed too much and one that replayed too little.
		for i := range got {
			if got[i] != payload[i] {
				t.Fatalf("echo diverges at byte %d of %d (sent %#x, got %#x)",
					i, len(payload), payload[i], got[i])
			}
		}
	}
}

func randomPayload(seed int64, n int) []byte {
	p := make([]byte, n)
	rand.New(rand.NewSource(seed)).Read(p)
	return p
}

// TestResumePreservesByteStream is the correctness bar for stream resumption:
// a worker conn is hard-reset mid-transfer and the *same* player connection must
// keep working, with a byte stream that neither gains nor loses a single byte.
//
// Byte-exactness is the whole difficulty. Frames handed to a dying socket are
// lost with no notification and the cipher cannot be resynchronized, so each
// side has to replay from the offset the other reports it accepted. Getting that
// offset wrong by any amount splices the stream mid-Minecraft-packet, which a
// round-trip test that only checked "traffic flows again" would happily pass.
func TestResumePreservesByteStream(t *testing.T) {
	const psk = "e2e-resume"
	hubPort := freePort(t)
	hubAddr := fmt.Sprintf("127.0.0.1:%d", hubPort)
	startHub(t, hubPort, psk)
	dest := newMockDest(t, modeEcho)

	// Only the tunnel runs through the relay; the player talks to the hub
	// directly, as it would from the internet. So the drop hits the middle leg
	// while both terminal sockets stay healthy — exactly the case resumption is
	// for.
	relay := newBlackholeRelay(t, hubAddr)
	c := startClientWithPing(t, relay.addr, psk, 1, 400, []client.Mapping{
		{Pattern: "mc.local", Destination: dest.addr},
	})

	pc := resumePlayer(t, hubAddr, "mc.local")
	payload := randomPayload(1, 3<<20)
	echoExchange(t, pc, payload, 512<<10, func() {
		t.Log("hard-resetting the tunnel mid-transfer")
		relay.dropAll()
	})
	t.Logf("stream survived the reset intact; worker conns now %d", c.WorkerConnCount())
}

// TestResumeWithConcurrentStreams covers the failure the single-stream test
// cannot reach: two independent tunnels resume at once. A torn rebind that
// writes a frame onto the wrong conn would credit or reset the other player's
// tunnel, and that only shows up when a second player is there to be corrupted.
func TestResumeWithConcurrentStreams(t *testing.T) {
	const psk = "e2e-resume-multi"
	hubPort := freePort(t)
	hubAddr := fmt.Sprintf("127.0.0.1:%d", hubPort)
	startHub(t, hubPort, psk)
	dest := newMockDest(t, modeEcho)

	relay := newBlackholeRelay(t, hubAddr)
	startClientWithPing(t, relay.addr, psk, 4, 400, []client.Mapping{
		{Pattern: "mc.local", Destination: dest.addr},
	})

	const players = 3
	conns := make([]net.Conn, players)
	for i := range conns {
		conns[i] = resumePlayer(t, hubAddr, "mc.local")
		// Distinct payloads: if a reattach crosses two streams the bytes land on
		// the wrong player, which an identical payload would hide.
		if _, err := conns[i].Write([]byte(fmt.Sprintf("hello-%d", i))); err != nil {
			t.Fatalf("player %d warmup write: %v", i, err)
		}
		echo := make([]byte, len("hello-0"))
		_ = conns[i].SetReadDeadline(time.Now().Add(15 * time.Second))
		if _, err := io.ReadFull(conns[i], echo); err != nil {
			t.Fatalf("player %d warmup echo: %v", i, err)
		}
	}

	var wg sync.WaitGroup
	for i := range conns {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			payload := randomPayload(int64(100+i), 768<<10)
			// Only the first player triggers the reset; the others are mid-flight
			// when it lands.
			disrupt := func() {}
			if i == 0 {
				disrupt = relay.dropAll
			}
			echoExchange(t, conns[i], payload, 128<<10, disrupt)
		}(i)
	}
	wg.Wait()
}

// TestResumeDisabledClosesImmediately pins the off switch. With resumption
// declined the hub must not hang the player waiting for a reattach that is never
// coming: the socket has to close as it did before the feature existed, so
// disabling it is a true revert rather than a slower failure.
func TestResumeDisabledClosesImmediately(t *testing.T) {
	const psk = "e2e-resume-off"
	hubPort := freePort(t)
	hubAddr := fmt.Sprintf("127.0.0.1:%d", hubPort)
	startHub(t, hubPort, psk)
	dest := newMockDest(t, modeEcho)

	relay := newBlackholeRelay(t, hubAddr)
	off := false
	startClientCfg(t, &client.Config{
		Server:         relay.addr,
		PSK:            psk,
		MaxTunnels:     1,
		PingIntervalMs: 400,
		StreamResume:   &off,
		Mappings:       []client.Mapping{{Pattern: "mc.local", Destination: dest.addr}},
	})

	pc := resumePlayer(t, hubAddr, "mc.local")
	if _, err := pc.Write([]byte("ping")); err != nil {
		t.Fatalf("warmup write: %v", err)
	}
	echo := make([]byte, 4)
	_ = pc.SetReadDeadline(time.Now().Add(15 * time.Second))
	if _, err := io.ReadFull(pc, echo); err != nil {
		t.Fatalf("warmup echo: %v", err)
	}

	relay.dropAll()

	// Well inside the hub's 20s grace: if the player is still open here, the hub
	// parked a stream for a client that never opted in.
	_ = pc.SetReadDeadline(time.Now().Add(10 * time.Second))
	if _, err := pc.Read(make([]byte, 1)); err == nil {
		t.Fatal("player socket stayed open after the tunnel dropped with resume disabled")
	}
}

// TestResumeGraceExpiryClosesPlayer covers the other end of the lifetime: when
// the tunnel never comes back, a hung player must not hang forever. The hub
// drops it once its grace expires, without leaking the stream or its buffers.
func TestResumeGraceExpiryClosesPlayer(t *testing.T) {
	const psk = "e2e-resume-grace"
	hubPort := freePort(t)
	hubAddr := fmt.Sprintf("127.0.0.1:%d", hubPort)
	startHubCfg(t, hubPort, psk, map[string]any{"resumeGraceMs": 3000})
	dest := newMockDest(t, modeEcho)

	relay := newBlackholeRelay(t, hubAddr)
	startClientCfg(t, &client.Config{
		Server:         relay.addr,
		PSK:            psk,
		MaxTunnels:     1,
		PingIntervalMs: 400,
		ResumeGraceMs:  2000,
		Mappings:       []client.Mapping{{Pattern: "mc.local", Destination: dest.addr}},
	})

	pc := resumePlayer(t, hubAddr, "mc.local")
	if _, err := pc.Write([]byte("ping")); err != nil {
		t.Fatalf("warmup write: %v", err)
	}
	echo := make([]byte, 4)
	_ = pc.SetReadDeadline(time.Now().Add(15 * time.Second))
	if _, err := io.ReadFull(pc, echo); err != nil {
		t.Fatalf("warmup echo: %v", err)
	}

	// Stop carrying the tunnel entirely, so every reattach attempt fails.
	relay.stop()
	relay.dropAll()

	start := time.Now()
	_ = pc.SetReadDeadline(time.Now().Add(30 * time.Second))
	if _, err := pc.Read(make([]byte, 1)); err == nil {
		t.Fatal("player socket never closed after the resume grace expired")
	}
	if elapsed := time.Since(start); elapsed < time.Second {
		t.Fatalf("player closed after %s, before any reattach could be attempted", elapsed)
	} else {
		t.Logf("hung player released after %s", elapsed.Round(100*time.Millisecond))
	}
}