package transport

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

const (
	x224TypeConnectionRequest = 0xE0
	x224TypeConnectionConfirm = 0xD0
	x224TypeData              = 0xF0

	negTypeRequest  = 0x01
	negTypeResponse = 0x02
	negTypeFailure  = 0x03

	protocolSSL     = 0x00000001
	protocolCredSSP = 0x00000002
)

const (
	ProtocolStandard uint32 = 0
	ProtocolTLS      uint32 = protocolSSL
	ProtocolCredSSP  uint32 = protocolCredSSP
)

// NegotiationResult 是 X.224 阶段的 RDP 安全协议协商结果 // Logan.Li
type NegotiationResult struct {
	SelectedProtocol uint32
	FailureCode      uint32
}

func (r NegotiationResult) ProtocolName() string {
	switch r.SelectedProtocol {
	case ProtocolCredSSP:
		return "CredSSP/NLA"
	case ProtocolTLS:
		return "TLS"
	case ProtocolStandard:
		return "Standard RDP"
	default:
		return fmt.Sprintf("unknown(0x%x)", r.SelectedProtocol)
	}
}

// BuildConnectionRequest 构造 X.224 Connection Request，并请求 NLA/TLS 协商 // Logan.Li
func BuildConnectionRequest() []byte {
	return BuildConnectionRequestWithCookie("aiops")
}

// BuildConnectionRequestWithCookie 按 MS-RDPBCGR 2.2.1.1.1 构造包含 mstshash 的请求 // Logan.Li
func BuildConnectionRequestWithCookie(username string) []byte {
	if username == "" {
		username = "aiops"
	}
	if len(username) > 9 {
		username = username[:9]
	}
	cookie := []byte("Cookie: mstshash=" + username + "\r\n")
	neg := make([]byte, 8)
	neg[0] = negTypeRequest
	neg[1] = 0
	binary.LittleEndian.PutUint16(neg[2:4], uint16(len(neg)))
	binary.LittleEndian.PutUint32(neg[4:8], protocolSSL|protocolCredSSP)

	length := 7 + len(cookie) + len(neg)
	pdu := make([]byte, 0, length)
	pdu = append(pdu, byte(length-1))
	pdu = append(pdu, x224TypeConnectionRequest)
	pdu = append(pdu, 0x00, 0x00) // dst-ref // Logan.Li
	pdu = append(pdu, 0x00, 0x00) // src-ref // Logan.Li
	pdu = append(pdu, 0x00)       // class/options // Logan.Li
	pdu = append(pdu, cookie...)
	pdu = append(pdu, neg...)
	return pdu
}

// ParseConnectionConfirm 解析 X.224 Connection Confirm 与 RDP Negotiation Response // Logan.Li
func ParseConnectionConfirm(payload []byte) (NegotiationResult, error) {
	if len(payload) < 7 {
		return NegotiationResult{}, fmt.Errorf("short X.224 confirm: %d", len(payload))
	}
	if payload[1] != x224TypeConnectionConfirm {
		return NegotiationResult{}, fmt.Errorf("unexpected X.224 pdu type: 0x%x", payload[1])
	}
	for i := 7; i+8 <= len(payload); i++ {
		switch payload[i] {
		case negTypeResponse:
			return NegotiationResult{SelectedProtocol: binary.LittleEndian.Uint32(payload[i+4 : i+8])}, nil
		case negTypeFailure:
			return NegotiationResult{FailureCode: binary.LittleEndian.Uint32(payload[i+4 : i+8])}, fmt.Errorf("RDP negotiation failure: 0x%x", binary.LittleEndian.Uint32(payload[i+4:i+8]))
		}
	}
	return NegotiationResult{SelectedProtocol: 0}, nil
}

// WriteX224Data 写入 X.224 Data TPDU，NLA 后的 MCS/GCC 数据仍经由该层承载 // Logan.Li
func WriteX224Data(w io.Writer, payload []byte) error {
	data := make([]byte, 0, len(payload)+3)
	data = append(data, 0x02, x224TypeData, 0x80)
	data = append(data, payload...)
	return WriteTPKT(w, data)
}

// ReadX224Data 读取并剥离 X.224 Data TPDU 头，返回上层 MCS 负载 // Logan.Li
func ReadX224Data(r io.Reader) ([]byte, error) {
	payload, err := ReadTPKT(r)
	if err != nil {
		return nil, err
	}
	if len(payload) < 3 {
		return nil, fmt.Errorf("X.224 data too short: %d", len(payload))
	}
	if payload[1] != x224TypeData {
		return nil, fmt.Errorf("unexpected X.224 type: 0x%x", payload[1])
	}
	return payload[3:], nil
}