package pack

import (
	"encoding/base32"
	"encoding/base64"
	"encoding/binary"
	"encoding/hex"
	"net"
	"net/netip"
	"strings"

	"codeberg.org/miekg/dns/internal/ddd"
)

const maxCompressionOffset = 2 << 13 // We have 14 bits for the compression pointer

// maybe this should all moved to cryptobyte as well...
// near future direction is clear all pack helpers should be here, not in msg_helpers.go

func Uint8(i uint8, msg []byte, off int) (off1 int, err error) {
	if off+1 > len(msg) {
		return len(msg), &Error{"overflow uint8"}
	}
	msg[off] = i
	return off + 1, nil
}

func Uint16(i uint16, msg []byte, off int) (off1 int, err error) {
	if off+2 > len(msg) {
		return len(msg), &Error{"overflow uint16"}
	}
	_ = msg[off+1]
	binary.BigEndian.PutUint16(msg[off:], i)
	return off + 2, nil
}

func Uint32(i uint32, msg []byte, off int) (off1 int, err error) {
	if off+4 > len(msg) {
		return len(msg), &Error{"overflow uint32"}
	}
	_ = msg[off+3]
	binary.BigEndian.PutUint32(msg[off:], i)
	return off + 4, nil
}

func Uint48(i uint64, msg []byte, off int) (off1 int, err error) {
	if off+6 > len(msg) {
		return len(msg), &Error{"overflow uint64 as uint48"}
	}
	_ = msg[off+5]
	msg[off] = byte(i >> 40)
	msg[off+1] = byte(i >> 32)
	msg[off+2] = byte(i >> 24)
	msg[off+3] = byte(i >> 16)
	msg[off+4] = byte(i >> 8)
	msg[off+5] = byte(i)
	off += 6
	return off, nil
}

func Uint64(i uint64, msg []byte, off int) (off1 int, err error) {
	if off+8 > len(msg) {
		return len(msg), &Error{"overflow uint64"}
	}
	_ = msg[off+7]
	binary.BigEndian.PutUint64(msg[off:], i)
	off += 8
	return off, nil
}

// StringAny packs a string as-is, no decoding or length bytes are written.
func StringAny(s string, msg []byte, off int) (int, error) {
	if off+len(s) > len(msg) {
		return len(msg), &Error{"overflow string anything"}
	}
	copy(msg[off:off+len(s)], s)
	off += len(s)
	return off, nil
}

func Strings(s []string, msg []byte, off int) (int, error) {
	if len(s) == 0 {
		if off >= len(msg) {
			return len(msg), ErrBuf
		}
		msg[off] = 0
		return off, nil
	}
	var err error
	for i := range s {
		off, err = String(s[i], msg, off)
		if err != nil {
			return len(msg), err
		}
	}
	return off, nil
}

func String(s string, msg []byte, off int) (int, error) {
	if strings.IndexByte(s, '\\') == -1 {
		l := len(s)
		if l > 255 {
			return len(msg), &Error{"overflow string"}
		}
		if off+1+l > len(msg) {
			return len(msg), &Error{"overflow string"}
		}
		msg[off] = byte(l)
		copy(msg[off+1:], s)
		return off + 1 + l, nil
	}

	lenByteoff := off
	if off >= len(msg) || len(s) > 256*4+1 /* If all \DDD */ {
		return len(msg), &Error{"overflow string"}
	}
	off++

	for i := 0; i < len(s); i++ {
		if len(msg) <= off {
			return len(msg), &Error{"overflow string"}
		}
		if s[i] == '\\' {
			i++
			if i == len(s) {
				break
			}
			// check for \DDD
			if ddd.Is(s[i:]) {
				msg[off] = ddd.ToByte(s[i:])
				i += 2
			} else {
				msg[off] = s[i]
			}
		} else {
			msg[off] = s[i]
		}
		off++
	}
	l := off - lenByteoff - 1
	if l > 255 {
		return len(msg), &Error{"string exceeded 255 bytes in txt"}
	}
	msg[lenByteoff] = byte(l)
	return off, nil
}

