package httpx

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/netip"
	"strings"
)

// ClientIPResolver trusts X-Forwarded-For only from explicitly configured
// proxy addresses. It is immutable and safe to share between requests.
// The zero value (and a nil resolver) trusts no proxy.
type ClientIPResolver struct {
	trusted []netip.Prefix
}

// NewClientIPResolver accepts IP addresses or CIDR networks, never hostnames.
func NewClientIPResolver(proxies ...string) (*ClientIPResolver, error) {
	r := &ClientIPResolver{}
	for _, value := range proxies {
		value = strings.TrimSpace(value)
		if address, err := netip.ParseAddr(value); err == nil && address.Zone() == "" {
			address = address.Unmap()
			r.trusted = append(r.trusted, netip.PrefixFrom(address, address.BitLen()))
			continue
		}
		prefix, err := netip.ParsePrefix(value)
		if err != nil {
			return nil, fmt.Errorf("invalid trusted proxy %q: expected an IP address or CIDR", value)
		}
		if prefix.Addr().Is4In6() {
			if prefix.Bits() < 96 {
				return nil, fmt.Errorf("invalid trusted proxy %q: mapped IPv4 prefix must be at least /96", value)
			}
			prefix = netip.PrefixFrom(prefix.Addr().Unmap(), prefix.Bits()-96)
		}
		r.trusted = append(r.trusted, prefix.Masked())
	}
	return r, nil
}

func (r *ClientIPResolver) trusts(address netip.Addr) bool {
	if r != nil {
		for _, prefix := range r.trusted {
			if prefix.Contains(address) {
				return true
			}
		}
	}
	return false
}

// ClientIP walks X-Forwarded-For from right to left, stopping at the first
// untrusted address. Missing or malformed forwarding data falls back to the
// socket peer. X-Real-IP, Forwarded and provider-specific headers are ignored.
func (r *ClientIPResolver) ClientIP(request *http.Request) string {
	if request == nil {
		return ""
	}
	host := request.RemoteAddr
	if value, _, err := net.SplitHostPort(host); err == nil {
		host = value
	}
	peer, err := netip.ParseAddr(host)
	if err != nil {
		return ""
	}
	peer = peer.WithZone("").Unmap()
	if !r.trusts(peer) {
		return peer.String()
	}
	values := request.Header.Values("X-Forwarded-For")
	var chain []netip.Addr
	for _, value := range values {
		for _, part := range strings.Split(value, ",") {
			address, err := netip.ParseAddr(strings.TrimSpace(part))
			if err != nil || address.Zone() != "" {
				return peer.String()
			}
			chain = append(chain, address.Unmap())
		}
	}
	current := peer
	for i := len(chain) - 1; i >= 0 && r.trusts(current); i-- {
		current = chain[i]
	}
	return current.String()
}

type clientIPResolverKey struct{}

// WithClientIPResolver configures ClientIP for Context values derived from ctx.
// Adapters should set this before running application middleware.
func WithClientIPResolver(ctx context.Context, resolver *ClientIPResolver) context.Context {
	return context.WithValue(ctx, clientIPResolverKey{}, resolver)
}

// ClientIP returns the verified client address without a port. By default it
// returns the socket peer; App's trusted proxy configuration enables XFF.
func (c *Context) ClientIP() string {
	resolver, _ := c.Context().Value(clientIPResolverKey{}).(*ClientIPResolver)
	return resolver.ClientIP(c.request)
}