package main

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"sync"
	"testing"
	"time"
)

// fakeSink records what the tap dispatched to one spectator.
type fakeSink struct {
	mu       sync.Mutex
	frames   []uint32
	reliable [][]byte
	tags     []byte
	audio    int
	controls []map[string]any
	ended    string
	full     bool
	backlog  int
}

func (f *fakeSink) enqueueFrame(msg frameMsg) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.frames = append(f.frames, msg.frame)
}

func (f *fakeSink) enqueueReliable(tag byte, data []byte) bool {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.full {
		return false
	}
	f.tags = append(f.tags, tag)
	f.reliable = append(f.reliable, data)
	return true
}

func (f *fakeSink) reliableBacklog() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.backlog
}

func (f *fakeSink) sendAudio([]byte) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.audio++
}

func (f *fakeSink) control(object any) {
	f.mu.Lock()
	defer f.mu.Unlock()
	raw, _ := json.Marshal(object)
	var decoded map[string]any
	_ = json.Unmarshal(raw, &decoded)
	f.controls = append(f.controls, decoded)
}

func (f *fakeSink) end(reason string) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.ended = reason
}

func (f *fakeSink) setFull(full bool) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.full = full
}

func (f *fakeSink) endedReason() string {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.ended
}

func (f *fakeSink) controlTypes() []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	var out []string
	for _, c := range f.controls {
		out = append(out, c["t"].(string))
	}
	return out
}

func (f *fakeSink) countTag(tag byte) int {
	f.mu.Lock()
	defer f.mu.Unlock()
	n := 0
	for _, t := range f.tags {
		if t == tag {
			n++
		}
	}
	return n
}

// frameFragment builds one frame-plane fragment (a minimal valid envelope).
func frameFragment(frame uint32, flags uint16, fragIdx, fragCount uint16) []byte {
	data := make([]byte, envelopeSize)
	binary.LittleEndian.PutUint32(data, envelopeMagic)
	data[4] = envelopeVersion
	data[5] = kindFrame
	binary.LittleEndian.PutUint16(data[6:], flags)
	binary.LittleEndian.PutUint32(data[16:], frame)
	binary.LittleEndian.PutUint16(data[20:], fragIdx)
	binary.LittleEndian.PutUint16(data[22:], fragCount)
	return data
}

func tagged(tag byte, payload []byte) []byte {
	return append([]byte{tag}, payload...)
}

func newTestTap() *seatTap {
	return newSeatTap(3, "t", "host-session", nil, nil)
}

func TestParseFrameMsgReportsKeyframesAndAnchors(t *testing.T) {
	key := parseFrameMsg(frameFragment(7, flagKeyframe, 0, 2))
	if !key.valid || !key.keyframe || key.anchor || !key.reliable ||
		!key.newFrame || key.last || key.fragment != 0 || key.fragments != 2 {
		t.Fatalf("keyframe fragment misparsed: %+v", key)
	}
	anchor := parseFrameMsg(frameFragment(8, flagAnchor, 1, 2))
	if !anchor.valid || anchor.keyframe || !anchor.anchor || !anchor.reliable ||
		anchor.newFrame || !anchor.last || anchor.fragment != 1 ||
		anchor.fragments != 2 {
		t.Fatalf("anchor fragment misparsed: %+v", anchor)
	}
	dependent := parseFrameMsg(frameFragment(9, 0, 0, 1))
	if !dependent.valid || dependent.keyframe || dependent.anchor ||
		dependent.reliable {
		t.Fatal("dependent frame flagged as reliable")
	}
	zeroFragments := parseFrameMsg(frameFragment(10, flagKeyframe, 0, 0))
	if zeroFragments.valid {
		t.Fatalf("zero-fragment envelope accepted: %+v", zeroFragments)
	}
	outOfRange := parseFrameMsg(frameFragment(11, flagKeyframe, 2, 2))
	if outOfRange.valid {
		t.Fatalf("out-of-range fragment accepted: %+v", outOfRange)
	}
}

