package middleware

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

func TestAnonymizeIPMiddlewareIPv4(t *testing.T) {
	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(AnonymizeIPMiddleware())

	r.GET("/test", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	req, _ := http.NewRequest(http.MethodGet, "/test", nil)
	req.RemoteAddr = "192.168.1.1:1234"

	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)
	assert.Equal(t, "192.168.1.0", req.RemoteAddr)
}

func TestAnonymizeIPMiddlewareIPv6(t *testing.T) {
	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(AnonymizeIPMiddleware())

	r.GET("/test", func(c *gin.Context) {
		c.String(http.StatusOK, "OK")
	})

	req, _ := http.NewRequest(http.MethodGet, "/test", nil)
	req.RemoteAddr = "[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:1234"

	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)

	assert.Equal(t, http.StatusOK, w.Code)
	assert.Equal(t, "2001:db8:85a3::", req.RemoteAddr)
}

func TestAnonymizeIP(t *testing.T) {
	tests := []struct {
		name     string
		ip       string
		expected string
	}{
		{"IPv4", "192.168.1.1", "192.168.1.0"},
		{"IPv6", "2001:0db8:85a3:0000:0000:8a2e:0370:7334", "2001:db8:85a3::"},
		{"InvalidIP", "invalid_ip", "invalid_ip"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := anonymizeIP(tt.ip)
			assert.Equal(t, tt.expected, result)
		})
	}
}