package policy

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/netip"
	"net/url"
	"strconv"
	"strings"
)

var (
	ErrInvalidAuthority = errors.New("invalid target authority")
	ErrTargetDenied     = errors.New("target address is denied by policy")
	ErrDNSResolution    = errors.New("target dns resolution failed")
)

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

type Config struct {
	Resolver             Resolver
	DenyCIDRs            []string
	AllowedPorts         []uint16
	AllowPrivateNetworks bool
	AllowLoopback        bool
	AllowLinkLocal       bool
}

type TargetPolicy struct {
	resolver             Resolver
	deny                 CIDRMatcher
	allowPrivateNetworks bool
	allowLoopback        bool
	allowLinkLocal       bool
	allowedPorts         map[uint16]struct{}
}

type Authority struct {
	Host       string
	Port       uint16
	LiteralIP  netip.Addr
	ResolvedIP netip.Addr
}

func (authority Authority) DialAddress() string {
	host := authority.Host
	if authority.ResolvedIP.IsValid() {
		host = authority.ResolvedIP.Unmap().String()
	} else if authority.LiteralIP.IsValid() {
		host = authority.LiteralIP.Unmap().String()
	}
	return net.JoinHostPort(host, strconv.FormatUint(uint64(authority.Port), 10))
}

type CIDRMatcher struct {
	prefixes []netip.Prefix
}

func NewTargetPolicy(config Config) (*TargetPolicy, error) {
	matcher, err := NewCIDRMatcher(config.DenyCIDRs)
	if err != nil {
		return nil, err
	}
	resolver := config.Resolver
	if resolver == nil {
		resolver = defaultResolver{}
	}
	ports := append([]uint16(nil), config.AllowedPorts...)
	if len(ports) == 0 {
		ports = []uint16{80, 443}
	}
	allowedPorts := make(map[uint16]struct{}, len(ports))
	for _, port := range ports {
		if port == 0 {
			return nil, fmt.Errorf("create target policy: allowed port must be positive")
		}
		allowedPorts[port] = struct{}{}
	}
	return &TargetPolicy{
		resolver:             resolver,
		deny:                 matcher,
		allowPrivateNetworks: config.AllowPrivateNetworks,
		allowLoopback:        config.AllowLoopback,
		allowLinkLocal:       config.AllowLinkLocal,
		allowedPorts:         allowedPorts,
	}, nil
}

func NewCIDRMatcher(cidrs []string) (CIDRMatcher, error) {
	prefixes := make([]netip.Prefix, 0, len(cidrs))
	for _, raw := range cidrs {
		raw = strings.TrimSpace(raw)
		if raw == "" {
			continue
		}
		prefix, err := netip.ParsePrefix(raw)
		if err != nil {
			return CIDRMatcher{}, fmt.Errorf("parse deny cidr %q: %w", raw, err)
		}
		prefixes = append(prefixes, prefix.Masked())
	}
	return CIDRMatcher{prefixes: prefixes}, nil
}

func (m CIDRMatcher) Match(addr netip.Addr) bool {
	if !addr.IsValid() {
		return false
	}
	addr = addr.Unmap()
	for _, prefix := range m.prefixes {
		if prefix.Contains(addr) {
			return true
		}
	}
	return false
}

func ParseURLAuthority(raw string) (Authority, error) {
	parsed, err := url.Parse(raw)
	if err != nil {
		return Authority{}, fmt.Errorf("%w: parse url: %v", ErrInvalidAuthority, err)
	}
	if parsed.Host == "" {
		return Authority{}, fmt.Errorf("%w: missing host", ErrInvalidAuthority)
	}

	port := parsed.Port()
	if port == "" {
		switch strings.ToLower(parsed.Scheme) {
		case "http":
			port = "80"
		case "https":
			port = "443"
		default:
			return Authority{}, fmt.Errorf("%w: unsupported url scheme %q", ErrInvalidAuthority, parsed.Scheme)
		}
	}
	return parseHostPort(parsed.Hostname(), port)
}

func ParseConnectAuthority(raw string) (Authority, error) {
	host, port, err := net.SplitHostPort(raw)
	if err != nil {
		return Authority{}, fmt.Errorf("%w: %v", ErrInvalidAuthority, err)
	}
	return parseHostPort(host, port)
}

func (p *TargetPolicy) EvaluateURL(ctx context.Context, raw string) (Authority, error) {
	authority, err := ParseURLAuthority(raw)
	if err != nil {
		return Authority{}, err
	}
	resolved, err := p.evaluateAuthority(ctx, authority)
	authority.ResolvedIP = resolved
	return authority, err
}

func (p *TargetPolicy) EvaluateConnectAuthority(ctx context.Context, raw string) (Authority, error) {
	authority, err := ParseConnectAuthority(raw)
	if err != nil {
		return Authority{}, err
	}
	resolved, err := p.evaluateAuthority(ctx, authority)
	authority.ResolvedIP = resolved
	return authority, err
}