func TestEpochGateAdmitsNothingBeforeTheEpochAndFramesOnlyAfterTheKeyframe(t *testing.T) {
	g := &gate{mode: gateEpoch}
	if g.admit(tagResource, frameMsg{}, 0) || g.admit(tagFrame, frameMsg{}, 0) ||
		g.admit(tagAudio, frameMsg{}, 0) {
		t.Fatal("a detached spectator must receive nothing")
	}
	g.attached = true
	if !g.admit(tagResource, frameMsg{}, 0) {
		t.Fatal("resources flow from the epoch marker on")
	}
	if g.admit(tagFrame, frameMsg{frame: 1}, 0) || g.admit(tagAudio, frameMsg{}, 0) {
		t.Fatal("frames and audio must wait for the keyframe")
	}
	if g.admit(tagRecovery, frameMsg{frame: 1, newFrame: true}, 0) {
		t.Fatal("a non-keyframe recovery message must not open the gate")
	}
	if !g.admit(tagRecovery, frameMsg{frame: 2, newFrame: true, keyframe: true}, 0) {
		t.Fatal("the first keyframe fragment opens the gate")
	}
	if !g.admit(tagRecovery, frameMsg{frame: 2, keyframe: true, last: true}, 0) ||
		!g.admit(tagFrame, frameMsg{frame: 3}, 0) || !g.admit(tagAudio, frameMsg{}, 0) {
		t.Fatal("later fragments, frames and audio flow once the keyframe is seen")
	}
}

func TestMirrorGateOpensFramesAtAnAnchorOnceTheBundleHasLeft(t *testing.T) {
	g := &gate{mode: gateMirror, attached: true}
	if !g.admit(tagResource, frameMsg{}, 0) || !g.admit(tagRecovery,
		frameMsg{frame: 5, newFrame: true, keyframe: true}, 0) {
		t.Fatal("reliable planes flow at once")
	}
	if g.admit(tagFrame, frameMsg{frame: 6, newFrame: true}, 0) {
		t.Fatal("a dependent frame cannot open the gate")
	}
	if g.admit(tagFrame, frameMsg{frame: 7, newFrame: true, anchor: true},
		frameGateBacklog+1) {
		t.Fatal("frames wait while the join bundle is still queued")
	}
	if g.admit(tagAudio, frameMsg{}, 0) {
		t.Fatal("audio waits with the frames")
	}
	if !g.admit(tagFrame, frameMsg{frame: 8, newFrame: true, anchor: true}, 0) {
		t.Fatal("the first anchor after the drain opens the gate")
	}
	if !g.admit(tagFrame, frameMsg{frame: 9, newFrame: true}, 0) ||
		!g.admit(tagAudio, frameMsg{}, 0) {
		t.Fatal("everything flows afterwards")
	}
}

// primeMirror drives a tap through a host bootstrap: epoch marker, one
// resource message, one keyframe. Returns the resource message's end
// sequence (the mirror's consistency point).
func primeMirror(t *testing.T, tap *seatTap) uint32 {
	t.Helper()
	tap.handleText([]byte(`{"t":"spectate","token":"t","seat":3}`))
	tap.handleText([]byte(`{"t":"epoch","id":1}`))
	for _, f := range resourceMessage(7, 4000,
		createTexture(2, 1, 4, 4, 1, 21),
		updateTexture(2, 2, 0, 0, 0, 4, 4, 21, fill(64, 0x5a))) {
		tap.handleBinary(tagged(tagResource, f))
	}
	for _, f := range keyframeMessage(7, 1, 100, fill(40, 0)) {
		tap.handleBinary(tagged(tagRecovery, f))
	}
	if !tap.mirror.ready() {
		t.Fatal("mirror must be ready after epoch + keyframe")
	}
	return 4000
}

