// Package rdpdr 实现 RDP 设备重定向静态虚拟通道的最小握手。// Logan.Li
package rdpdr

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"log"
	"sync"
	"unicode/utf16"
)

type Sender interface {
	SendChannelPDU(channelID uint16, data []byte) error
}

const ChannelName = "rdpdr"

const (
	componentRDPDR = 0x4472

	pakIDServerAnnounce     = 0x496e
	pakIDClientIDConfirm    = 0x4343
	pakIDClientName         = 0x434e
	pakIDServerCoreCap      = 0x5350
	pakIDClientCoreCap      = 0x4350
	pakIDDeviceListAnnounce = 0x4441
	pakIDDeviceReply        = 0x6472
	pakIDDeviceIORequest    = 0x4952
	pakIDDeviceIOCompletion = 0x4943
	pakIDUserLoggedOn       = 0x554c

	capGeneralType = 0x0001
	capDriveType   = 0x0004

	deviceTypeDisk = 0x00000008
)

type Download struct {
	Name string
	Data []byte
}

type Handler struct {
	sender             Sender
	channelID          uint16
	clientID           uint32
	gotServerCap       bool
	gotIDConfirm       bool
	serverIOCode1      uint32
	computerName       string
	driveName          string
	devicePreannounced bool
	deviceAnnounced    bool
	mu                 sync.Mutex
	nextFileID         uint32
	files              map[uint32]*openFile
	uploads            map[string][]byte
	onDownload         func(Download)
}

func NewHandler(computerName string, driveName string) *Handler {
	if computerName == "" {
		computerName = "goclient"
	}
	if driveName == "" {
		driveName = "AIOPS"
	}
	return &Handler{computerName: computerName, driveName: driveName, nextFileID: 1, files: make(map[uint32]*openFile), uploads: make(map[string][]byte)}
}

func (h *Handler) SetDownloadHandler(fn func(Download)) {
	h.onDownload = fn
}

func (h *Handler) AddUpload(name string, data []byte) {
	if name == "" || len(data) == 0 {
		return
	}
	h.mu.Lock()
	h.uploads[name] = append([]byte(nil), data...)
	h.mu.Unlock()
	log.Printf("[rdpdr] upload staged name=%s size=%d", name, len(data))
}

func (h *Handler) listUploads() []dirEntry {
	h.mu.Lock()
	defer h.mu.Unlock()
	out := make([]dirEntry, 0, len(h.uploads))
	for name, data := range h.uploads {
		out = append(out, dirEntry{name: name, size: uint64(len(data))})
	}
	return out
}

func (h *Handler) getUpload(name string) ([]byte, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	data, ok := h.uploads[name]
	if !ok {
		return nil, false
	}
	return append([]byte(nil), data...), true
}

func (h *Handler) Attach(sender Sender, channelID uint16) {
	h.sender = sender
	h.channelID = channelID
	log.Printf("[rdpdr] attached channelID=%d", channelID)
}

func (h *Handler) PreannounceDevices() {
	h.mu.Lock()
	if h.devicePreannounced || h.sender == nil || h.channelID == 0 {
		h.mu.Unlock()
		return
	}
	h.devicePreannounced = true
	h.mu.Unlock()
	// 部分 Windows 会在图形首帧后才启动 rdpdr；预公告客户端能力和磁盘设备，尽早完成驱动器重定向。// Logan.Li
	log.Printf("[rdpdr] preannounce devices client=%s drive=%s", h.computerName, h.driveName)
	h.sendClientName()
	h.sendCoreCapResponse()
	h.sendDeviceList(true)
}

func (h *Handler) Process(data []byte) {
	if len(data) < 4 {
		return
	}
	component := binary.LittleEndian.Uint16(data[0:2])
	packetID := binary.LittleEndian.Uint16(data[2:4])
	body := data[4:]
	if component != componentRDPDR {
		log.Printf("[rdpdr] ignore component=0x%x packet=0x%x len=%d", component, packetID, len(data))
		return
	}
	log.Printf("[rdpdr] recv packet=0x%x len=%d", packetID, len(body))
	switch packetID {
	case pakIDServerAnnounce:
		h.handleServerAnnounce(body)
	case pakIDServerCoreCap:
		h.handleServerCoreCap(body)
	case pakIDClientIDConfirm:
		h.handleClientIDConfirm(body)
	case pakIDDeviceReply:
		h.handleDeviceReply(body)
	case pakIDDeviceIORequest:
		h.handleIORequest(body)
	case pakIDUserLoggedOn:
		h.sendDeviceListOnce()
	default:
		log.Printf("[rdpdr] unhandled packet=0x%x len=%d", packetID, len(body))
	}
}

