package pdu

import (
	"bytes"
	"encoding/binary"
	"testing"
)

// buildSurfCmd 构造一条 SET/STREAM_SURFACE_BITS 命令。
// win10=true 用 Win10 实测 22 字节头（codecID@13），否则用
// FreeRDP 经典 21 字节头（codecID@12）。
func buildSurfCmd(cmdType uint16, codecID, bpp byte, w, h uint16, payload []byte, win10 bool) []byte {
	hdrLen := 21
	extra := 0
	if win10 {
		hdrLen = 22
		extra = 1
	}
	buf := make([]byte, hdrLen+len(payload))
	binary.LittleEndian.PutUint16(buf[0:], cmdType)
	binary.LittleEndian.PutUint16(buf[2:], 10)     // destLeft
	binary.LittleEndian.PutUint16(buf[4:], 20)     // destTop
	binary.LittleEndian.PutUint16(buf[6:], 10+w-1) // destRight
	binary.LittleEndian.PutUint16(buf[8:], 20+h-1) // destBottom
	buf[10] = bpp                                  // bpp
	buf[11] = 0                                    // reserved
	buf[12+extra] = codecID
	binary.LittleEndian.PutUint16(buf[13+extra:], w)
	binary.LittleEndian.PutUint16(buf[15+extra:], h)
	binary.LittleEndian.PutUint32(buf[17+extra:], uint32(len(payload)))
	copy(buf[hdrLen:], payload)
	return buf
}

func buildFrameMarker(action uint16, frameID uint32) []byte {
	buf := make([]byte, 8)
	binary.LittleEndian.PutUint16(buf[0:], CMDTYPE_FRAME_MARKER)
	binary.LittleEndian.PutUint16(buf[2:], action)
	binary.LittleEndian.PutUint32(buf[4:], frameID)
	return buf
}

// Win10 实测 22 字节头布局回归：cmdType(2)+dest(8)+bpp(1)+?@11+?@12+
// codecID@13+width@14+height@16+len u32@18。偏移若错位，
// codecID/width/length 断言即失败。
func TestParseSurfaceCommandsWin10Layout(t *testing.T) {
	payload := make([]byte, 2*2*4) // 2x2 32bpp 原始像素
	buf := buildSurfCmd(CMDTYPE_SET_SURFACE_BITS, 0, 32, 2, 2, payload, true)

	result, consumed, needMore, valid, _, _ := parseSurfaceCommandsIncremental(buf, 0)
	if !valid || needMore {
		t.Fatalf("valid=%v needMore=%v", valid, needMore)
	}
	if consumed != len(buf) {
		t.Fatalf("consumed=%d, want %d", consumed, len(buf))
	}
	if len(result.Rects) != 1 {
		t.Fatalf("rects=%d, want 1", len(result.Rects))
	}
	r := result.Rects[0]
	if r.Width != 2 || r.Height != 2 || r.BitsPerPixel != 32 {
		t.Fatalf("size=%dx%d bpp=%d, want 2x2/32", r.Width, r.Height, r.BitsPerPixel)
	}
	if r.DestLeft != 10 || r.DestTop != 20 {
		t.Fatalf("dest=(%d,%d), want (10,20)", r.DestLeft, r.DestTop)
	}
}

// FreeRDP 经典 21 字节头（codecID@12）的自适应回退。
func TestParseSurfaceCommandsClassicLayout(t *testing.T) {
	buf := buildSurfCmd(CMDTYPE_SET_SURFACE_BITS, 0, 32, 2, 2,
		make([]byte, 2*2*4), false)
	result, consumed, needMore, valid, _, _ := parseSurfaceCommandsIncremental(buf, 0)
	if !valid || needMore || consumed != len(buf) {
		t.Fatalf("valid=%v needMore=%v consumed=%d/%d", valid, needMore, consumed, len(buf))
	}
	if len(result.Rects) != 1 {
		t.Fatalf("rects=%d, want 1", len(result.Rects))
	}
	if result.Rects[0].Width != 2 || result.Rects[0].Height != 2 {
		t.Fatalf("size=%dx%d, want 2x2", result.Rects[0].Width, result.Rects[0].Height)
	}
}