func TestJoinIsServedFromTheMirrorThenFollowsTheLiveTap(t *testing.T) {
	tap := newTestTap()
	s0 := primeMirror(t, tap)
	sink := &fakeSink{}
	tap.add(sink)
	tap.servePending()

	if got := sink.controlTypes(); len(got) != 1 || got[0] != "joined" {
		t.Fatalf("expected a joined control, got %v", got)
	}
	snapshot := sink.countTag(tagResource)
	if snapshot == 0 || sink.countTag(tagRecovery) != 1 {
		t.Fatalf("bundle: %d resource, %d recovery fragments", snapshot,
			sink.countTag(tagRecovery))
	}
	// The snapshot fragments reproduce the mirrored texture and end at s0.
	clone := decodeSnapshot(t, sink.reliable[:snapshot], s0)
	if clone.textures[2] == nil || clone.textures[2].mips[0].bytes[0] != 0x5a {
		t.Fatal("snapshot does not carry the texture")
	}
	last := sink.reliable[len(sink.reliable)-1]
	if binary.LittleEndian.Uint16(last[6:])&flagDictionaryOnly == 0 {
		t.Fatal("the keyframe copy must be dictionary-only")
	}
	tap.mu.Lock()
	g := tap.sinks[sink]
	tap.mu.Unlock()
	if g.mode != gateMirror || !g.attached || g.framesOpen {
		t.Fatalf("gate after join: %+v", *g)
	}

	// Live: a resource message flows; a dependent frame is held; the
	// first anchor opens frames (and audio); the host never heard a thing.
	for _, f := range resourceMessage(7, s0+1, defineBlob(9, fill(8, 9))) {
		tap.handleBinary(tagged(tagResource, f))
	}
	if sink.countTag(tagResource) != snapshot+1 {
		t.Fatal("live resource messages flow after the bundle")
	}
	tap.handleBinary(tagged(tagFrame, frameFragment(101, 0, 0, 1)))
	tap.handleBinary(tagged(tagAudio, []byte{1, 2, 3}))
	if len(sink.frames) != 0 || sink.audio != 0 {
		t.Fatal("dependent frames and audio wait for the first anchor")
	}
	tap.handleBinary(tagged(tagFrame, frameFragment(102, flagAnchor, 0, 1)))
	tap.handleBinary(tagged(tagFrame, frameFragment(103, 0, 0, 1)))
	tap.handleBinary(tagged(tagAudio, []byte{1, 2, 3}))
	if len(sink.frames) != 2 || sink.frames[0] != 102 || sink.audio != 1 {
		t.Fatalf("frames after the anchor: %v audio=%d", sink.frames, sink.audio)
	}
	if !tap.lastJoin.IsZero() {
		t.Fatal("a mirror join must never ask the relay for an epoch")
	}
	if tap.joinsMirror != 1 || tap.joinsEpoch != 0 {
		t.Fatalf("join accounting %d/%d", tap.joinsMirror, tap.joinsEpoch)
	}
}

func TestJoinWaitsForAKeyframeInFlight(t *testing.T) {
	tap := newTestTap()
	primeMirror(t, tap)
	// A new keyframe's first fragment is on the tap; its second is not.
	parts := keyframeMessage(7, 2, 200, fill(100*1024, 1))
	if len(parts) < 2 {
		t.Fatal("test needs a multi-fragment keyframe")
	}
	tap.handleBinary(tagged(tagRecovery, parts[0]))
	sink := &fakeSink{}
	tap.add(sink)
	tap.servePending()
	if len(sink.controls) != 0 {
		t.Fatal("no join while a keyframe is mid-flight")
	}
	for _, part := range parts[1:] {
		tap.handleBinary(tagged(tagRecovery, part))
	}
	tap.servePending()
	if got := sink.controlTypes(); len(got) != 1 || got[0] != "joined" {
		t.Fatalf("join after the keyframe completed: %v", got)
	}
	if sink.countTag(tagRecovery) != len(parts) {
		t.Fatalf("the bundle carries the NEW keyframe (%d fragments), got %d",
			len(parts), sink.countTag(tagRecovery))
	}
	first := sink.reliable[len(sink.reliable)-len(parts)]
	if e, _, _ := parseEnvelope(first); e.frame != 200 {
		t.Fatalf("bundle keyframe frame %d, want 200", e.frame)
	}
}

