// Package fastlyxff provides Fastly CDN IP range management for trusted proxy handling.
//
// This package parses Fastly's public IP ranges JSON file and provides middleware
// for both Echo framework and standard net/http for proper client IP extraction
// from X-Forwarded-For headers. It's designed specifically for services deployed
// behind Fastly's CDN that need to identify real client IPs for logging, rate
// limiting, and security purposes.
//
// Fastly publishes their edge server IP ranges in a JSON format that this package
// consumes to automatically configure trusted proxy ranges. This ensures that
// X-Forwarded-For headers are only trusted when they originate from legitimate
// Fastly edge servers.
//
// Key features:
//   - Automatic parsing of Fastly's IP ranges JSON format
//   - Support for both IPv4 and IPv6 address ranges
//   - Echo framework integration via TrustOption generation
//   - Standard net/http middleware support
//   - CIDR notation parsing and validation
//
// # Echo Framework Usage
//
//	fastlyRanges, err := fastlyxff.New("fastly.json")
//	if err != nil {
//		return err
//	}
//	options, err := fastlyRanges.EchoTrustOption()
//	if err != nil {
//		return err
//	}
//	e.IPExtractor = echo.ExtractIPFromXFFHeader(options...)
//
// # Net/HTTP Usage
//
//	fastlyRanges, err := fastlyxff.New("fastly.json")
//	if err != nil {
//		return err
//	}
//	middleware := fastlyRanges.HTTPMiddleware()
//
//	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
//		// Both methods work - middleware updates r.RemoteAddr (with port 0) and stores in context
//		realIP := fastlyxff.GetRealIP(r)  // Preferred method
//		// OR: host, _, _ := net.SplitHostPort(r.RemoteAddr)  // Direct access (port will be "0")
//		fmt.Fprintf(w, "Real IP: %s\n", realIP)
//	})
//
//	http.ListenAndServe(":8080", middleware(handler))
//
// # Net/HTTP with Additional Trusted Ranges
//
//	fastlyRanges, err := fastlyxff.New("fastly.json")
//	if err != nil {
//		return err
//	}
//
//	// Add custom trusted CIDRs (e.g., internal load balancers)
//	// Note: For Echo framework, use the ekko package for additional ranges
//	err = fastlyRanges.AddTrustedCIDR("10.0.0.0/8")
//	if err != nil {
//		return err
//	}
//
//	middleware := fastlyRanges.HTTPMiddleware()
//	handler := middleware(yourHandler)
//
// The JSON file typically contains IP ranges in this format:
//
//	{
//	  "addresses": ["23.235.32.0/20", "43.249.72.0/22", ...],
//	  "ipv6_addresses": ["2a04:4e40::/32", "2a04:4e42::/32", ...]
//	}
package fastlyxff

import (
	"context"
	"encoding/json"
	"net"
	"net/http"
	"net/netip"
	"os"
	"strings"

	"github.com/labstack/echo/v4"
)

// FastlyXFF represents Fastly's published IP ranges for their CDN edge servers.
// This structure matches the JSON format provided by Fastly for their public IP ranges.
// It contains separate lists for IPv4 and IPv6 CIDR ranges, plus additional trusted CIDRs.
type FastlyXFF struct {
	IPv4       []string `json:"addresses"`      // IPv4 CIDR ranges (e.g., "23.235.32.0/20")
	IPv6       []string `json:"ipv6_addresses"` // IPv6 CIDR ranges (e.g., "2a04:4e40::/32")
	extraCIDRs []string // Additional trusted CIDRs added via AddTrustedCIDR
}

// TrustedNets holds parsed network prefixes for efficient IP range checking.
type TrustedNets struct {
	prefixes []netip.Prefix // Parsed network prefixes for efficient lookups
}

// contextKey is used for storing the real client IP in request context
type contextKey string

const realIPKey contextKey = "fastly-real-ip"

// New loads and parses Fastly IP ranges from a JSON file.
// The file should contain Fastly's published IP ranges in their standard JSON format.
//
// Parameters:
//   - fileName: Path to the Fastly IP ranges JSON file
//
// Returns the parsed FastlyXFF structure or an error if the file cannot be
// read or the JSON format is invalid.
func New(fileName string) (*FastlyXFF, error) {
	b, err := os.ReadFile(fileName)
	if err != nil {
		return nil, err
	}

	d := FastlyXFF{}

	err = json.Unmarshal(b, &d)
	if err != nil {
		return nil, err
	}

	return &d, nil
}

