package main

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

type PolicyError struct {
	Code   string
	Reason string
}

func (e *PolicyError) Error() string {
	if e.Reason == "" {
		return e.Code
	}
	return e.Code + ": " + e.Reason
}

func refuse(code, reason string) error {
	return &PolicyError{Code: code, Reason: reason}
}

type Target struct {
	URL    string
	Method string
	Host   string
	Port   int
	Dest   Destination
}

type Policy struct {
	mu     sync.RWMutex
	dests  []Destination
	lookup func(ctx context.Context, host string) ([]netip.Addr, error)
	public func(netip.Addr) bool
	dial   func(ctx context.Context, network, address string) (net.Conn, error)
}

func NewPolicy() *Policy {
	dialer := &net.Dialer{Timeout: time.Duration(TunnelDialSeconds) * time.Second}
	return &Policy{
		dests:  cloneDestinations(BakedDestinations),
		lookup: systemLookup,
		public: IsPublicAddr,
		dial:   dialer.DialContext,
	}
}

func systemLookup(ctx context.Context, host string) ([]netip.Addr, error) {
	return net.DefaultResolver.LookupNetIP(ctx, "ip", host)
}

func cloneDestinations(list []Destination) []Destination {
	out := make([]Destination, len(list))
	for i, item := range list {
		item.Methods = slices.Clone(item.Methods)
		item.Ports = slices.Clone(item.Ports)
		item.Modes = slices.Clone(item.Modes)
		out[i] = item
	}
	return out
}

func normalizeHost(host string) string {
	return strings.TrimRight(strings.ToLower(strings.TrimSpace(host)), ".")
}

func suffixMatches(suffix, host string) bool {
	return host == suffix || strings.HasSuffix(host, "."+suffix)
}

func intersectStrings(a, b []string) []string {
	out := []string{}
	for _, item := range a {
		if slices.ContainsFunc(b, func(other string) bool { return strings.EqualFold(item, other) }) {
			out = append(out, item)
		}
	}
	return out
}

func intersectPorts(a, b []int) []int {
	out := []int{}
	for _, item := range a {
		if slices.Contains(b, item) && slices.Contains(AllowedPorts, item) {
			out = append(out, item)
		}
	}
	return out
}

func (p *Policy) Refine(offered []Destination) []string {
	var rejected []string
	narrowed := []Destination{}
	seen := map[string]bool{}
	for _, item := range offered {
		suffix := normalizeHost(item.Suffix)
		var base *Destination
		for i := range BakedDestinations {
			if suffixMatches(BakedDestinations[i].Suffix, suffix) {
				base = &BakedDestinations[i]
				break
			}
		}
		if base == nil || seen[suffix] {
			rejected = append(rejected, item.Suffix)
			continue
		}
		seen[suffix] = true
		narrowed = append(narrowed, Destination{
			Suffix:  suffix,
			Adapter: base.Adapter,
			Methods: intersectStrings(base.Methods, item.Methods),
			Ports:   intersectPorts(base.Ports, item.Ports),
			Modes:   intersectStrings(base.Modes, item.Modes),
			MaxRPS:  item.MaxRPS,
		})
	}
	p.mu.Lock()
	p.dests = narrowed
	p.mu.Unlock()
	return rejected
}

func (p *Policy) Destinations() []Destination {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return cloneDestinations(p.dests)
}

func (p *Policy) destinationFor(host string) (Destination, bool) {
	p.mu.RLock()
	defer p.mu.RUnlock()
	var best Destination
	found := false
	for _, item := range p.dests {
		if suffixMatches(item.Suffix, host) && (!found || len(item.Suffix) > len(best.Suffix)) {
			best, found = item, true
		}
	}
	return best, found
}

func isAddressLiteral(host string) bool {
	_, err := netip.ParseAddr(strings.Trim(host, "[]"))
	return err == nil
}

func (p *Policy) CheckHostPort(host string, port int, mode string) (Destination, error) {
	host = normalizeHost(host)
	if host == "" {
		return Destination{}, refuse(RefuseBadURL, "no host")
	}
	if isAddressLiteral(host) {
		return Destination{}, refuse(RefuseHost, "address literals are not allowed")
	}
	dest, ok := p.destinationFor(host)
	if !ok {
		return Destination{}, refuse(RefuseHost, "destination is not on the allowlist")
	}
	if !slices.Contains(AllowedPorts, port) || !slices.Contains(dest.Ports, port) {
		return Destination{}, refuse(RefusePort, "port is not allowed")
	}
	if mode != "" && mode != "any" && !slices.Contains(dest.Modes, mode) {
		return Destination{}, refuse(RefuseHost, "mode is not allowed for this destination")
	}
	return dest, nil
}

