package protocol

import (
	"encoding/binary"
	"errors"
	"fmt"
	"net/netip"
)

const (
	VPNCmdOpen  byte = 0x30
	VPNCmdPush  byte = 0x31
	VPNCmdPull  byte = 0x32
	VPNCmdClose byte = 0x33

	VPNRespOpened byte = 0x40
	VPNRespAck    byte = 0x41
	VPNRespData   byte = 0x42
	VPNRespWait   byte = 0x43
	VPNRespClosed byte = 0x44
	VPNRespError  byte = 0x7f

	VPNNoAck       uint32 = 0xffffffff
	VPNMaxFragment        = 65535
)

type VPNSessionID [16]byte

func VPNError(message string) []byte {
	b := []byte(message)
	if len(b) > 4096 {
		b = b[:4096]
	}
	out := make([]byte, 1+len(b))
	out[0] = VPNRespError
	copy(out[1:], b)
	return out
}

func ParseVPNError(payload []byte) error {
	if len(payload) == 0 {
		return errors.New("empty DragonTCP VPN response")
	}
	if payload[0] == VPNRespError {
		return errors.New(string(payload[1:]))
	}
	return nil
}

// OPEN request:
// cmd(1) sid(16) tokenLen(2) token(N) ipv4(4) ipv6(16) mtu(2)
func BuildVPNOpen(sid VPNSessionID, token string, ipv4, ipv6 netip.Addr, mtu int) ([]byte, error) {
	if len(token) > 4096 {
		return nil, errors.New("token too long")
	}
	if !ipv4.Is4() || !ipv6.Is6() {
		return nil, errors.New("invalid VPN client addresses")
	}
	if mtu < 576 || mtu > 65535 {
		return nil, errors.New("invalid VPN MTU")
	}
	out := make([]byte, 1+16+2+len(token)+4+16+2)
	out[0] = VPNCmdOpen
	copy(out[1:17], sid[:])
	binary.BigEndian.PutUint16(out[17:19], uint16(len(token)))
	pos := 19
	copy(out[pos:pos+len(token)], token)
	pos += len(token)
	v4 := ipv4.As4()
	copy(out[pos:pos+4], v4[:])
	pos += 4
	v6 := ipv6.As16()
	copy(out[pos:pos+16], v6[:])
	pos += 16
	binary.BigEndian.PutUint16(out[pos:pos+2], uint16(mtu))
	return out, nil
}

func ParseVPNOpen(payload []byte) (sid VPNSessionID, token string, ipv4, ipv6 netip.Addr, mtu int, err error) {
	if len(payload) < 1+16+2+4+16+2 || payload[0] != VPNCmdOpen {
		err = errors.New("bad VPN OPEN")
		return
	}
	copy(sid[:], payload[1:17])
	tokenLen := int(binary.BigEndian.Uint16(payload[17:19]))
	need := 1 + 16 + 2 + tokenLen + 4 + 16 + 2
	if tokenLen < 0 || len(payload) != need {
		err = errors.New("bad VPN OPEN length")
		return
	}
	pos := 19
	token = string(payload[pos : pos+tokenLen])
	pos += tokenLen
	var a4 [4]byte
	copy(a4[:], payload[pos:pos+4])
	ipv4 = netip.AddrFrom4(a4)
	pos += 4
	var a6 [16]byte
	copy(a6[:], payload[pos:pos+16])
	ipv6 = netip.AddrFrom16(a6)
	pos += 16
	mtu = int(binary.BigEndian.Uint16(payload[pos : pos+2]))
	return
}

func BuildVPNOpened(maxChunk int) []byte {
	if maxChunk > VPNMaxFragment {
		maxChunk = VPNMaxFragment
	}
	if maxChunk < 1 {
		maxChunk = 1
	}
	out := make([]byte, 3)
	out[0] = VPNRespOpened
	binary.BigEndian.PutUint16(out[1:3], uint16(maxChunk))
	return out
}

func ParseVPNOpened(payload []byte) (int, error) {
	if err := ParseVPNError(payload); err != nil {
		return 0, err
	}
	if len(payload) != 3 || payload[0] != VPNRespOpened {
		return 0, errors.New("bad VPN OPENED response")
	}
	return int(binary.BigEndian.Uint16(payload[1:3])), nil
}

// PUSH request: cmd(1) sid(16) seq(4) offset(2) total(2) data(N)
func BuildVPNPush(sid VPNSessionID, seq uint32, offset, total int, data []byte) ([]byte, error) {
	if total < 1 || total > 65535 || offset < 0 || offset > total || len(data) < 1 || offset+len(data) > total || len(data) > VPNMaxFragment {
		return nil, errors.New("invalid VPN PUSH fragment")
	}
	out := make([]byte, 25+len(data))
	out[0] = VPNCmdPush
	copy(out[1:17], sid[:])
	binary.BigEndian.PutUint32(out[17:21], seq)
	binary.BigEndian.PutUint16(out[21:23], uint16(offset))
	binary.BigEndian.PutUint16(out[23:25], uint16(total))
	copy(out[25:], data)
	return out, nil
}