func (p *TargetPolicy) EvaluateAuthority(ctx context.Context, authority Authority) error {
	_, err := p.evaluateAuthority(ctx, authority)
	return err
}

func (p *TargetPolicy) evaluateAuthority(ctx context.Context, authority Authority) (netip.Addr, error) {
	if p == nil {
		return netip.Addr{}, fmt.Errorf("evaluate target: nil policy")
	}
	if err := ctx.Err(); err != nil {
		return netip.Addr{}, err
	}
	if _, allowed := p.allowedPorts[authority.Port]; !allowed {
		return netip.Addr{}, fmt.Errorf("%w: destination port %d is not allowed", ErrTargetDenied, authority.Port)
	}
	if authority.LiteralIP.IsValid() {
		return authority.LiteralIP.Unmap(), p.validateAddress(authority.Host, authority.LiteralIP)
	}
	addrs, err := p.resolver.LookupNetIP(ctx, authority.Host)
	if err != nil {
		return netip.Addr{}, err
	}
	if len(addrs) == 0 {
		return netip.Addr{}, fmt.Errorf("%w: no addresses returned for %q", ErrDNSResolution, authority.Host)
	}
	for _, addr := range addrs {
		if err := p.validateAddress(authority.Host, addr); err != nil {
			return netip.Addr{}, err
		}
	}
	return addrs[0].Unmap(), nil
}

func (p *TargetPolicy) validateAddress(host string, addr netip.Addr) error {
	addr = addr.Unmap()
	if p.deny.Match(addr) {
		return fmt.Errorf("%w: %s matched deny cidr", ErrTargetDenied, addr)
	}
	if isSpecialUse(addr) {
		return fmt.Errorf("%w: %s is a special-use address", ErrTargetDenied, addr)
	}
	if p.allowLoopback && addr.IsLoopback() {
		return nil
	}
	if p.allowLinkLocal && addr.IsLinkLocalUnicast() {
		return nil
	}
	if p.allowPrivateNetworks && addr.IsPrivate() {
		return nil
	}
	if isDefaultDenied(addr) {
		return fmt.Errorf("%w: %s for host %q", ErrTargetDenied, addr, host)
	}
	return nil
}

func isSpecialUse(addr netip.Addr) bool {
	addr = addr.Unmap()
	for _, prefix := range specialUsePrefixes {
		if prefix.Contains(addr) {
			return true
		}
	}
	return false
}

var specialUsePrefixes = []netip.Prefix{
	netip.MustParsePrefix("0.0.0.0/8"),
	netip.MustParsePrefix("100.64.0.0/10"),
	netip.MustParsePrefix("168.63.129.16/32"),
	netip.MustParsePrefix("169.254.169.254/32"),
	netip.MustParsePrefix("169.254.170.2/32"),
	netip.MustParsePrefix("192.0.0.0/24"),
	netip.MustParsePrefix("192.0.2.0/24"),
	netip.MustParsePrefix("192.88.99.0/24"),
	netip.MustParsePrefix("198.18.0.0/15"),
	netip.MustParsePrefix("198.51.100.0/24"),
	netip.MustParsePrefix("203.0.113.0/24"),
	netip.MustParsePrefix("240.0.0.0/4"),
	netip.MustParsePrefix("64:ff9b::/96"),
	netip.MustParsePrefix("64:ff9b:1::/48"),
	netip.MustParsePrefix("100::/64"),
	netip.MustParsePrefix("2001::/23"),
	netip.MustParsePrefix("2001:db8::/32"),
	netip.MustParsePrefix("2002::/16"),
	netip.MustParsePrefix("3fff::/20"),
	netip.MustParsePrefix("fd00:ec2::254/128"),
}

func parseHostPort(host, port string) (Authority, error) {
	host = strings.TrimSpace(host)
	port = strings.TrimSpace(port)
	if host == "" || port == "" {
		return Authority{}, fmt.Errorf("%w: missing host or port", ErrInvalidAuthority)
	}
	numericPort, err := strconv.ParseUint(port, 10, 16)
	if err != nil || numericPort == 0 {
		return Authority{}, fmt.Errorf("%w: invalid port %q", ErrInvalidAuthority, port)
	}

	authority := Authority{
		Host: strings.TrimSuffix(host, "."),
		Port: uint16(numericPort),
	}
	if ip, err := netip.ParseAddr(authority.Host); err == nil {
		authority.Host = ip.String()
		authority.LiteralIP = ip.Unmap()
	}
	return authority, nil
}

func isDefaultDenied(addr netip.Addr) bool {
	return addr.IsLoopback() ||
		addr.IsPrivate() ||
		addr.IsLinkLocalUnicast() ||
		addr.IsLinkLocalMulticast() ||
		addr.IsUnspecified() ||
		addr.IsMulticast()
}

type defaultResolver struct{}

func (defaultResolver) LookupNetIP(ctx context.Context, host string) ([]netip.Addr, error) {
	addrs, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrDNSResolution, err)
	}
	return addrs, nil
}