package main

import (
	"testing"

	"github.com/yuanyuanxiang/SimpleRemoter/server/go/connection"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/hub"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/protocol"
)

type frameRecorder struct {
	frames []recordedFrame
}

type recordedFrame struct {
	deviceID string
	packet   []byte
	isKey    bool
}

func (r *frameRecorder) OnScreenFrame(deviceID string, packet []byte, isKeyframe bool) {
	r.frames = append(r.frames, recordedFrame{deviceID, packet, isKeyframe})
}
func (r *frameRecorder) OnDeviceOnline(hub.DeviceInfo)       {}
func (r *frameRecorder) OnDeviceOffline(string)              {}
func (r *frameRecorder) OnDeviceUpdate(string, int, string)  {}
func (r *frameRecorder) OnResolutionChange(string, int, int) {}
func (r *frameRecorder) OnCursorChange(string, byte)         {}
func (r *frameRecorder) OnClipboardText(string, string)      {}
func (r *frameRecorder) OnTerminalReady(string, bool)        {}
func (r *frameRecorder) OnTerminalData(string, []byte)       {}
func (r *frameRecorder) OnTerminalClosed(string, string)     {}

// TestOnReceiveRelaysKeyframeToken guards the forced-IDR recovery path: when
// the browser asks for a keyframe, the client replies with TOKEN_KEYFRAME.
// Dropping that token (the old no-op) leaves the browser decoding P-frames
// against a reference chain it never received, i.e. persistent corruption.
func TestOnReceiveRelaysKeyframeToken(t *testing.T) {
	h := hub.New()
	rec := &frameRecorder{}
	h.Subscribe(rec)

	h.Register(&hub.Device{ID: "dev", PeerIP: "10.0.0.1"},
		&connection.Context{RemoteAddr: "10.0.0.1:1000"})
	screenCtx := &connection.Context{RemoteAddr: "10.0.0.1:1001"}
	if !h.BindScreenConn("dev", screenCtx) {
		t.Fatal("screen bind failed")
	}

	// [TOKEN_KEYFRAME][algo=H264][cursorPos:8][cursorIdx:1][IDR NAL]
	data := []byte{
		protocol.TokenKeyframe, 2,
		0, 0, 0, 0, 0, 0, 0, 0,
		0,
		0x00, 0x00, 0x00, 0x01, 0x65, 0x88, 0x84,
	}
	(&MyHandler{hub: h}).OnReceive(screenCtx, data)

	if len(rec.frames) != 1 {
		t.Fatalf("frames relayed = %d, want 1 (TOKEN_KEYFRAME must not be dropped)", len(rec.frames))
	}
	f := rec.frames[0]
	if f.deviceID != "dev" {
		t.Fatalf("deviceID = %q, want dev", f.deviceID)
	}
	if !f.isKey || f.packet[4] != 1 {
		t.Fatalf("isKeyframe=%v frameType=%d, want true/1", f.isKey, f.packet[4])
	}
	if want := 9 + len(data) - (1 + protocol.ScreenFrameHeaderLen); len(f.packet) != want {
		t.Fatalf("packet len = %d, want %d", len(f.packet), want)
	}
}