func (h *Handler) handleServerAnnounce(body []byte) {
	if len(body) < 8 {
		return
	}
	h.clientID = binary.LittleEndian.Uint32(body[4:8])
	log.Printf("[rdpdr] server announce clientID=%d version=%d.%d", h.clientID, binary.LittleEndian.Uint16(body[0:2]), binary.LittleEndian.Uint16(body[2:4]))

	buf := make([]byte, 12)
	binary.LittleEndian.PutUint16(buf[0:2], componentRDPDR)
	binary.LittleEndian.PutUint16(buf[2:4], pakIDClientIDConfirm)
	binary.LittleEndian.PutUint16(buf[4:6], 1)
	binary.LittleEndian.PutUint16(buf[6:8], 0x000c)
	binary.LittleEndian.PutUint32(buf[8:12], h.clientID)
	h.send(buf)
	h.sendClientName()
}

func (h *Handler) sendClientName() {
	runes := utf16.Encode([]rune(h.computerName + "\x00"))
	name := make([]byte, len(runes)*2)
	for i, r := range runes {
		binary.LittleEndian.PutUint16(name[i*2:], r)
	}
	buf := make([]byte, 16+len(name))
	binary.LittleEndian.PutUint16(buf[0:2], componentRDPDR)
	binary.LittleEndian.PutUint16(buf[2:4], pakIDClientName)
	binary.LittleEndian.PutUint32(buf[4:8], 1)
	binary.LittleEndian.PutUint32(buf[12:16], uint32(len(name)))
	copy(buf[16:], name)
	h.send(buf)
}

func (h *Handler) handleServerCoreCap(body []byte) {
	if len(body) < 4 {
		return
	}
	count := binary.LittleEndian.Uint16(body[0:2])
	offset := 4
	for i := 0; i < int(count) && offset+8 <= len(body); i++ {
		capType := binary.LittleEndian.Uint16(body[offset:])
		capLen := int(binary.LittleEndian.Uint16(body[offset+2:]))
		if capLen < 8 || offset+capLen > len(body) {
			break
		}
		if capType == capGeneralType && capLen >= 24 {
			h.serverIOCode1 = binary.LittleEndian.Uint32(body[offset+20 : offset+24])
		}
		offset += capLen
	}
	h.gotServerCap = true
	h.trySendReady()
}

func (h *Handler) handleClientIDConfirm(body []byte) {
	if len(body) >= 8 {
		h.clientID = binary.LittleEndian.Uint32(body[4:8])
	}
	h.gotIDConfirm = true
	h.trySendReady()
}

func (h *Handler) trySendReady() {
	if !h.gotServerCap || !h.gotIDConfirm {
		return
	}
	h.sendCoreCapResponse()
	// 能力协商完成后立即公告虚拟盘，避免首次登录后资源管理器等待网络位置刷新。// Logan.Li
	h.sendDeviceListOnce()
}

func (h *Handler) sendCoreCapResponse() {
	const generalCapLen = 44
	const driveCapLen = 8
	buf := bytes.Buffer{}
	writeLE16(&buf, componentRDPDR)
	writeLE16(&buf, pakIDClientCoreCap)
	writeLE16(&buf, 2)
	writeLE16(&buf, 0)

	writeLE16(&buf, capGeneralType)
	writeLE16(&buf, generalCapLen)
	writeLE32(&buf, 2)
	writeLE32(&buf, 0)
	writeLE32(&buf, 0)
	writeLE16(&buf, 1)
	writeLE16(&buf, 0x000c)
	ioCode1 := h.serverIOCode1 & 0x0000ffff
	if ioCode1 == 0 {
		ioCode1 = 0x0000ffff
	}
	writeLE32(&buf, ioCode1)
	writeLE32(&buf, 0)
	writeLE32(&buf, 0x00000007)
	writeLE32(&buf, 0x00000001)
	writeLE32(&buf, 0)
	writeLE32(&buf, 0)

	writeLE16(&buf, capDriveType)
	writeLE16(&buf, driveCapLen)
	writeLE32(&buf, 2)
	h.send(buf.Bytes())
}

