package magna

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

// decodeDomain decodes a domain name from a buffer starting at offset.
// It returns the domain name along with the offset and error.
func decodeDomain(buf []byte, offset int) (string, int, error) {
	var builder strings.Builder
	firstLabel := true

	seenOffsets := make(map[int]struct{})
	finalOffsetAfterJump := -1

	currentOffset := offset

	for {
		if _, found := seenOffsets[currentOffset]; found {
			return "", len(buf), &DomainCompressionError{}
		}
		seenOffsets[currentOffset] = struct{}{}

		length, nextOffsetAfterLen, err := getU8(buf, currentOffset)
		if err != nil {
			return "", len(buf), fmt.Errorf("failed to read domain label length: %w", err)
		}

		if length == 0 {
			currentOffset = nextOffsetAfterLen
			break
		}

		if (length & 0xC0) == 0xC0 {
			sec, nextOffsetAfterPtr, err := getU8(buf, nextOffsetAfterLen)
			if err != nil {
				return "", len(buf), fmt.Errorf("failed to read domain compression pointer offset byte: %w", err)
			}

			jumpTargetOffset := int(length&0x3F)<<8 | int(sec)

			if jumpTargetOffset >= len(buf) {
				return "", len(buf), &BufferOverflowError{Length: len(buf), Offset: jumpTargetOffset}
			}
			if _, found := seenOffsets[jumpTargetOffset]; found {
				return "", len(buf), &DomainCompressionError{}
			}

			if finalOffsetAfterJump == -1 {
				finalOffsetAfterJump = nextOffsetAfterPtr
			}

			currentOffset = jumpTargetOffset
			continue
		}

		if length > 63 {
			return "", len(buf), &InvalidLabelError{Length: int(length)}
		}

		labelBytes, nextOffsetAfterLabel, err := getSlice(buf, nextOffsetAfterLen, int(length))
		if err != nil {
			return "", len(buf), fmt.Errorf("failed to read domain label data: %w", err)
		}

		if !firstLabel {
			builder.WriteByte('.')
		}
		builder.Write(labelBytes)
		firstLabel = false

		currentOffset = nextOffsetAfterLabel

	}

	finalReadOffset := currentOffset
	if finalOffsetAfterJump != -1 {
		finalReadOffset = finalOffsetAfterJump
	}

	return builder.String(), finalReadOffset, nil
}

// encodeDomain returns the bytes of the input bytes appened with the encoded domain name.
func encodeDomain(bytes []byte, domainName string, offsets *map[string]uint16) ([]byte, error) {
	if domainName == "." || domainName == "" {
		return append(bytes, 0), nil
	}

	cleanDomain := strings.TrimSuffix(domainName, ".")
	if cleanDomain == "" {
		return append(bytes, 0), nil
	}

	start := 0
	for start < len(cleanDomain) {
		suffix := cleanDomain[start:]

		if offset, found := (*offsets)[suffix]; found {
			pointer := 0xC000 | offset
			bytes = binary.BigEndian.AppendUint16(bytes, pointer)
			return bytes, nil
		}

		currentPos := uint16(len(bytes))
		if currentPos <= 0x3FFF {
			(*offsets)[suffix] = currentPos
		}

		end := strings.IndexByte(suffix, '.')
		var label string
		nextStart := len(cleanDomain)

		if end == -1 {
			label = suffix
			start = nextStart
		} else {
			label = suffix[:end]
			nextStart = start + end + 1
			start = nextStart
		}

		labelBytes := []byte(label)

		if len(labelBytes) > 63 {
			return nil, &InvalidLabelError{Length: int(len(labelBytes))}
		}

		if len(labelBytes) == 0 && start < len(cleanDomain) {
			return nil, &InvalidLabelError{Length: 0}
		}

		bytes = append(bytes, byte(len(labelBytes)))
		bytes = append(bytes, labelBytes...)
	}

	bytes = append(bytes, 0)
	return bytes, nil
}