// Package protocol implements the SoftEther VPN wire protocol.
//
// SoftEther uses a custom binary serialization format called "Pack" for all
// control-plane messages (hello, auth, welcome). A Pack contains Elements,
// each with a name, type, and one or more Values.
//
// Wire format (all big-endian):
//
//	Pack:    uint32(numElements) + Element...
//	Element: BufStr(name) + uint32(type) + uint32(numValues) + Value...
//	BufStr:  uint32(len) + bytes (null-terminated)
//	Value:   type-dependent (see below)
//
// Reference: https://github.com/SoftEtherVPN/SoftEtherVPN/blob/v5.02.5187/src/Mayaqua/Pack.c
package protocol

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"math/rand"
	"strings"
)

// Value types matching SoftEther's VALUE_* constants.
// See: https://github.com/SoftEtherVPN/SoftEtherVPN/blob/v5.02.5187/src/Mayaqua/Pack.h#L20-L28
const (
	ValueInt    uint32 = 0 // 32-bit unsigned integer
	ValueData   uint32 = 1 // Arbitrary binary data (prefixed with uint32 length)
	ValueStr    uint32 = 2 // ANSI string (prefixed with uint32 length)
	ValueUniStr uint32 = 3 // Unicode string as UTF-8 (prefixed with uint32 length, null-terminated)
	ValueInt64  uint32 = 4 // 64-bit unsigned integer
)

// Value holds a single typed value within an Element.
type Value struct {
	IntValue   uint32
	Int64Value uint64
	Data       []byte
	Str        string
}

// Element is a named, typed collection of values within a Pack.
type Element struct {
	Name   string
	Type   uint32
	Values []Value
}

// Pack is a collection of named Elements, used for all SoftEther control messages.
type Pack struct {
	Elements []*Element
}

// --- Reading ---
// See: https://github.com/SoftEtherVPN/SoftEtherVPN/blob/v5.02.5187/src/Mayaqua/Pack.c#L831 (PackRead)

// ReadPack deserializes a Pack from a binary reader.
func ReadPack(r io.Reader) (*Pack, error) {
	var num uint32
	if err := binary.Read(r, binary.BigEndian, &num); err != nil {
		return nil, err
	}
	if num > 262144 {
		return nil, fmt.Errorf("too many elements: %d", num)
	}

	p := &Pack{}
	for i := uint32(0); i < num; i++ {
		e, err := readElement(r)
		if err != nil {
			return nil, fmt.Errorf("element %d: %w", i, err)
		}
		p.Elements = append(p.Elements, e)
	}
	return p, nil
}

func readElement(r io.Reader) (*Element, error) {
	name, err := readBufStr(r)
	if err != nil {
		return nil, err
	}

	var typ, numValues uint32
	if err := binary.Read(r, binary.BigEndian, &typ); err != nil {
		return nil, err
	}
	if err := binary.Read(r, binary.BigEndian, &numValues); err != nil {
		return nil, err
	}

	e := &Element{Name: name, Type: typ}
	for i := uint32(0); i < numValues; i++ {
		v, err := readValue(r, typ)
		if err != nil {
			return nil, err
		}
		e.Values = append(e.Values, v)
	}
	return e, nil
}

func readValue(r io.Reader, typ uint32) (v Value, err error) {
	switch typ {
	case ValueInt:
		err = binary.Read(r, binary.BigEndian, &v.IntValue)
	case ValueInt64:
		err = binary.Read(r, binary.BigEndian, &v.Int64Value)
	case ValueData:
		var size uint32
		if err = binary.Read(r, binary.BigEndian, &size); err != nil {
			return
		}
		v.Data = make([]byte, size)
		_, err = io.ReadFull(r, v.Data)
	case ValueStr:
		var size uint32
		if err = binary.Read(r, binary.BigEndian, &size); err != nil {
			return
		}
		buf := make([]byte, size)
		if _, err = io.ReadFull(r, buf); err != nil {
			return
		}
		v.Str = string(buf)
	case ValueUniStr:
		var size uint32
		if err = binary.Read(r, binary.BigEndian, &size); err != nil {
			return
		}
		buf := make([]byte, size)
		if _, err = io.ReadFull(r, buf); err != nil {
			return
		}
		// SoftEther sends UTF-8 for UniStr on the wire
		v.Str = strings.TrimRight(string(buf), "\x00")
	default:
		err = fmt.Errorf("unknown value type: %d", typ)
	}
	return
}

