package transport

import (
	"context"
	"fmt"
	"net"
	"net/netip"
	"regexp"
	"strings"

	"golang.org/x/net/idna"
)

const maxResolvedAddresses = 16

var dnsLabelPattern = regexp.MustCompile(
	`^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$`,
)

var forbiddenAddressPrefixes = mustPrefixes(
	"0.0.0.0/8",
	"100.64.0.0/10",
	"127.0.0.0/8",
	"169.254.0.0/16",
	"192.0.0.0/24",
	"192.0.2.0/24",
	"192.88.99.0/24",
	"198.18.0.0/15",
	"198.51.100.0/24",
	"203.0.113.0/24",
	"224.0.0.0/4",
	"240.0.0.0/4",
	"64:ff9b::/96",
	"64:ff9b:1::/48",
	"100::/64",
	"2001::/23",
	"2001:db8::/32",
	"2002::/16",
	"fc00::/7",
	"fe80::/10",
	"ff00::/8",
)

type Resolver interface {
	LookupNetIP(
		context.Context,
		string,
		string,
	) ([]netip.Addr, error)
}

type NetResolver struct {
	Resolver *net.Resolver
}

func (resolver NetResolver) LookupNetIP(
	ctx context.Context,
	network string,
	host string,
) ([]netip.Addr, error) {
	selected := resolver.Resolver
	if selected == nil {
		selected = net.DefaultResolver
	}
	return selected.LookupNetIP(ctx, network, host)
}

func NormalizeHostname(value string) (string, error) {
	trimmed := strings.TrimSuffix(strings.TrimSpace(value), ".")
	if trimmed == "" || len(trimmed) > 253 {
		return "", fmt.Errorf("mail endpoint host is invalid")
	}
	if _, err := netip.ParseAddr(trimmed); err == nil {
		return "", fmt.Errorf("mail endpoint IP literals are forbidden")
	}
	ascii, err := idna.Lookup.ToASCII(trimmed)
	if err != nil {
		return "", fmt.Errorf("mail endpoint host is invalid: %w", err)
	}
	ascii = strings.ToLower(ascii)
	if ascii == "localhost" ||
		strings.HasSuffix(ascii, ".localhost") ||
		strings.HasSuffix(ascii, ".local") ||
		!strings.Contains(ascii, ".") {
		return "", fmt.Errorf("mail endpoint local host is forbidden")
	}
	for _, label := range strings.Split(ascii, ".") {
		if !dnsLabelPattern.MatchString(label) {
			return "", fmt.Errorf("mail endpoint host label is invalid")
		}
	}
	return ascii, nil
}

func ResolveAllowed(
	ctx context.Context,
	resolver Resolver,
	host string,
) (string, []netip.Addr, error) {
	normalized, err := NormalizeHostname(host)
	if err != nil {
		return "", nil, err
	}
	if resolver == nil {
		resolver = NetResolver{}
	}
	addresses, err := resolver.LookupNetIP(ctx, "ip", normalized)
	if err != nil {
		return "", nil, fmt.Errorf("resolve mail endpoint: %w", err)
	}
	if len(addresses) == 0 || len(addresses) > maxResolvedAddresses {
		return "", nil, fmt.Errorf(
			"mail endpoint address set is invalid",
		)
	}
	result := make([]netip.Addr, 0, len(addresses))
	seen := make(map[netip.Addr]struct{}, len(addresses))
	for _, address := range addresses {
		address = address.Unmap()
		if err := ValidateAddress(address); err != nil {
			return "", nil, err
		}
		if _, duplicate := seen[address]; duplicate {
			continue
		}
		seen[address] = struct{}{}
		result = append(result, address)
	}
	if len(result) == 0 {
		return "", nil, fmt.Errorf("mail endpoint has no usable address")
	}
	return normalized, result, nil
}

func ValidateAddress(address netip.Addr) error {
	address = address.Unmap()
	if !address.IsValid() ||
		!address.IsGlobalUnicast() ||
		address.IsPrivate() ||
		address.IsLoopback() ||
		address.IsLinkLocalUnicast() ||
		address.IsLinkLocalMulticast() ||
		address.IsMulticast() ||
		address.IsUnspecified() {
		return fmt.Errorf("mail endpoint address is forbidden")
	}
	for _, prefix := range forbiddenAddressPrefixes {
		if prefix.Contains(address) {
			return fmt.Errorf("mail endpoint address is forbidden")
		}
	}
	return nil
}

func mustPrefixes(values ...string) []netip.Prefix {
	result := make([]netip.Prefix, 0, len(values))
	for _, value := range values {
		result = append(result, netip.MustParsePrefix(value))
	}
	return result
}