// EchoTrustOption converts Fastly IP ranges into Echo framework trust options.
// This method generates trust configurations that tell Echo to accept X-Forwarded-For
// headers only from Fastly's edge servers, ensuring accurate client IP extraction.
//
// The generated trust options should be used with Echo's IP extractor:
//
//	options, err := fastlyRanges.EchoTrustOption()
//	if err != nil {
//		return err
//	}
//	e.IPExtractor = echo.ExtractIPFromXFFHeader(options...)
//
// Returns a slice of Echo trust options or an error if any CIDR range cannot be parsed.
func (xff *FastlyXFF) EchoTrustOption() ([]echo.TrustOption, error) {
	ranges := []echo.TrustOption{}

	for _, s := range append(xff.IPv4, xff.IPv6...) {
		_, cidr, err := net.ParseCIDR(s)
		if err != nil {
			return nil, err
		}

		trust := echo.TrustIPRange(cidr)
		ranges = append(ranges, trust)
	}

	return ranges, nil
}

// AddTrustedCIDR adds an additional CIDR to the list of trusted proxies.
// This allows trusting proxies beyond Fastly's published ranges.
// The cidr parameter must be a valid CIDR notation (e.g., "10.0.0.0/8", "192.168.1.0/24").
// Returns an error if the CIDR format is invalid.
func (xff *FastlyXFF) AddTrustedCIDR(cidr string) error {
	// Validate CIDR format
	_, _, err := net.ParseCIDR(cidr)
	if err != nil {
		return err
	}

	// Add to extra CIDRs
	xff.extraCIDRs = append(xff.extraCIDRs, cidr)
	return nil
}

// isTrustedProxy checks if the given IP address belongs to Fastly's trusted IP ranges
// or any additional CIDRs added via AddTrustedCIDR.
func (xff *FastlyXFF) isTrustedProxy(ip string) bool {
	addr, err := netip.ParseAddr(ip)
	if err != nil {
		return false
	}

	// Check all IPv4 and IPv6 ranges (Fastly + additional)
	allRanges := append(append(xff.IPv4, xff.IPv6...), xff.extraCIDRs...)
	for _, s := range allRanges {
		_, cidr, err := net.ParseCIDR(s)
		if err != nil {
			continue
		}
		if cidr.Contains(net.IP(addr.AsSlice())) {
			return true
		}
	}
	return false
}

// extractRealIP extracts the real client IP from X-Forwarded-For header.
// It returns the rightmost IP that is not from a trusted Fastly proxy.
func (xff *FastlyXFF) extractRealIP(r *http.Request) string {
	// Get the immediate peer IP
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}

	// If the immediate peer is not a trusted Fastly proxy, return it
	if !xff.isTrustedProxy(host) {
		return host
	}

	// Check X-Forwarded-For header
	xff_header := r.Header.Get("X-Forwarded-For")
	if xff_header == "" {
		return host
	}

	// Parse comma-separated IP list
	ips := strings.Split(xff_header, ",")
	if len(ips) == 0 {
		return host
	}

	// Find the leftmost IP that is not from a trusted proxy
	// This represents the original client IP
	for i := 0; i < len(ips); i++ {
		ip := strings.TrimSpace(ips[i])
		if ip != "" && !xff.isTrustedProxy(ip) {
			return ip
		}
	}

	// Fallback to the immediate peer
	return host
}

// HTTPMiddleware returns a net/http middleware that extracts real client IP
// from X-Forwarded-For headers when the request comes from trusted Fastly proxies.
// The real IP is stored in the request context and also updates r.RemoteAddr
// with port 0 (since the original port is from the proxy, not the real client).
func (xff *FastlyXFF) HTTPMiddleware() func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			realIP := xff.extractRealIP(r)

			// Store in context for GetRealIP function
			ctx := context.WithValue(r.Context(), realIPKey, realIP)
			r = r.WithContext(ctx)

			// Update RemoteAddr to be consistent with extracted IP
			// Use port 0 since the original port is from the proxy, not the real client
			r.RemoteAddr = net.JoinHostPort(realIP, "0")

			next.ServeHTTP(w, r)
		})
	}
}

// GetRealIP retrieves the real client IP from the request context.
// This should be used after the HTTPMiddleware has processed the request.
// Returns the remote address if no real IP was extracted.
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
}