// readBufStr reads a string written by SoftEther's WriteBufStr.
// Wire format: uint32(strlen+1) + strlen bytes (no null terminator on wire).
// The size field includes the null terminator in its count, but the null
// is NOT actually written to the wire.
// See: https://github.com/SoftEtherVPN/SoftEtherVPN/blob/v5.02.5187/src/Mayaqua/Memory.c#L2903
func readBufStr(r io.Reader) (string, error) {
	var size uint32
	if err := binary.Read(r, binary.BigEndian, &size); err != nil {
		return "", err
	}
	if size == 0 {
		return "", nil
	}
	// size includes the null terminator count, but only (size-1) bytes are on the wire
	buf := make([]byte, size-1)
	if _, err := io.ReadFull(r, buf); err != nil {
		return "", err
	}
	return string(buf), nil
}

// --- Writing ---
// See: https://github.com/SoftEtherVPN/SoftEtherVPN/blob/v5.02.5187/src/Mayaqua/Pack.c#L873 (PackWrite)

// Bytes serializes the Pack to binary format.
func (p *Pack) Bytes() ([]byte, error) {
	var buf bytes.Buffer
	if err := binary.Write(&buf, binary.BigEndian, uint32(len(p.Elements))); err != nil {
		return nil, err
	}
	for _, e := range p.Elements {
		if err := writeElement(&buf, e); err != nil {
			return nil, err
		}
	}
	return buf.Bytes(), nil
}

func writeElement(w io.Writer, e *Element) error {
	if err := writeBufStr(w, e.Name); err != nil {
		return err
	}
	if err := binary.Write(w, binary.BigEndian, e.Type); err != nil {
		return err
	}
	if err := binary.Write(w, binary.BigEndian, uint32(len(e.Values))); err != nil {
		return err
	}
	for _, v := range e.Values {
		if err := writeValue(w, &v, e.Type); err != nil {
			return err
		}
	}
	return nil
}

func writeValue(w io.Writer, v *Value, typ uint32) error {
	switch typ {
	case ValueInt:
		return binary.Write(w, binary.BigEndian, v.IntValue)
	case ValueInt64:
		return binary.Write(w, binary.BigEndian, v.Int64Value)
	case ValueData:
		if err := binary.Write(w, binary.BigEndian, uint32(len(v.Data))); err != nil {
			return err
		}
		_, err := w.Write(v.Data)
		return err
	case ValueStr:
		b := []byte(v.Str)
		if err := binary.Write(w, binary.BigEndian, uint32(len(b))); err != nil {
			return err
		}
		_, err := w.Write(b)
		return err
	case ValueUniStr:
		b := append([]byte(v.Str), 0) // UTF-8 with null terminator
		if err := binary.Write(w, binary.BigEndian, uint32(len(b))); err != nil {
			return err
		}
		_, err := w.Write(b)
		return err
	default:
		return fmt.Errorf("unknown value type: %d", typ)
	}
}

// writeBufStr writes a string in SoftEther's WriteBufStr format.
// Wire format: uint32(strlen+1) + strlen bytes (no null terminator on wire).
func writeBufStr(w io.Writer, s string) error {
	b := []byte(s)
	if err := binary.Write(w, binary.BigEndian, uint32(len(b)+1)); err != nil {
		return err
	}
	_, err := w.Write(b)
	return err
}

// --- Accessors ---

// GetElement finds an element by name (case-insensitive).
func (p *Pack) GetElement(name string) *Element {
	upper := strings.ToUpper(name)
	for _, e := range p.Elements {
		if strings.ToUpper(e.Name) == upper {
			return e
		}
	}
	return nil
}

