// Package xff provides trusted proxy handling and real client IP extraction
// from X-Forwarded-For headers.
//
// This package has no external dependencies — it uses only the Go standard library.
//
// The XFF extraction algorithm walks right-to-left through the X-Forwarded-For
// header, skipping trusted proxy IPs, and returns the first untrusted IP as the
// real client address. This follows the MDN-recommended approach for secure
// client IP extraction.
//
// # Usage with net/http middleware
//
//	tp, err := xff.NewFromCIDRs([]string{"10.0.0.0/8", "192.168.0.0/16"})
//	if err != nil {
//		return err
//	}
//	handler := tp.HTTPMiddleware()(yourHandler)
//
// # Direct extraction
//
//	realIP := tp.ExtractRealIP(r)
package xff

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

// TrustedProxies holds a set of trusted proxy network prefixes and provides
// methods for extracting the real client IP from X-Forwarded-For headers.
type TrustedProxies struct {
	prefixes []netip.Prefix
}

type contextKey string

const realIPKey contextKey = "xff-real-ip"

// New creates a TrustedProxies from already-parsed prefixes.
func New(prefixes ...netip.Prefix) *TrustedProxies {
	return &TrustedProxies{prefixes: prefixes}
}

// NewFromCIDRs creates a TrustedProxies from CIDR strings (e.g., "10.0.0.0/8").
func NewFromCIDRs(cidrs []string) (*TrustedProxies, error) {
	prefixes := make([]netip.Prefix, 0, len(cidrs))
	for _, s := range cidrs {
		p, err := netip.ParsePrefix(s)
		if err != nil {
			return nil, err
		}
		prefixes = append(prefixes, p)
	}
	return &TrustedProxies{prefixes: prefixes}, nil
}

// AddCIDR adds a CIDR string to the trusted proxy list.
func (tp *TrustedProxies) AddCIDR(cidr string) error {
	p, err := netip.ParsePrefix(cidr)
	if err != nil {
		return err
	}
	tp.prefixes = append(tp.prefixes, p)
	return nil
}

// AddPrefix adds a parsed prefix to the trusted proxy list.
func (tp *TrustedProxies) AddPrefix(prefix netip.Prefix) {
	tp.prefixes = append(tp.prefixes, prefix)
}

// Prefixes returns a copy of the trusted proxy prefixes.
func (tp *TrustedProxies) Prefixes() []netip.Prefix {
	out := make([]netip.Prefix, len(tp.prefixes))
	copy(out, tp.prefixes)
	return out
}

// IsTrusted reports whether ip belongs to any of the trusted proxy ranges.
func (tp *TrustedProxies) IsTrusted(ip string) bool {
	addr, err := netip.ParseAddr(ip)
	if err != nil {
		return false
	}
	return tp.isTrustedAddr(addr)
}

func (tp *TrustedProxies) isTrustedAddr(addr netip.Addr) bool {
	for _, p := range tp.prefixes {
		if p.Contains(addr) {
			return true
		}
	}
	return false
}

// ExtractRealIP extracts the real client IP from a request by walking the
// X-Forwarded-For header right-to-left, skipping trusted proxy IPs.
// If the immediate peer is not a trusted proxy, its IP is returned.
func (tp *TrustedProxies) ExtractRealIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}

	hostAddr, err := netip.ParseAddr(host)
	if err != nil || !tp.isTrustedAddr(hostAddr) {
		return host
	}

	xffHeader := r.Header.Get("X-Forwarded-For")
	if xffHeader == "" {
		return host
	}

	ips := strings.Split(xffHeader, ",")

	// Walk right-to-left: skip trusted proxies, return first untrusted IP.
	for i := len(ips) - 1; i >= 0; i-- {
		ip := strings.TrimSpace(ips[i])
		if ip == "" {
			continue
		}
		addr, err := netip.ParseAddr(ip)
		if err != nil || !tp.isTrustedAddr(addr) {
			return ip
		}
	}

	return host
}

// HTTPMiddleware returns a net/http middleware that extracts the real client IP
// from X-Forwarded-For headers and stores it in the request context and
// RemoteAddr. The port in RemoteAddr is set to 0 because the original port
// belongs to the proxy connection, not the real client.
func (tp *TrustedProxies) HTTPMiddleware() func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			realIP := tp.ExtractRealIP(r)
			ctx := context.WithValue(r.Context(), realIPKey, realIP)
			r = r.WithContext(ctx)
			r.RemoteAddr = net.JoinHostPort(realIP, "0")
			next.ServeHTTP(w, r)
		})
	}
}

// GetRealIP retrieves the real client IP from the request context.
// Returns the remote address host if no real IP was extracted by middleware.
func GetRealIP(r *http.Request) string {
	if ip, ok := r.Context().Value(realIPKey).(string); ok {
		return ip
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}