package middleware

import (
	"net/netip"

	"github.com/gin-gonic/gin"
)

// AnonymizeIPMiddleware masks the last segment of IPv4 addresses
// and the last 80 bits of IPv6 addresses.
func AnonymizeIPMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		anonymizedIP := anonymizeIP(ip)
		c.Request.RemoteAddr = anonymizedIP
		c.Next()
	}
}

// anonymizeIP masks the last segment of IPv4 addresses
// and the last 80 bits of IPv6 addresses.
func anonymizeIP(ip string) string {
	parsedIP, err := netip.ParseAddr(ip)
	if err != nil {
		// If there's an error parsing the IP, return the original IP
		return ip
	}

	if parsedIP.Is4() {
		// Mask the last octet for IPv4
		ipPrefix, err := parsedIP.Prefix(24)
		if err != nil {
			return ip
		}

		return ipPrefix.Masked().Addr().String()
	} else if parsedIP.Is6() {
		// Mask the last 80 bits for IPv6
		ipPrefix, err := parsedIP.Prefix(48)
		if err != nil {
			return ip
		}
		return ipPrefix.Masked().Addr().String()
	}

	// If it's neither IPv4 nor IPv6, return the original IP
	return ip
}