// GetInt returns the first uint32 value of the named element.
func (p *Pack) GetInt(name string) uint32 {
	if e := p.GetElement(name); e != nil && e.Type == ValueInt && len(e.Values) > 0 {
		return e.Values[0].IntValue
	}
	return 0
}

// GetInt64 returns the first uint64 value of the named element.
func (p *Pack) GetInt64(name string) uint64 {
	if e := p.GetElement(name); e != nil && e.Type == ValueInt64 && len(e.Values) > 0 {
		return e.Values[0].Int64Value
	}
	return 0
}

// GetStr returns the first string value of the named element (STR or UNISTR).
func (p *Pack) GetStr(name string) string {
	if e := p.GetElement(name); e != nil && len(e.Values) > 0 {
		if e.Type == ValueStr || e.Type == ValueUniStr {
			return e.Values[0].Str
		}
	}
	return ""
}

// GetData returns the first data value of the named element.
func (p *Pack) GetData(name string) []byte {
	if e := p.GetElement(name); e != nil && e.Type == ValueData && len(e.Values) > 0 {
		return e.Values[0].Data
	}
	return nil
}

// GetBool returns the first integer value as a boolean (non-zero = true).
func (p *Pack) GetBool(name string) bool {
	return p.GetInt(name) != 0
}

// GetError returns the "error" field, which holds SoftEther error codes.
// See: https://github.com/SoftEtherVPN/SoftEtherVPN/blob/v5.02.5187/src/Cedar/Cedar.h#L590
func (p *Pack) GetError() uint32 {
	return p.GetInt("error")
}

// --- Mutators ---

// AddInt adds a uint32 element.
func (p *Pack) AddInt(name string, val uint32) {
	p.Elements = append(p.Elements, &Element{
		Name:   name,
		Type:   ValueInt,
		Values: []Value{{IntValue: val}},
	})
}

// AddInt64 adds a uint64 element.
func (p *Pack) AddInt64(name string, val uint64) {
	p.Elements = append(p.Elements, &Element{
		Name:   name,
		Type:   ValueInt64,
		Values: []Value{{Int64Value: val}},
	})
}

// AddStr adds a string element.
func (p *Pack) AddStr(name string, val string) {
	p.Elements = append(p.Elements, &Element{
		Name:   name,
		Type:   ValueStr,
		Values: []Value{{Str: val}},
	})
}

// AddData adds a binary data element.
func (p *Pack) AddData(name string, val []byte) {
	p.Elements = append(p.Elements, &Element{
		Name:   name,
		Type:   ValueData,
		Values: []Value{{Data: val}},
	})
}

// AddBool adds a boolean element (stored as uint32: 0 or 1).
func (p *Pack) AddBool(name string, val bool) {
	v := uint32(0)
	if val {
		v = 1
	}
	p.AddInt(name, v)
}

// AddDummyValue adds the "pencore" random padding element.
// SoftEther servers expect this padding in Pack messages sent via HTTP.
// See: https://github.com/SoftEtherVPN/SoftEtherVPN/blob/v5.02.5187/src/Mayaqua/HTTP.c#L1165
func (p *Pack) AddDummyValue() {
	size := rand.Intn(1000)
	buf := make([]byte, size)
	rand.Read(buf)
	p.AddData("pencore", buf)
}

// AddIP4 adds an IPv4 address with the required companion @ipv6_* elements.
// SoftEther always expects these three companion elements for any IP field.
// See: https://github.com/SoftEtherVPN/SoftEtherVPN/blob/v5.02.5187/src/Mayaqua/Pack.c#L575 (PackAddIp32)
func (p *Pack) AddIP4(name string, ip uint32) {
	p.AddBool(name+"@ipv6_bool", false)
	p.AddData(name+"@ipv6_array", make([]byte, 16))
	p.AddInt(name+"@ipv6_scope_id", 0)
	p.AddInt(name, ip)
}