package middleware

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

	chimw "github.com/go-chi/chi/v5/middleware"
)

func TestLocalOnlyUsesOriginalTCPPeer(t *testing.T) {
	allowed := http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	})
	// 与生产顺序一致：先保存 TCP 对端，再由 RealIP 处理日志所用地址，最后做本机限制。
	chain := CapturePeerAddress(chimw.RealIP(LocalOnly(allowed)))

	cases := []struct {
		name       string
		remoteAddr string
		forwarded  string
		want       int
	}{
		{name: "IPv4 回环允许", remoteAddr: "127.0.0.1:41000", want: http.StatusNoContent},
		{name: "IPv6 回环允许", remoteAddr: "[::1]:41000", want: http.StatusNoContent},
		{name: "局网地址拒绝", remoteAddr: "192.168.2.106:41000", want: http.StatusForbidden},
		{
			name:       "伪造 X-Forwarded-For 仍拒绝",
			remoteAddr: "192.168.2.106:41000",
			forwarded:  "127.0.0.1",
			want:       http.StatusForbidden,
		},
		{name: "非法地址拒绝", remoteAddr: "localhost", want: http.StatusForbidden},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			r := httptest.NewRequest(http.MethodPost, "/api/v1/setup/admin", nil)
			r.RemoteAddr = tc.remoteAddr
			if tc.forwarded != "" {
				r.Header.Set("X-Forwarded-For", tc.forwarded)
			}
			w := httptest.NewRecorder()
			chain.ServeHTTP(w, r)
			if w.Code != tc.want {
				t.Fatalf("状态码 = %d，期望 %d；响应：%s", w.Code, tc.want, w.Body.String())
			}
		})
	}
}

func TestLocalOnlyFailsClosedWithoutCapturedPeer(t *testing.T) {
	h := LocalOnly(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNoContent)
	}))
	r := httptest.NewRequest(http.MethodPost, "/api/v1/setup/admin", nil)
	r.RemoteAddr = "127.0.0.1:41000"
	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)
	if w.Code != http.StatusForbidden {
		t.Fatalf("未安装 CapturePeerAddress 时必须拒绝，实际状态码 %d", w.Code)
	}
}