func TestResetRegatesEveryoneAndTheNextEpochServesThemAgain(t *testing.T) {
	tap := newTestTap()
	primeMirror(t, tap)
	sink := &fakeSink{}
	tap.add(sink)
	tap.servePending()
	before := len(sink.reliable)

	tap.handleText([]byte(`{"t":"reset"}`))
	if got := sink.controlTypes(); got[len(got)-1] != "epochReset" {
		t.Fatalf("reset must tell the spectator: %v", got)
	}
	tap.handleBinary(tagged(tagFrame, frameFragment(300, flagAnchor, 0, 1)))
	if len(sink.frames) != 0 {
		t.Fatal("nothing flows while detached")
	}
	tap.servePending()
	if len(sink.reliable) != before {
		t.Fatal("no bundle before the host's new epoch is complete")
	}
	if !tap.lastJoin.IsZero() {
		t.Fatal("a reset is the host's own epoch; never ask for another")
	}
	primeMirror(t, tap)
	tap.servePending()
	if len(sink.reliable) <= before {
		t.Fatal("the new epoch's state is served from the mirror")
	}
	if got := sink.controlTypes(); got[len(got)-1] != "joined" {
		t.Fatalf("expected joined after the new epoch: %v", got)
	}
	tap.handleText([]byte(`{"t":"sessionEnded","reason":"replay-ended"}`))
	if sink.endedReason() != "replay-ended" {
		t.Fatalf("ended=%q", sink.endedReason())
	}
}

func TestLegacyRecoveryRequiresFreshSpectatorSession(t *testing.T) {
	tap := newTestTap()
	primeMirror(t, tap)
	sink := &fakeSink{}
	tap.add(sink)
	tap.servePending()
	first := len(sink.reliable)
	if !tap.retireLegacyRecovery(sink, "anchor-stall") {
		t.Fatal("an anchor stall is a recognized legacy recovery")
	}
	tap.servePending()
	if len(sink.reliable) != first {
		t.Fatal("the mirror must not append behind the old reliable queue")
	}
	if sink.endedReason() != "stream-ended" {
		t.Fatalf("legacy spectator was not retired: %q", sink.endedReason())
	}
	other := &fakeSink{}
	if tap.retireLegacyRecovery(other, "missing-blob") {
		t.Fatal("an unrecognized reason must not close the spectator")
	}
	if other.endedReason() != "" {
		t.Fatal("an ignored request changed the spectator")
	}
}

