package utils

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
)

type LookupIPAddrFunc func(context.Context, string) ([]net.IPAddr, error)

func ValidatePublicEndpointHost(ctx context.Context, host string) error {
	return ValidatePublicEndpointHostWithResolver(ctx, host, net.DefaultResolver.LookupIPAddr)
}

func ValidatePublicEndpointHostWithResolver(
	ctx context.Context,
	host string,
	lookupIPAddr LookupIPAddrFunc,
) error {
	host = strings.Trim(strings.TrimSpace(host), "[]")
	if host == "" {
		return errors.New("host is empty")
	}
	if ip := net.ParseIP(host); ip != nil {
		return validatePublicEndpointIP(host, ip)
	}
	if lookupIPAddr == nil {
		return errors.New("resolver is nil")
	}
	ips, err := lookupIPAddr(ctx, host)
	if err != nil {
		return fmt.Errorf("resolve host %s failed: %w", host, err)
	}
	if len(ips) == 0 {
		return fmt.Errorf("resolve host %s returned no ip", host)
	}
	for _, resolved := range ips {
		if err := validatePublicEndpointIP(host, resolved.IP); err != nil {
			return err
		}
	}
	return nil
}

func validatePublicEndpointIP(host string, ip net.IP) error {
	if ip == nil {
		return fmt.Errorf("host %s resolves to empty ip", host)
	}
	if IsDisallowedPublicEndpointIP(ip) {
		return fmt.Errorf("host %s resolves to non-public ip %s", host, ip.String())
	}
	return nil
}

func IsDisallowedPublicEndpointIP(ip net.IP) bool {
	if ip == nil {
		return true
	}
	if ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast() ||
		ip.IsLinkLocalMulticast() || ip.IsUnspecified() || ip.IsMulticast() {
		return true
	}
	if ipv4 := ip.To4(); ipv4 != nil {
		return ipv4[0] == 100 && ipv4[1] >= 64 && ipv4[1] <= 127
	}
	return false
}