// 真实抓包（Win10 19041, 24bpp 会话, 64x64 NSCodec 瓦片）：
// 平面长度 4096+448+339+0=4883, bitmapDataLength=4903=20+4883。
func TestParseSurfaceCommandsRealCapture(t *testing.T) {
	// 真实命令前 48 字节（SURFCMD_DUMP 捕获），payload 截取自转储
	head := []byte{
		0x01, 0x00, 0x40, 0x04, 0x40, 0x02, 0x80, 0x04, 0x80, 0x02,
		0x20, 0x00, 0x00, 0x01, 0x40, 0x00, 0x40, 0x00, 0x27, 0x13,
		0x00, 0x00,
	}
	nsHdr := []byte{ // NSCodec 流头：四平面长度+颜色参数
		0x00, 0x10, 0x00, 0x00, 0xC0, 0x01, 0x00, 0x00,
		0x53, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		0x03, 0x01, 0x00, 0x00,
	}
	// bitmapData = NSCodec 流头(20) + 四平面数据(4883) = 4903
	buf := append(head, nsHdr...)
	buf = append(buf, make([]byte, 4883)...)
	if len(buf) != 22+4903 {
		t.Fatalf("capture length %d, want %d", len(buf), 22+4903)
	}
	result, consumed, needMore, valid, _, _ := parseSurfaceCommandsIncremental(buf, 0)
	if !valid || needMore {
		t.Fatalf("valid=%v needMore=%v", valid, needMore)
	}
	if consumed != len(buf) {
		t.Fatalf("consumed=%d, want %d", consumed, len(buf))
	}
	if len(result.Rects) != 1 {
		t.Fatalf("rects=%d, want 1 (codecID=1 NSCodec)", len(result.Rects))
	}
	r := result.Rects[0]
	if r.DestLeft != 1088 || r.DestTop != 576 || r.Width != 64 || r.Height != 64 {
		t.Fatalf("rect=(%d,%d %dx%d), want (1088,576 64x64)",
			r.DestLeft, r.DestTop, r.Width, r.Height)
	}
	if r.BitsPerPixel != 32 {
		t.Fatalf("bpp=%d, want 32 (NSCodec 输出)", r.BitsPerPixel)
	}
}

func TestParseSurfaceCommandsRoundtrip(t *testing.T) {
	payload := make([]byte, 2*2*4) // 2x2 32bpp 原始像素
	buf := append(buildFrameMarker(SURFCMD_FRAMEACTION_BEGIN, 7),
		buildSurfCmd(CMDTYPE_SET_SURFACE_BITS, 0, 32, 2, 2, payload, true)...)
	buf = append(buf, buildFrameMarker(SURFCMD_FRAMEACTION_END, 7)...)

	result, consumed, needMore, valid, _, _ := parseSurfaceCommandsIncremental(buf, 0)
	if !valid {
		t.Fatal("buffer should parse as valid command stream")
	}
	if needMore {
		t.Fatal("complete stream should not need more data")
	}
	if consumed != len(buf) {
		t.Fatalf("consumed=%d, want %d", consumed, len(buf))
	}
	if len(result.Rects) != 1 {
		t.Fatalf("rects=%d, want 1", len(result.Rects))
	}
	r := result.Rects[0]
	if r.DestLeft != 10 || r.DestTop != 20 {
		t.Fatalf("dest=(%d,%d), want (10,20)", r.DestLeft, r.DestTop)
	}
	if r.Width != 2 || r.Height != 2 {
		t.Fatalf("size=%dx%d, want 2x2", r.Width, r.Height)
	}
	if r.BitsPerPixel != 32 {
		t.Fatalf("bpp=%d, want 32", r.BitsPerPixel)
	}
	if len(result.FrameIDs) != 1 || result.FrameIDs[0] != 7 {
		t.Fatalf("frameIDs=%v, want [7]", result.FrameIDs)
	}
}

// Win10 会把一条命令拆到多个 fast-path PDU：前半段必须 needMore（且
// 缓冲不得被调用方丢弃），后半段到齐后完整解析且 consumed 精确。
func TestParseSurfaceCommandsSplitAcrossSegments(t *testing.T) {
	cmd := buildSurfCmd(CMDTYPE_STREAM_SURFACE_BITS, 0, 32, 4, 2,
		make([]byte, 4*2*4), true)
	cut := 24 // 落在 payload 中间

	r1, c1, more1, valid1, _, _ := parseSurfaceCommandsIncremental(cmd[:cut], 0)
	if !valid1 || !more1 || c1 != 0 || len(r1.Rects) != 0 {
		t.Fatalf("seg1: valid=%v more=%v consumed=%d rects=%d",
			valid1, more1, c1, len(r1.Rects))
	}

	result, consumed, needMore, valid, _, _ := parseSurfaceCommandsIncremental(cmd, 0)
	if !valid || needMore {
		t.Fatalf("seg2: valid=%v needMore=%v", valid, needMore)
	}
	if consumed != len(cmd) || len(result.Rects) != 1 {
		t.Fatalf("seg2: consumed=%d/%d rects=%d",
			consumed, len(cmd), len(result.Rects))
	}
}