func TestWithoutAnEpochTheTapFallsBackToTheInBandJoin(t *testing.T) {
	tap := newTestTap()
	sink := &fakeSink{}
	tap.add(sink)
	tap.servePending()
	if tap.lastJoin.IsZero() {
		t.Fatal("a tap that never saw an epoch must ask the relay for one")
	}
	tap.mu.Lock()
	mode := tap.sinks[sink].mode
	tap.mu.Unlock()
	if mode != gateEpoch {
		t.Fatalf("spectator must wait at the marker, mode=%d", mode)
	}
	// Resource bytes before the marker are not for this spectator.
	for _, f := range resourceMessage(7, 10, createTexture(1, 1, 4, 4, 1, 21)) {
		tap.handleBinary(tagged(tagResource, f))
	}
	if len(sink.reliable) != 0 {
		t.Fatal("nothing before the epoch marker")
	}
	tap.handleText([]byte(`{"t":"epoch","id":1}`))
	for _, f := range resourceMessage(7, 11, createTexture(2, 1, 4, 4, 1, 21)) {
		tap.handleBinary(tagged(tagResource, f))
	}
	if sink.countTag(tagResource) == 0 {
		t.Fatal("the in-band snapshot flows from the marker on")
	}
	tap.handleBinary(tagged(tagFrame, frameFragment(50, 0, 0, 1)))
	for _, f := range keyframeMessage(7, 1, 51, fill(40, 0)) {
		tap.handleBinary(tagged(tagRecovery, f))
	}
	tap.handleBinary(tagged(tagFrame, frameFragment(52, 0, 0, 1)))
	if len(sink.frames) != 1 || sink.frames[0] != 52 {
		t.Fatalf("frames after the epoch's keyframe only: %v", sink.frames)
	}
	if tap.joinsEpoch != 1 {
		t.Fatalf("epoch join accounting %d", tap.joinsEpoch)
	}
	// That epoch also primed the mirror: the next spectator joins from it.
	late := &fakeSink{}
	tap.add(late)
	tap.servePending()
	if got := late.controlTypes(); len(got) != 1 || got[0] != "joined" {
		t.Fatalf("late spectator must join from the mirror: %v", got)
	}
}

func TestSlowSpectatorIsDroppedNotWaitedFor(t *testing.T) {
	tap := newTestTap()
	primeMirror(t, tap)
	sink := &fakeSink{}
	tap.add(sink)
	tap.servePending()
	sink.setFull(true)
	for _, f := range resourceMessage(7, 4001, defineBlob(9, fill(8, 9))) {
		tap.handleBinary(tagged(tagResource, f))
	}
	time.Sleep(20 * time.Millisecond)
	if sink.endedReason() != "too-slow" || tap.dropped != 1 {
		t.Fatalf("ended=%q dropped=%d", sink.endedReason(), tap.dropped)
	}
	tap.mu.Lock()
	_, still := tap.sinks[sink]
	tap.mu.Unlock()
	if still {
		t.Fatal("a dropped spectator leaves the fan-out")
	}
}

func TestHostSessionsNeverTouchTheSpectatorCaps(t *testing.T) {
	if spectatorsPerSeat != 0 || spectatorsPerBox != 0 || spectatorsPerIP != 0 {
		t.Fatal("spectators are unlimited by default (soft seats)")
	}
	if spectatorTotal() != 0 {
		t.Fatal("no spectators before any tap")
	}
}

// A tab closing kills the host's signaling socket and its QUIC session at
// once, so the relay's {"t":"sessionEnded"} marker and the bridge's own
// teardown race. Either way a spectator must be told the host left (a
// terminal end) and never "stream-ended", which asks the page to rejoin a
// seat nobody holds.
func TestHostSessionEndTellsSpectatorsTheHostLeft(t *testing.T) {
	for _, source := range []string{"marker", "context"} {
		tap := newTestTap()
		primeMirror(t, tap)
		sink := &fakeSink{}
		tap.add(sink)
		tap.servePending()
		if source == "marker" {
			tap.handleText([]byte(`{"t":"sessionEnded","reason":"host-left"}`))
		} else {
			// The tap's own context ends with the bridge's host session.
			ctx, cancel := context.WithCancel(context.Background())
			cancel()
			tap.shutdown(tap.endReason(ctx))
		}
		if sink.ended != "host-left" {
			t.Fatalf("%s: spectator ended with %q, want host-left",
				source, sink.ended)
		}
	}
}

// A tap that loses its relay connection while the host session lives (a
// relay restart, a replaced tap) is a fan-out fault: the spectator rejoins.
func TestTapLossWhileTheHostLivesAsksForARejoin(t *testing.T) {
	tap := newTestTap()
	primeMirror(t, tap)
	sink := &fakeSink{}
	tap.add(sink)
	tap.servePending()
	tap.shutdown(tap.endReason(context.Background()))
	if sink.ended != "stream-ended" {
		t.Fatalf("spectator ended with %q, want stream-ended", sink.ended)
	}
}