func A(a netip.Addr, msg []byte, off int) (int, error) {
	if off+net.IPv4len > len(msg) {
		return len(msg), &Error{"overflow a"}
	}
	val := a.As4()
	_ = msg[off+net.IPv4len-1]
	msg[off] = val[0]
	msg[off+1] = val[1]
	msg[off+2] = val[2]
	msg[off+3] = val[3]
	return off + net.IPv4len, nil
}

func AAAA(aaaa netip.Addr, msg []byte, off int) (int, error) {
	if off+net.IPv6len > len(msg) {
		return len(msg), &Error{"overflow aaaa"}
	}
	val := aaaa.As16()
	_ = msg[off+net.IPv6len-1]
	msg[off] = val[0]
	msg[off+1] = val[1]
	msg[off+2] = val[2]
	msg[off+3] = val[3]
	msg[off+4] = val[4]
	msg[off+5] = val[5]
	msg[off+6] = val[6]
	msg[off+7] = val[7]
	msg[off+8] = val[8]
	msg[off+9] = val[9]
	msg[off+10] = val[10]
	msg[off+11] = val[11]
	msg[off+12] = val[12]
	msg[off+13] = val[13]
	msg[off+14] = val[14]
	msg[off+15] = val[15]
	return off + net.IPv6len, nil
}

// Name packs the string s into msg.
func Name(s string, msg []byte, off int, compression map[string]uint16, compress bool) (off1 int, err error) {
	// XXX: A logical copy of this function exists in dnsutil.IsName and should be kept in sync with this function.

	lenmsg := len(msg)
	ls := len(s)

	if ls == 1 && s[0] == '.' {
		msg[off] = 0
		return off + 1, nil
	}
	if ls > 1 && s[0] == '.' { // leading dots are not legal except for the root zone
		return len(msg), &Error{"leading dot in name"}
	}
	if s[ls-1] != '.' {
		return len(msg), &Error{"name must be fully qualified"}
	}

	// Each dot ends a segment of the name. We trade each dot byte for a length byte.
	// There is also a trailing zero. Emit sequence of counted strings, chopping at dots.
	var (
		begin    int
		labelLen int
	)

	for begin < ls {
		i := strings.IndexByte(s[begin:], '.')
		if i == -1 {
			break
		}
		i += begin

		labelLen = i - begin
		if labelLen >= 1<<6 { // top two bits of length must be clear
			return lenmsg, &Error{"illegal label type in name"}
		}
		if labelLen == 0 {
			return lenmsg, &Error{"consecutive dots in name"}
		}

		// off can already (we're in a loop) be bigger than len(msg)
		if off+labelLen >= lenmsg {
			return lenmsg, &Error{"overflow name"}
		}

		if compress && labelLen > 1 { // don't try to compress '.'
			if p, ok := compression[s[begin:]]; ok {
				binary.BigEndian.PutUint16(msg[off:], 0xC000|p)
				return off + 2, nil
			}
		}
		if compression != nil && off < maxCompressionOffset {
			compression[s[begin:]] = uint16(off)
		}

		// the following is covered by the length check above
		msg[off] = byte(labelLen)
		copy(msg[off+1:], s[begin:i])

		off += 1 + labelLen
		begin = i + 1
	}
	if off >= lenmsg {
		return lenmsg, &Error{"overflow name"}
	}

	msg[off] = 0
	return off + 1, nil
}