func (h *Handler) sendDeviceList(includeDevices bool) {
	deviceData := []byte(nil)
	deviceCount := uint32(0)
	totalLen := 8
	if includeDevices {
		deviceCount = 1
		deviceData = []byte(h.driveName)
		totalLen += 20 + len(deviceData)
	}
	if deviceCount == 0 {
		// FreeRDP 在没有预登录设备时直接跳过空公告，避免 Windows 记录一个无设备列表。// Logan.Li
		log.Printf("[rdpdr] skip empty device list announce")
		return
	}
	buf := make([]byte, totalLen)
	binary.LittleEndian.PutUint16(buf[0:2], componentRDPDR)
	binary.LittleEndian.PutUint16(buf[2:4], pakIDDeviceListAnnounce)
	binary.LittleEndian.PutUint32(buf[4:8], deviceCount)
	if includeDevices {
		// 公告一个磁盘设备，Windows 里显示为 \\goclient\AIOPS。// Logan.Li
		binary.LittleEndian.PutUint32(buf[8:12], deviceTypeDisk)
		binary.LittleEndian.PutUint32(buf[12:16], 1)
		copy(buf[16:24], []byte(dosDeviceName(h.driveName)))
		binary.LittleEndian.PutUint32(buf[24:28], uint32(len(deviceData)))
		copy(buf[28:], deviceData)
	}
	h.send(buf)
	log.Printf("[rdpdr] device list announced devices=%d name=%s", deviceCount, h.driveName)
}

func (h *Handler) sendDeviceListOnce() {
	h.mu.Lock()
	if h.deviceAnnounced {
		h.mu.Unlock()
		return
	}
	h.deviceAnnounced = true
	h.mu.Unlock()
	h.sendDeviceList(true)
}

func dosDeviceName(name string) string {
	if name == "" {
		return "AIOPS"
	}
	if len(name) > 8 {
		return name[:8]
	}
	return name
}

func (h *Handler) send(data []byte) {
	if h.sender == nil || h.channelID == 0 {
		return
	}
	if err := h.sender.SendChannelPDU(h.channelID, data); err != nil {
		log.Printf("[rdpdr] send failed: %v", err)
	}
}

func (h *Handler) sendIOCompletion(deviceID, completionID, status uint32, output []byte) {
	buf := make([]byte, 16+len(output))
	binary.LittleEndian.PutUint16(buf[0:2], componentRDPDR)
	binary.LittleEndian.PutUint16(buf[2:4], pakIDDeviceIOCompletion)
	binary.LittleEndian.PutUint32(buf[4:8], deviceID)
	binary.LittleEndian.PutUint32(buf[8:12], completionID)
	binary.LittleEndian.PutUint32(buf[12:16], status)
	copy(buf[16:], output)
	h.send(buf)
}

func (h *Handler) allocFileID() uint32 {
	h.mu.Lock()
	defer h.mu.Unlock()
	id := h.nextFileID
	h.nextFileID++
	return id
}

func (h *Handler) storeOpenFile(file *openFile) uint32 {
	id := h.allocFileID()
	h.mu.Lock()
	h.files[id] = file
	h.mu.Unlock()
	return id
}

func (h *Handler) getOpenFile(id uint32) (*openFile, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	file, ok := h.files[id]
	return file, ok
}

func (h *Handler) removeOpenFile(id uint32) (*openFile, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	file, ok := h.files[id]
	if ok {
		delete(h.files, id)
	}
	return file, ok
}

func (h *Handler) handleDeviceReply(body []byte) {
	if len(body) < 8 {
		return
	}
	deviceID := binary.LittleEndian.Uint32(body[0:4])
	status := binary.LittleEndian.Uint32(body[4:8])
	if status == statusSuccess {
		h.mu.Lock()
		h.deviceAnnounced = true
		h.mu.Unlock()
	}
	log.Printf("[rdpdr] device reply deviceID=%d status=0x%08x", deviceID, status)
}

func writeLE16(w *bytes.Buffer, v uint16) {
	_ = binary.Write(w, binary.LittleEndian, v)
}

func writeLE32(w *bytes.Buffer, v uint32) {
	_ = binary.Write(w, binary.LittleEndian, v)
}

func (h *Handler) String() string {
	return fmt.Sprintf("rdpdr(channel=%d)", h.channelID)
}