// codecID=1 (NSCodec) 的 payload 即使解码失败，分帧也必须精确消费，
// 否则累积缓冲错位后整条流报废（黑屏的伴随症状）。
func TestParseSurfaceCommandsNSCodecFraming(t *testing.T) {
	cmd := buildSurfCmd(CMDTYPE_SET_SURFACE_BITS, 1, 32, 800, 480,
		make([]byte, 100), true) // 非 NSCodec 合法数据 → 解码可能失败
	result, consumed, needMore, valid, _, _ := parseSurfaceCommandsIncremental(cmd, 0)
	if !valid || needMore {
		t.Fatalf("valid=%v needMore=%v", valid, needMore)
	}
	if consumed != len(cmd) {
		t.Fatalf("consumed=%d, want %d", consumed, len(cmd))
	}
	if len(result.Rects) > 1 {
		t.Fatalf("rects=%d, want at most 1", len(result.Rects))
	}
}

// 长度字段读到离谱值时应判无效并丢弃缓冲，而不是无限累积。
func TestParseSurfaceCommandsRejectsBadLength(t *testing.T) {
	cmd := buildSurfCmd(CMDTYPE_SET_SURFACE_BITS, 0, 32, 2, 2,
		make([]byte, 16), true)
	binary.LittleEndian.PutUint32(cmd[18:], 1<<30) // 1GB，非法
	_, _, _, valid, _, _ := parseSurfaceCommandsIncremental(cmd, 22)
	if valid {
		t.Fatal("absurd bitmapDataLength should invalidate the buffer")
	}
}

// TestFastPathFragmentConstants 锁定分片字段的线路值：MS-RDPBCGR
// 2.2.9.1.1.3.1 / FreeRDP fastpath.h 枚举（SINGLE=0, LAST=1, FIRST=2,
// NEXT=3）左移 4 位。此前 FIRST/NEXT/LAST 三个值轮换错位——首片被当
// 孤儿丢弃、续片提前冲刷、尾片被缓存到下一条更新，是 24bpp 花屏与
// surface cmd resync 的总根源。
func TestFastPathFragmentConstants(t *testing.T) {
	if FASTPATH_FRAGMENT_SINGLE != 0x00 {
		t.Fatalf("SINGLE = %#x, want 0x00", FASTPATH_FRAGMENT_SINGLE)
	}
	if FASTPATH_FRAGMENT_LAST != 0x10 {
		t.Fatalf("LAST = %#x, want 0x10", FASTPATH_FRAGMENT_LAST)
	}
	if FASTPATH_FRAGMENT_FIRST != 0x20 {
		t.Fatalf("FIRST = %#x, want 0x20", FASTPATH_FRAGMENT_FIRST)
	}
	if FASTPATH_FRAGMENT_NEXT != 0x30 {
		t.Fatalf("NEXT = %#x, want 0x30", FASTPATH_FRAGMENT_NEXT)
	}
}

// TestFastPathPointerNewLayout 用 Win10 实测指针字节锁定 TS_POINTER_NEW
// 布局：xorBpp u16 首字段 + 全 u16 颜色指针属性（FreeRDP
// update_read_pointer_new → s_update_read_pointer_color）。此前漏读
// xorBpp 导致整体后移 2 字节，指针宽高/掩码长度全错位、光标无法显示。
func TestFastPathPointerNewLayout(t *testing.T) {
	raw := []byte{
		0x20, 0x00, // xorBpp = 32
		0x00, 0x00, // cacheIndex = 0
		0x03, 0x00, // hotSpot.xPos = 3
		0x03, 0x00, // hotSpot.yPos = 3
		0x29, 0x00, // width = 41
		0x27, 0x00, // height = 39
		0xEA, 0x00, // lengthAndMask = 234 = ((41+15)/16)*2 * 39
		0xFC, 0x18, // lengthXorMask = 6396 = ((41*32+15)/16)*2 * 39
	}
	raw = append(raw, make([]byte, 234+6396)...)
	p := &FastPathUpdatePointerPDU{}
	if err := p.Unpack(bytes.NewReader(raw)); err != nil {
		t.Fatalf("Unpack: %v", err)
	}
	if p.XorBpp != 32 || p.CacheIdx != 0 || p.HotX != 3 || p.HotY != 3 ||
		p.Width != 41 || p.Height != 39 {
		t.Fatalf("header fields wrong: %+v", p)
	}
	if p.MaskLen != 234 || len(p.Mask) != 234 {
		t.Fatalf("AND mask: len=%d MaskLen=%d, want 234", len(p.Mask), p.MaskLen)
	}
	if p.XorLen != 6396 || len(p.Data) != 6396 {
		t.Fatalf("XOR mask: len=%d XorLen=%d, want 6396", len(p.Data), p.XorLen)
	}
}