func ParseVPNPush(payload []byte) (sid VPNSessionID, seq uint32, offset, total int, data []byte, err error) {
	if len(payload) < 26 || payload[0] != VPNCmdPush {
		err = errors.New("bad VPN PUSH")
		return
	}
	copy(sid[:], payload[1:17])
	seq = binary.BigEndian.Uint32(payload[17:21])
	offset = int(binary.BigEndian.Uint16(payload[21:23]))
	total = int(binary.BigEndian.Uint16(payload[23:25]))
	data = payload[25:]
	if total < 1 || offset < 0 || offset > total || len(data) < 1 || offset+len(data) > total {
		err = errors.New("bad VPN PUSH fragment bounds")
	}
	return
}

func BuildVPNAck(seq uint32, accepted int) []byte {
	out := make([]byte, 7)
	out[0] = VPNRespAck
	binary.BigEndian.PutUint32(out[1:5], seq)
	binary.BigEndian.PutUint16(out[5:7], uint16(accepted))
	return out
}

func ParseVPNAck(payload []byte) (seq uint32, accepted int, err error) {
	if e := ParseVPNError(payload); e != nil {
		err = e
		return
	}
	if len(payload) != 7 || payload[0] != VPNRespAck {
		err = errors.New("bad VPN ACK")
		return
	}
	seq = binary.BigEndian.Uint32(payload[1:5])
	accepted = int(binary.BigEndian.Uint16(payload[5:7]))
	return
}

// PULL request: cmd(1) sid(16) ack(4) want(4) offset(2) limit(2)
func BuildVPNPull(sid VPNSessionID, ack, want uint32, offset, limit int) ([]byte, error) {
	if offset < 0 || offset > 65535 || limit < 1 || limit > VPNMaxFragment {
		return nil, errors.New("invalid VPN PULL")
	}
	out := make([]byte, 29)
	out[0] = VPNCmdPull
	copy(out[1:17], sid[:])
	binary.BigEndian.PutUint32(out[17:21], ack)
	binary.BigEndian.PutUint32(out[21:25], want)
	binary.BigEndian.PutUint16(out[25:27], uint16(offset))
	binary.BigEndian.PutUint16(out[27:29], uint16(limit))
	return out, nil
}

func ParseVPNPull(payload []byte) (sid VPNSessionID, ack, want uint32, offset, limit int, err error) {
	if len(payload) != 29 || payload[0] != VPNCmdPull {
		err = errors.New("bad VPN PULL")
		return
	}
	copy(sid[:], payload[1:17])
	ack = binary.BigEndian.Uint32(payload[17:21])
	want = binary.BigEndian.Uint32(payload[21:25])
	offset = int(binary.BigEndian.Uint16(payload[25:27]))
	limit = int(binary.BigEndian.Uint16(payload[27:29]))
	if limit < 1 {
		err = errors.New("bad VPN PULL limit")
	}
	return
}

// DATA response: cmd(1) seq(4) offset(2) total(2) data(N)
func BuildVPNData(seq uint32, offset, total int, data []byte) []byte {
	out := make([]byte, 9+len(data))
	out[0] = VPNRespData
	binary.BigEndian.PutUint32(out[1:5], seq)
	binary.BigEndian.PutUint16(out[5:7], uint16(offset))
	binary.BigEndian.PutUint16(out[7:9], uint16(total))
	copy(out[9:], data)
	return out
}

func ParseVPNData(payload []byte) (seq uint32, offset, total int, data []byte, wait bool, err error) {
	if e := ParseVPNError(payload); e != nil {
		err = e
		return
	}
	if len(payload) == 1 && payload[0] == VPNRespWait {
		wait = true
		return
	}
	if len(payload) < 10 || payload[0] != VPNRespData {
		err = fmt.Errorf("bad VPN DATA response type/length")
		return
	}
	seq = binary.BigEndian.Uint32(payload[1:5])
	offset = int(binary.BigEndian.Uint16(payload[5:7]))
	total = int(binary.BigEndian.Uint16(payload[7:9]))
	data = payload[9:]
	if total < 1 || offset < 0 || offset+len(data) > total || len(data) < 1 {
		err = errors.New("bad VPN DATA bounds")
	}
	return
}

func BuildVPNClose(sid VPNSessionID) []byte {
	out := make([]byte, 17)
	out[0] = VPNCmdClose
	copy(out[1:17], sid[:])
	return out
}

func ParseVPNClose(payload []byte) (sid VPNSessionID, err error) {
	if len(payload) != 17 || payload[0] != VPNCmdClose {
		return sid, errors.New("bad VPN CLOSE")
	}
	copy(sid[:], payload[1:17])
	return sid, nil
}

func IsVPNCommand(payload []byte) bool {
	if len(payload) == 0 {
		return false
	}
	return payload[0] >= VPNCmdOpen && payload[0] <= VPNCmdClose
}