package rdp

import (
	"encoding/binary"
	"fmt"
	"io"

	"ezops/pkg/rdp/transport"
)

const (
	fastPathActionFastPath = 0x0
	fastPathActionX224     = 0x3
)

// WirePacket 表示从 RDP 连接读取到的一层帧 // Logan.Li
type WirePacket struct {
	FastPath bool
	SecFlag  byte
	Payload  []byte
}

// ReadWirePacket 读取 TPKT/X.224 或 FastPath 输出帧 // Logan.Li
func ReadWirePacket(r io.Reader) (WirePacket, error) {
	var hdr [2]byte
	if _, err := io.ReadFull(r, hdr[:]); err != nil {
		return WirePacket{}, err
	}
	version := hdr[0]
	if version == fastPathActionX224 {
		var extHdr [2]byte
		if _, err := io.ReadFull(r, extHdr[:]); err != nil {
			return WirePacket{}, err
		}
		size := binary.BigEndian.Uint16(extHdr[:])
		if size < 4 {
			return WirePacket{}, fmt.Errorf("TPKT 长度非法: %d", size)
		}
		body := make([]byte, int(size)-4)
		if _, err := io.ReadFull(r, body); err != nil {
			return WirePacket{}, err
		}
		if len(body) < 3 || body[1] != 0xF0 {
			return WirePacket{}, fmt.Errorf("非 X.224 Data PDU")
		}
		return WirePacket{Payload: body[3:]}, nil
	}
	secFlag := (version >> 6) & 0x3
	length := int(hdr[1])
	packetSize := 0
	if length&0x80 != 0 {
		var extByte [1]byte
		if _, err := io.ReadFull(r, extByte[:]); err != nil {
			return WirePacket{}, err
		}
		packetSize = ((length &^ 0x80) << 8) + int(extByte[0]) - 3
	} else {
		packetSize = length - 2
	}
	if packetSize < 0 {
		return WirePacket{FastPath: true, SecFlag: secFlag}, nil
	}
	body := make([]byte, packetSize)
	if _, err := io.ReadFull(r, body); err != nil {
		return WirePacket{}, err
	}
	return WirePacket{FastPath: true, SecFlag: secFlag, Payload: body}, nil
}

// WriteFastPathInput 发送 Fast-Path Input 事件批次 // Logan.Li
func WriteFastPathInput(w io.Writer, events []byte) error {
	hdr := uint16(len(events)+3) | 0x8000
	frame := make([]byte, 0, len(events)+3)
	frame = append(frame, fastPathActionFastPath|0x04, byte(hdr>>8), byte(hdr))
	frame = append(frame, events...)
	return transport.WriteRaw(w, frame)
}