package transport

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

const (
	tpktVersion    = 3
	tpktHeaderSize = 4
	maxTPKTSize    = 64 * 1024
)

// WriteTPKT 写入 TPKT 包头与负载，RDP 初始握手基于 TPKT 承载 // Logan.Li
func WriteTPKT(w io.Writer, payload []byte) error {
	total := len(payload) + tpktHeaderSize
	if total > maxTPKTSize {
		return fmt.Errorf("TPKT payload too large: %d", len(payload))
	}
	header := []byte{tpktVersion, 0, 0, 0}
	binary.BigEndian.PutUint16(header[2:4], uint16(total))
	if _, err := w.Write(header); err != nil {
		return err
	}
	_, err := w.Write(payload)
	return err
}

// ReadTPKT 读取一个完整 TPKT 包并返回负载 // Logan.Li
func ReadTPKT(r io.Reader) ([]byte, error) {
	header := make([]byte, tpktHeaderSize)
	if _, err := io.ReadFull(r, header); err != nil {
		return nil, err
	}
	if header[0] != tpktVersion {
		return nil, fmt.Errorf("invalid TPKT version: %d", header[0])
	}
	total := int(binary.BigEndian.Uint16(header[2:4]))
	if total < tpktHeaderSize || total > maxTPKTSize {
		return nil, fmt.Errorf("invalid TPKT length: %d", total)
	}
	payload := make([]byte, total-tpktHeaderSize)
	if _, err := io.ReadFull(r, payload); err != nil {
		return nil, err
	}
	return payload, nil
}

// WriteRaw 直接写入已封装的 TPKT/FastPath 帧，供会话层发送输入事件 // Logan.Li
func WriteRaw(w io.Writer, frame []byte) error {
	_, err := w.Write(frame)
	return err
}