// MName packs the string s into msg taking escaped dots into account.
func MName(s string, msg []byte, off int) (off1 int, err error) {
	lenmsg := len(msg)
	ls := len(s)

	if ls == 1 && s[0] == '.' {
		msg[off] = 0
		return off + 1, nil
	}
	if ls > 1 && s[0] == '.' { // leading dots are not legal except for the root zone
		return len(msg), &Error{"leading dot in name"}
	}
	if s[ls-1] != '.' {
		return len(msg), &Error{"name must be fully qualified"}
	}

	// Each dot ends a segment of the name. We trade each dot byte for a length byte.
	// Except for escaped dots (\.), which are normal dots. There is also a trailing zero.
	// Emit sequence of counted strings, chopping at dots.
	var (
		begin    int
		labelLen int
		escape   int
	)

	for begin < ls {
		i := strings.IndexByte(s[begin:], '.')
		if i == -1 {
			break
		}

		escape = 0
		if i > 0 && s[begin+i-1] == '\\' { // escaped dot
			if begin+i+2 > ls {
				return lenmsg, &Error{"overflow name after escape"}
			}
			j := strings.IndexByte(s[begin+i+1:], '.') // search ahead for another dot
			if j == -1 {
				break // no dot found, break as above
			}
			// i is the end of the label, 1 for \
			i += 1 + j
			escape = 1
		}

		i += begin
		labelLen = i - begin - escape

		if labelLen >= 1<<6 { // top two bits of length must be clear
			return lenmsg, &Error{"illegal label type in name"}
		}
		if labelLen == 0 {
			return lenmsg, &Error{"consecutive dots in name"}
		}

		// off can already (we're in a loop) be bigger than len(msg)
		if off+labelLen >= lenmsg {
			return lenmsg, &Error{"overflow name"}
		}

		msg[off] = byte(labelLen)
		if escape == 0 {
			copy(msg[off+1:], s[begin:i])
			off += 1 + labelLen
			begin = i + 1
			continue
		}

		// slow(er) copy path to skip the escaped dot
		l := 1
		for k := begin; k < i; k++ {
			if s[k] == '\\' {
				continue
			}
			msg[off+l] = s[k]
			l++
		}
		off += 1 + labelLen
		begin = i + 1
	}

	if off >= lenmsg {
		return lenmsg, &Error{"overflow name"}
	}

	msg[off] = 0
	return off + 1, nil
}

func StringBase32(s string, msg []byte, off int) (int, error) {
	b32, err := Base32([]byte(s))
	if err != nil {
		return len(msg), err
	}
	if off+len(b32) > len(msg) {
		return len(msg), &Error{Err: "overflow base32"}
	}
	copy(msg[off:off+len(b32)], b32)
	off += len(b32)
	return off, nil
}

func StringBase64(s string, msg []byte, off int) (int, error) {
	b64, err := Base64([]byte(s))
	if err != nil {
		return len(msg), err
	}
	if off+len(b64) > len(msg) {
		return len(msg), &Error{Err: "overflow base64"}
	}
	copy(msg[off:off+len(b64)], b64)
	off += len(b64)
	return off, nil
}

func Base32(s []byte) (buf []byte, err error) {
	for i, b := range s {
		if b >= 'a' && b <= 'z' {
			s[i] = b - 32
		}
	}
	b32hex := base32.HexEncoding.WithPadding(base32.NoPadding)
	buflen := b32hex.DecodedLen(len(s))
	buf = make([]byte, buflen)
	n, err := b32hex.Decode(buf, s)
	buf = buf[:n]
	return
}

func Base64(s []byte) (buf []byte, err error) {
	buflen := base64.StdEncoding.DecodedLen(len(s))
	buf = make([]byte, buflen)
	n, err := base64.StdEncoding.Decode(buf, s)
	buf = buf[:n]
	return
}

func StringHex(s string, msg []byte, off int) (int, error) {
	h, err := hex.DecodeString(s)
	if err != nil {
		return len(msg), err
	}
	if off+len(h) > len(msg) {
		return len(msg), &Error{Err: "overflow hex"}
	}
	copy(msg[off:off+len(h)], h)
	off += len(h)
	return off, nil
}

func Names(names []string, msg []byte, off int, compress map[string]uint16) (int, error) {
	var err error
	for i := range names {
		off, err = Name(names[i], msg, off, compress, false)
		if err != nil {
			return len(msg), err
		}
	}
	return off, nil
}
