package guard

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

const (
	maxURLLength = 4096
	maxRedirects = 5
)

var (
	ErrURLTooLong        = errors.New("URL is too long")
	ErrUnsupportedScheme = errors.New("only http and https URLs are allowed")
	ErrMissingHost       = errors.New("URL must include a host")
	ErrBlockedAddress    = errors.New("destination address is not allowed")
	ErrBlockedPort       = errors.New("destination port is not allowed")
	ErrTooManyRedirects  = errors.New("too many redirects")
)

type Policy struct {
	mutex        sync.RWMutex
	allowedPorts map[uint16]bool
}

func NewPolicy(allowedPorts []uint16) *Policy {
	policy := &Policy{}
	policy.SetAllowedPorts(allowedPorts)
	return policy
}

func (policy *Policy) SetAllowedPorts(allowedPorts []uint16) {
	ports := make(map[uint16]bool, len(allowedPorts))
	for _, port := range allowedPorts {
		ports[port] = true
	}
	policy.mutex.Lock()
	defer policy.mutex.Unlock()
	policy.allowedPorts = ports
}

func (policy *Policy) ValidateURL(raw string) (*url.URL, error) {
	raw = strings.TrimSpace(raw)
	if len(raw) > maxURLLength {
		return nil, ErrURLTooLong
	}

	parsed, err := url.Parse(raw)
	if err != nil {
		return nil, fmt.Errorf("invalid URL: %w", err)
	}
	if parsed.Scheme != "http" && parsed.Scheme != "https" {
		return nil, ErrUnsupportedScheme
	}

	host := parsed.Hostname()
	if host == "" {
		return nil, ErrMissingHost
	}
	if address, err := netip.ParseAddr(host); err == nil && !IsPublicAddress(address) {
		return nil, ErrBlockedAddress
	}

	port, err := portOf(parsed)
	if err != nil {
		return nil, err
	}
	if !policy.portAllowed(port) {
		return nil, ErrBlockedPort
	}
	return parsed, nil
}

func (policy *Policy) portAllowed(port uint16) bool {
	policy.mutex.RLock()
	defer policy.mutex.RUnlock()
	return len(policy.allowedPorts) == 0 || policy.allowedPorts[port]
}

func portOf(parsed *url.URL) (uint16, error) {
	if raw := parsed.Port(); raw != "" {
		value, err := strconv.ParseUint(raw, 10, 16)
		if err != nil {
			return 0, fmt.Errorf("invalid port: %w", err)
		}
		return uint16(value), nil
	}
	if parsed.Scheme == "https" {
		return 443, nil
	}
	return 80, nil
}