func (p *Policy) Check(rawURL, method, mode string) (Target, error) {
	if rawURL == "" || len(rawURL) > MaxURLChars {
		return Target{}, refuse(RefuseBadURL, "url is missing or too long")
	}
	parsed, err := url.Parse(strings.TrimSpace(rawURL))
	if err != nil {
		return Target{}, refuse(RefuseBadURL, "url cannot be parsed")
	}
	defaultPort := 0
	switch parsed.Scheme {
	case "http":
		defaultPort = 80
	case "https":
		defaultPort = 443
	default:
		return Target{}, refuse(RefuseBadURL, "only http and https are allowed")
	}
	if parsed.User != nil {
		return Target{}, refuse(RefuseBadURL, "credentials in a url are not allowed")
	}
	port := defaultPort
	if text := parsed.Port(); text != "" {
		port, err = strconv.Atoi(text)
		if err != nil {
			return Target{}, refuse(RefuseBadURL, "port is not a number")
		}
	}
	dest, err := p.CheckHostPort(parsed.Hostname(), port, mode)
	if err != nil {
		return Target{}, err
	}
	verb := strings.ToUpper(strings.TrimSpace(method))
	if !slices.Contains(upperAll(dest.Methods), verb) {
		return Target{}, refuse(RefuseMethod, "method is not allowed here")
	}
	return Target{URL: parsed.String(), Method: verb, Host: normalizeHost(parsed.Hostname()), Port: port, Dest: dest}, nil
}

func upperAll(list []string) []string {
	out := make([]string, len(list))
	for i, item := range list {
		out[i] = strings.ToUpper(item)
	}
	return out
}

func (p *Policy) CheckRedirect(location, method string) (Target, error) {
	target, err := p.Check(location, method, "fetch")
	if err != nil {
		var policyErr *PolicyError
		if errors.As(err, &policyErr) {
			return Target{}, refuse(RefuseRedirect, policyErr.Reason)
		}
		return Target{}, refuse(RefuseRedirect, err.Error())
	}
	return target, nil
}

func (p *Policy) Resolve(ctx context.Context, host string) ([]netip.Addr, error) {
	host = normalizeHost(host)
	addrs, err := p.lookup(ctx, host)
	if err != nil {
		return nil, refuse(FailDNS, "host did not resolve")
	}
	if len(addrs) == 0 {
		return nil, refuse(FailDNS, "host has no addresses")
	}
	if len(addrs) > 16 {
		addrs = addrs[:16]
	}
	for _, addr := range addrs {
		if !p.public(addr.Unmap()) {
			return nil, refuse(RefuseAddress, "host resolves to a non public address")
		}
	}
	return addrs, nil
}

func (p *Policy) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	host, portText, err := net.SplitHostPort(address)
	if err != nil {
		return nil, refuse(RefuseBadURL, "bad address")
	}
	port, err := strconv.Atoi(portText)
	if err != nil {
		return nil, refuse(RefuseBadURL, "bad port")
	}
	if _, err := p.CheckHostPort(host, port, ""); err != nil {
		return nil, err
	}
	addrs, err := p.Resolve(ctx, host)
	if err != nil {
		return nil, err
	}
	var last error
	for _, addr := range addrs {
		conn, err := p.dial(ctx, network, net.JoinHostPort(addr.Unmap().String(), portText))
		if err == nil {
			return conn, nil
		}
		last = err
	}
	return nil, fmt.Errorf("connect failed: %w", last)
}

var blockedPrefixes = mustPrefixes(
	"0.0.0.0/8", "10.0.0.0/8", "100.64.0.0/10", "127.0.0.0/8", "169.254.0.0/16",
	"172.16.0.0/12", "192.0.0.0/24", "192.0.2.0/24", "192.88.99.0/24", "192.168.0.0/16",
	"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",
	"::/128", "::1/128", "::ffff:0:0/96", "64:ff9b::/96", "64:ff9b:1::/48", "100::/64",
	"2001::/32", "2001:db8::/32", "2002::/16", "fc00::/7", "fe80::/10", "ff00::/8",
)

func mustPrefixes(list ...string) []netip.Prefix {
	out := make([]netip.Prefix, len(list))
	for i, item := range list {
		out[i] = netip.MustParsePrefix(item)
	}
	return out
}

func IsPublicAddr(addr netip.Addr) bool {
	addr = addr.Unmap()
	if !addr.IsValid() || addr.Zone() != "" {
		return false
	}
	if addr.IsUnspecified() || addr.IsLoopback() || addr.IsPrivate() || addr.IsMulticast() ||
		addr.IsLinkLocalUnicast() || addr.IsLinkLocalMulticast() || addr.IsInterfaceLocalMulticast() ||
		!addr.IsGlobalUnicast() {
		return false
	}
	for _, prefix := range blockedPrefixes {
		if prefix.Contains(addr) {
			return false
		}
	}
	return true
}