// TestParseSurfaceCommandsDropsUnknownTail 验证尾部未知结构（实测 Win10
// 整屏重绘批次末尾的 11 字节尾巴）不再拖垮整批：前缀命令照常解析上屏。
func TestParseSurfaceCommandsDropsUnknownTail(t *testing.T) {
	cmd := buildSurfCmd(CMDTYPE_SET_SURFACE_BITS, 0, 32, 2, 2,
		make([]byte, 16), true)
	bad := []byte{0x80, 0x00, 0x05, 0x00, 0x01, 0x00, 0x36, 0x01, 0x00, 0x00, 0x00}
	buf := append(append([]byte{}, cmd...), bad...)
	result, consumed, needMore, valid, dropped, _ := parseSurfaceCommandsIncremental(buf, 22)
	if !valid {
		t.Fatal("合法前缀不应判无效")
	}
	if !dropped {
		t.Fatal("未知尾巴应标记 dropped")
	}
	if needMore {
		t.Fatal("丢弃尾巴后不应再等待更多数据")
	}
	if consumed != len(cmd) {
		t.Fatalf("consumed=%d，期望 %d", consumed, len(cmd))
	}
	if len(result.Rects) != 1 {
		t.Fatalf("前缀命令应产出 1 个矩形，实得 %d", len(result.Rects))
	}
}

// TestPersistentKeyListPDU：键按单元容量分配、头 20 字节布局、
// bBitMask=FIRST|LAST、键条目 uint64 小端，超 2042 键在发送层截断
// （此处仅验证 PDU 自身，截断在 maybeSendPersistentKeyList）。
func TestPersistentKeyListPDU(t *testing.T) {
	// 2000 键、单元容量 600/1024/4096/0/0 → 600+1024+376=2000
	keys := make([]uint64, 2000)
	for i := range keys {
		keys[i] = uint64(i) + 0x10000
	}
	cells := [5]uint16{600, 1024, 4096, 0, 0}
	m := NewPersistentKeyListPDU(0x103EA, keys, cells)
	if m.numEntries != [5]uint16{600, 1024, 376, 0, 0} {
		t.Fatalf("numEntries=%v", m.numEntries)
	}
	b := m.Serialize()
	// ShareDataHeader(12) + 头 24（num/total×10+bBitMask+pad）+ 2000×8
	if len(b) != 12+24+2000*8 {
		t.Fatalf("len=%d", len(b))
	}
	if b[12+20] != 0x03 || b[12+21] != 0 || b[12+22] != 0 {
		t.Fatalf("bBitMask/pad wrong: %v", b[12+20:12+23])
	}
	if got := binary.LittleEndian.Uint16(b[12:]); got != 600 {
		t.Fatalf("numEntriesCache0=%d", got)
	}
	if got := binary.LittleEndian.Uint16(b[12+10:]); got != 600 {
		t.Fatalf("totalEntriesCache0=%d", got)
	}
	if got := binary.LittleEndian.Uint16(b[12+4:]); got != 376 {
		t.Fatalf("numEntriesCache2=%d", got)
	}
	// 第一条键
	if got := binary.LittleEndian.Uint64(b[12+24:]); got != 0x10000 {
		t.Fatalf("first key=%x", got)
	}
	// 空键列表也合法（totalEntries 全 0）
	empty := NewPersistentKeyListPDU(0x103EA, nil, cells).Serialize()
	if len(empty) != 12+24 {
		t.Fatalf("empty len=%d", len(empty))
	}
}