package rdpdr

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

// fakeSender 捕获发往通道的消息。
type fakeSender struct{ msgs [][]byte }

func (f *fakeSender) SendToChannel(ch string, s []byte) (int, error) {
	f.msgs = append(f.msgs, append([]byte(nil), s...))
	return len(s), nil
}

// fakeFS 记录派发并支持手动完成。
type fakeFS struct {
	h        *Handler
	opens    []string
	lists    int
	complet  []uint32
}

func (f *fakeFS) Open(completionID, fileID uint32, path string) {
	f.opens = append(f.opens, path)
}
func (f *fakeFS) Read(completionID, fileID uint32, offset uint64, length uint32) {}
func (f *fakeFS) Close(completionID, fileID uint32)                              {}
func (f *fakeFS) List(completionID, fileID uint32, initial bool) {
	f.complet = append(f.complet, completionID)
	f.lists++
}
func (f *fakeFS) Stat(completionID, fileID uint32)                             {}
func (f *fakeFS) Volume(completionID uint32, infoClass uint32)                 {}

func newTestHandler() (*Handler, *fakeSender, *fakeFS) {
	h := NewHandler("webrdp", "local")
	s := &fakeSender{}
	h.Sender(s)
	fs := &fakeFS{h: h}
	h.SetFilesystem(fs)
	return h, s, fs
}

func TestServerAnnounceFlow(t *testing.T) {
	h, s, _ := newTestHandler()
	// Server Announce：major 1 minor 5 clientID 0x1234
	req := make([]byte, 12)
	binary.LittleEndian.PutUint16(req[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(req[2:], PAKID_CORE_SERVER_ANNOUNCE)
	binary.LittleEndian.PutUint16(req[4:], 1)
	binary.LittleEndian.PutUint16(req[6:], 5)
	binary.LittleEndian.PutUint32(req[8:], 0x1234)
	h.Process(req)

	if len(s.msgs) != 2 {
		t.Fatalf("应发 2 条消息（announce reply/name），实发 %d", len(s.msgs))
	}
	// 1) Client Announce Reply：v1.5 + 原 ClientId
	if got := binary.LittleEndian.Uint16(s.msgs[0][2:]); got != PAKID_CORE_CLIENTID_CONFIRM {
		t.Fatalf("msg0 packetID=%#x", got)
	}
	if binary.LittleEndian.Uint16(s.msgs[0][4:]) != 1 || binary.LittleEndian.Uint16(s.msgs[0][6:]) != 5 {
		t.Fatal("announce reply 版本应为 1.5")
	}
	if binary.LittleEndian.Uint32(s.msgs[0][8:]) != 0x1234 {
		t.Fatal("announce reply 应原样回 ClientId")
	}
	// 2) Client Name Request：UnicodeFlag=1
	if got := binary.LittleEndian.Uint16(s.msgs[1][2:]); got != PAKID_CORE_CLIENT_NAME {
		t.Fatalf("msg1 packetID=%#x", got)
	}
	if binary.LittleEndian.Uint32(s.msgs[1][4:]) != 1 {
		t.Fatal("UnicodeFlag 应为 1")
	}

	// 服务端 CLIENTID_CONFIRM：仅记录状态，设备列表还不发
	req2 := make([]byte, 8)
	binary.LittleEndian.PutUint16(req2[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(req2[2:], PAKID_CORE_CLIENTID_CONFIRM)
	h.Process(req2)
	if len(s.msgs) != 2 {
		t.Fatalf("CLIENTID_CONFIRM 后不应发消息，实发 %d", len(s.msgs))
	}
	// 服务端 USER_LOGGEDON → 此刻才宣告设备列表（MS-RDPEFS 3.2.5.1.3）
	req3 := make([]byte, 8)
	binary.LittleEndian.PutUint16(req3[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(req3[2:], PAKID_CORE_USER_LOGGEDON)
	h.Process(req3)
	if len(s.msgs) != 3 {
		t.Fatalf("USER_LOGGEDON 后应发设备列表，实发 %d", len(s.msgs))
	}
	// 3) Device List：count=1, type=FILESYSTEM, DosName "webrdp"
	if got := binary.LittleEndian.Uint16(s.msgs[2][2:]); got != PAKID_CORE_DEVICELIST_ANNOUNCE {
		t.Fatalf("msg2 packetID=%#x", got)
	}
	if binary.LittleEndian.Uint32(s.msgs[2][4:]) != 1 {
		t.Fatal("DeviceCount 应为 1")
	}
	if binary.LittleEndian.Uint32(s.msgs[2][8:]) != RDPDR_DTYP_FILESYSTEM {
		t.Fatal("DeviceType 应为 FILESYSTEM(8)")
	}
	if got := binary.LittleEndian.Uint32(s.msgs[2][12:]); got != 1 {
		t.Fatalf("DeviceId 应为 1（非 0），实 %d", got)
	}
	if name := string(bytes.TrimRight(s.msgs[2][16:24], "\x00")); name != "webrdp" {
		t.Fatalf("DosName=%q", name)
	}
	// V02 协商下 DeviceData = ASCII 全名 + 1 字节 NUL（FreeRDP drive_main.c
	// 同款；服务端 UNC 映射键取 DosName，DeviceData 内容须非空否则拒装）
	if got := len(s.msgs[2]); got != 35 { // 28 + "webrdp" ASCII 6 + NUL 1
		t.Fatalf("设备列表报文应为 35 字节，实 %d", got)
	}
	if got := binary.LittleEndian.Uint32(s.msgs[2][24:]); got != 7 {
		t.Fatalf("DeviceDataLength 应为 7，实 %d", got)
	}
	if !bytes.Equal(s.msgs[2][28:35], []byte("webrdp\x00")) {
		t.Fatalf("DeviceData 应为 ASCII 全名+NUL，实 % x", s.msgs[2][28:35])
	}
}

func TestAnnounceVersionNegotiation(t *testing.T) {
	h, s, _ := newTestHandler()
	// 服务端 1.13（Win10）→ 客户端应回 1.MIN(0x000D, 0x000D)=1.13
	req := make([]byte, 12)
	binary.LittleEndian.PutUint16(req[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(req[2:], PAKID_CORE_SERVER_ANNOUNCE)
	binary.LittleEndian.PutUint16(req[4:], 1)
	binary.LittleEndian.PutUint16(req[6:], 0x000D)
	binary.LittleEndian.PutUint32(req[8:], 0x1234)
	h.Process(req)
	if binary.LittleEndian.Uint16(s.msgs[0][4:]) != 1 ||
		binary.LittleEndian.Uint16(s.msgs[0][6:]) != 0x000D {
		t.Fatalf("服务端 1.13 时应回 1.13，实 %d.%d",
			binary.LittleEndian.Uint16(s.msgs[0][4:]),
			binary.LittleEndian.Uint16(s.msgs[0][6:]))
	}
	// 服务端 1.5（老语义）→ 回 1.5
	h2, s2, _ := newTestHandler()
	binary.LittleEndian.PutUint16(req[6:], 5)
	h2.Process(req)
	if binary.LittleEndian.Uint16(s2.msgs[0][6:]) != 5 {
		t.Fatal("服务端 1.5 时应回 1.5")
	}
}

func TestSingleAnnounceAtLogon(t *testing.T) {
	h, s, _ := newTestHandler()
	// Server Announce → Server Caps → ClientID Confirm：1.0x 语义下
	// 均不应发设备列表（FreeRDP：非登录阶段跳过 FS 设备，count=0 不发）
	ann := make([]byte, 12)
	binary.LittleEndian.PutUint16(ann[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(ann[2:], PAKID_CORE_SERVER_ANNOUNCE)
	binary.LittleEndian.PutUint16(ann[4:], 1)
	binary.LittleEndian.PutUint16(ann[6:], 0x000D)
	binary.LittleEndian.PutUint32(ann[8:], 0x1234)
	h.Process(ann)
	caps := make([]byte, 8)
	binary.LittleEndian.PutUint16(caps[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(caps[2:], PAKID_CORE_SERVER_CAPABILITY)
	caps[4] = 0
	h.Process(caps)
	conf := make([]byte, 12)
	binary.LittleEndian.PutUint16(conf[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(conf[2:], PAKID_CORE_CLIENTID_CONFIRM)
	binary.LittleEndian.PutUint16(conf[4:], 1)
	binary.LittleEndian.PutUint16(conf[6:], 0x000D)
	binary.LittleEndian.PutUint32(conf[8:], 0x99) // 服务端回显采纳
	h.Process(conf)
	// msgs: announce reply, name request, caps response——无设备列表
	if n := len(s.msgs); n != 3 {
		t.Fatalf("登录前不应发设备列表（3 条），实 %d", n)
	}
	if h.clientID != 0x99 {
		t.Fatalf("应采纳 Confirm 回显的 ClientId，实 %#x", h.clientID)
	}
	// USER_LOGGEDON：唯一一次设备列表宣告
	logon := make([]byte, 8)
	binary.LittleEndian.PutUint16(logon[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(logon[2:], PAKID_CORE_USER_LOGGEDON)
	h.Process(logon)
	if n := len(s.msgs); n != 4 {
		t.Fatalf("USER_LOGGEDON 后应发设备列表（4 条），实 %d", n)
	}
	if got := binary.LittleEndian.Uint16(s.msgs[3][2:]); got != PAKID_CORE_DEVICELIST_ANNOUNCE {
		t.Fatalf("msg3 应为设备列表，实 %#x", got)
	}
	// 重复 USER_LOGGEDON 不重发
	h.Process(logon)
	if n := len(s.msgs); n != 4 {
		t.Fatalf("重复 USER_LOGGEDON 不应重发，实 %d", n)
	}
}

func TestSanitizeShareName(t *testing.T) {
	if got := sanitizeShareName(`a:b<c>d"e/f\g|h i,j`); got != "a_b_c_d_e_f_g_h_i_j" {
		t.Fatalf("sanitizeShareName=%q", got)
	}
	if got := sanitizeShareName("rdpdrive-test"); got != "rdpdrive-test" {
		t.Fatalf("合法名不应改动，实 %q", got)
	}
}

func TestClientCapabilityResponse(t *testing.T) {
	h, s, _ := newTestHandler()
	// 先握手（服务端 1.13 → 客户端版本 1.13），再发带 GENERAL capset 的
	// 服务端能力请求（ioCode1 = 全部常见 IRP 位）
	ann := make([]byte, 12)
	binary.LittleEndian.PutUint16(ann[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(ann[2:], PAKID_CORE_SERVER_ANNOUNCE)
	binary.LittleEndian.PutUint16(ann[4:], 1)
	binary.LittleEndian.PutUint16(ann[6:], 0x000D)
	binary.LittleEndian.PutUint32(ann[8:], 0x1234)
	h.Process(ann)

	req := make([]byte, 8+44)
	binary.LittleEndian.PutUint16(req[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(req[2:], PAKID_CORE_SERVER_CAPABILITY)
	binary.LittleEndian.PutUint16(req[4:], 1) // numCapabilities
	binary.LittleEndian.PutUint16(req[8:], CAP_GENERAL_TYPE)
	binary.LittleEndian.PutUint16(req[10:], 44)
	binary.LittleEndian.PutUint32(req[12:], 2)
	binary.LittleEndian.PutUint32(req[28:], 0xFFFFFFFF) // ioCode1
	h.Process(req)

	if len(s.msgs) != 3 || len(s.msgs[2]) != 8+44+8 {
		t.Fatalf("能力响应长度应 60，实 %v", len(s.msgs[2]))
	}
	m := s.msgs[2]
	if binary.LittleEndian.Uint16(m[2:]) != PAKID_CORE_CLIENT_CAPABILITY {
		t.Fatal("packetID 应为 CLIENT_CAPABILITY")
	}
	if binary.LittleEndian.Uint16(m[4:]) != 2 {
		t.Fatal("numCapabilities 应为 2")
	}
	// GENERAL：版本字段各 2 字节（旧实现误写 4 字节，服务端解析错位）
	if binary.LittleEndian.Uint16(m[8:]) != CAP_GENERAL_TYPE || binary.LittleEndian.Uint16(m[10:]) != 44 {
		t.Fatal("GENERAL capset 头错误")
	}
	if binary.LittleEndian.Uint16(m[24:]) != 1 || binary.LittleEndian.Uint16(m[26:]) != 0x000D {
		t.Fatal("GENERAL 协议版本应为 1.13（各 2 字节）")
	}
	// ioCode1 = 客户端掩码 ∩ 服务端掩码
	if got := binary.LittleEndian.Uint32(m[28:]); got != clientIOCode1 {
		t.Fatalf("ioCode1 应为求交结果 %#x，实 %#x", clientIOCode1, got)
	}
	// extendedPDU = REMOVE|DISPLAY_NAME|USER_LOGGEDON；extraFlags1 = ENABLE_ASYNCIO
	if got := binary.LittleEndian.Uint32(m[36:]); got != 7 {
		t.Fatalf("extendedPDU 应为 7，实 %#x", got)
	}
	if got := binary.LittleEndian.Uint32(m[40:]); got != RDPDR_ENABLE_ASYNCIO {
		t.Fatalf("extraFlags1 应为 ENABLE_ASYNCIO，实 %#x", got)
	}
	// DRIVE：仅 8 字节头 {type=4, len=8, version=2}（FreeRDP 同款）
	if binary.LittleEndian.Uint16(m[52:]) != CAP_DRIVE_TYPE || binary.LittleEndian.Uint16(m[54:]) != 8 {
		t.Fatal("DRIVE capset 头错误（FreeRDP 对齐：仅 8 字节头，len=8）")
	}
	if binary.LittleEndian.Uint32(m[56:]) != 2 {
		t.Fatal("DRIVE capset Version 应为 2")
	}
}

func TestCreateDispatchAndComplete(t *testing.T) {
	h, s, fs := newTestHandler()

	// 构造 CREATE 请求：path "\hello.txt"（FileId 由 handler 分配，预期为 1）
	path := utf16Bytes("\\hello.txt")
	req := make([]byte, 56+len(path))
	binary.LittleEndian.PutUint16(req[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(req[2:], PAKID_CORE_DEVICE_IOREQUEST)
	binary.LittleEndian.PutUint32(req[4:], h.deviceID)
	binary.LittleEndian.PutUint32(req[12:], 77) // CompletionId
	binary.LittleEndian.PutUint32(req[16:], IRP_MJ_CREATE)
	binary.LittleEndian.PutUint32(req[52:], uint32(len(path)))
	copy(req[56:], path)
	h.Process(req)

	if len(fs.opens) != 1 || fs.opens[0] != "\\hello.txt" {
		t.Fatalf("Open 未正确派发: %v", fs.opens)
	}

	// 完成路径 A：失败 → 纯状态响应
	h.CompleteJSON(77, STATUS_NO_SUCH_FILE, "create", "file")
	if len(s.msgs) != 1 {
		t.Fatalf("失败完成应发 1 条响应，实 %d", len(s.msgs))
	}
	resp := s.msgs[0]
	if binary.LittleEndian.Uint16(resp[2:]) != PAKID_CORE_DEVICE_IOCOMPLETION ||
		binary.LittleEndian.Uint32(resp[8:]) != 77 ||
		binary.LittleEndian.Uint32(resp[12:]) != STATUS_NO_SUCH_FILE {
		t.Fatal("失败 CREATE 响应头错误")
	}

	// 完成路径 B：成功 → 16 字节头 + FileId(4) + Information(1)
	h.track(78, &pending{deviceID: h.deviceID, fileID: 1, major: IRP_MJ_CREATE})
	h.CompleteJSON(78, STATUS_SUCCESS, "create", "file")
	resp = s.msgs[1]
	if len(resp) != 21 {
		t.Fatalf("成功 CREATE 响应应 21 字节，实 %d", len(resp))
	}
	if binary.LittleEndian.Uint32(resp[16:]) != 1 || resp[20] != 1 {
		t.Fatal("CREATE 响应 FileId/Information 错误")
	}
}

func TestDirectoryListPagination(t *testing.T) {
	h, s, fs := newTestHandler()
	dirID := uint32(1)

	req := make([]byte, 56)
	binary.LittleEndian.PutUint16(req[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(req[2:], PAKID_CORE_DEVICE_IOREQUEST)
	binary.LittleEndian.PutUint32(req[4:], h.deviceID)
	binary.LittleEndian.PutUint32(req[8:], dirID)
	binary.LittleEndian.PutUint32(req[12:], 500)
	binary.LittleEndian.PutUint32(req[16:], IRP_MJ_DIRECTORY_CONTROL)
	binary.LittleEndian.PutUint32(req[20:], IRP_MN_QUERY_DIRECTORY)
	binary.LittleEndian.PutUint32(req[24:], FileBothDirectoryInformation)
	req[28] = 1 // InitialQuery
	h.Process(req)

	if len(fs.complet) != 1 {
		t.Fatal("List 未派发")
	}
	// 桥接返回 2 条目
	h.CompleteJSON(fs.complet[0], STATUS_SUCCESS, "list",
		`[{"name":"a.txt","size":5,"mtime":1700000000000},{"name":"sub","dir":true,"mtime":1700000000000}]`)
	resp := s.msgs[0]
	if binary.LittleEndian.Uint32(resp[12:]) != STATUS_SUCCESS {
		t.Fatal("首次枚举应成功")
	}
	dataLen := binary.LittleEndian.Uint32(resp[16:])
	data := resp[20:]
	if int(dataLen) != len(data)-0 || dataLen == 0 {
		t.Fatalf("数据长度字段不一致: %d", dataLen)
	}
	// 条目链校验：第一条 NextEntryOffset 指向第二条，第二条为 0
	first := binary.LittleEndian.Uint32(data[0:])
	if first == 0 || int(first)+4 > len(data) {
		t.Fatalf("NextEntryOffset 链错误: %d", first)
	}
	// a.txt 名字长度（UTF-16 字节数）在 class3: FileNameLength @60
	nameLen := binary.LittleEndian.Uint32(data[60:])
	if nameLen != 10 { // "a.txt" 5 chars × 2
		t.Fatalf("FileNameLength=%d", nameLen)
	}
	second := binary.LittleEndian.Uint32(data[first:])
	if second != 0 {
		t.Fatalf("末条 NextEntryOffset 应为 0，实 %d", second)
	}

	// 续枚举（InitialQuery=0）：桥接再次返回全量，游标已到末尾 → NO_MORE_FILES
	req[28] = 0
	h.Process(req)
	if len(fs.complet) != 2 {
		t.Fatal("续枚举 List 未派发")
	}
	h.CompleteJSON(fs.complet[1], STATUS_SUCCESS, "list",
		`[{"name":"a.txt","size":5,"mtime":1700000000000},{"name":"sub","dir":true,"mtime":1700000000000}]`)
	resp = s.msgs[1]
	if binary.LittleEndian.Uint32(resp[12:]) != STATUS_NO_MORE_FILES {
		t.Fatalf("续枚举应返回 NO_MORE_FILES，实 %#x", binary.LittleEndian.Uint32(resp[12:]))
	}
}

func TestReadCompletionEncoding(t *testing.T) {
	h, s, _ := newTestHandler()
	h.track(9, &pending{deviceID: h.deviceID, fileID: 3, major: IRP_MJ_READ})
	h.CompleteBytes(9, STATUS_SUCCESS, []byte("hello"))
	resp := s.msgs[0]
	if len(resp) != 16+4+5 {
		t.Fatalf("READ 响应应 25 字节，实 %d", len(resp))
	}
	if binary.LittleEndian.Uint32(resp[16:]) != 5 {
		t.Fatal("Length 前缀应为 5")
	}
	if !bytes.Equal(resp[20:], []byte("hello")) {
		t.Fatal("数据内容不一致")
	}
}

func TestWriteDeniedInReadOnlyMode(t *testing.T) {
	h, s, _ := newTestHandler()
	req := make([]byte, 24)
	binary.LittleEndian.PutUint16(req[0:], RDPDR_CTYP_CORE)
	binary.LittleEndian.PutUint16(req[2:], PAKID_CORE_DEVICE_IOREQUEST)
	binary.LittleEndian.PutUint32(req[4:], h.deviceID)
	binary.LittleEndian.PutUint32(req[12:], 42)
	binary.LittleEndian.PutUint32(req[16:], IRP_MJ_WRITE)
	h.Process(req)
	if len(s.msgs) != 1 {
		t.Fatal("WRITE 应直接拒绝")
	}
	if binary.LittleEndian.Uint32(s.msgs[0][12:]) != STATUS_ACCESS_DENIED {
		t.Fatal("WRITE 应回 ACCESS_DENIED")
	}
}

func TestDirEntryEncodingLayouts(t *testing.T) {
	entries := []DirEntry{{Name: "x", Dir: true, Size: 0, Mtime: 1700000000000}}
	// 单条即末条：NextEntryOffset=0，无对齐填充
	if d := encodeDirEntries(FileDirectoryInformation, entries); len(d) != 64+2 {
		t.Fatalf("class1 单条应为 66，实 %d", len(d))
	}
	if d := encodeDirEntries(FileBothDirectoryInformation, entries); len(d) != 93+2 {
		t.Fatalf("class3 单条应为 95，实 %d", len(d))
	}
	if d := encodeDirEntries(FileNamesInformation, entries); len(d) != 12+2 {
		t.Fatalf("names 单条应为 14，实 %d", len(d))
	}
	if d := encodeDirEntries(0x99, entries); d != nil {
		t.Fatal("未知信息类应返回 nil")
	}
	// 非末条需 4 字节对齐：两条 class1（"x"=2 字节名）→ 68 + 66
	two := []DirEntry{{Name: "x", Mtime: 1}, {Name: "y", Mtime: 1}}
	if d := encodeDirEntries(FileDirectoryInformation, two); len(d) != 68+66 {
		t.Fatalf("class1 两条应 68+66，实 %d", len(d))
	}
	// FileBasicInformation 40 字节
	if d := encodeFileInfo(FileBasicInformation, entries[0]); len(d) != 40 {
		t.Fatalf("FileBasicInformation 应 40 字节，实 %d", len(d))
	}
	// FileFsSizeInformation 24 字节
	if d := encodeVolumeInfo(FileFsSizeInformation, "local"); len(d) != 24 {
		t.Fatalf("FileFsSizeInformation 应 24 字节，实 %d", len(d))
	}
}