package httpx

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

func prefixes(t *testing.T, cidrs ...string) []netip.Prefix {
	t.Helper()
	var out []netip.Prefix
	for _, c := range cidrs {
		p, err := netip.ParsePrefix(c)
		if err != nil {
			t.Fatal(err)
		}
		out = append(out, p)
	}
	return out
}

func TestResolveClient(t *testing.T) {
	trusted := "172.16.0.0/12,127.0.0.0/8"
	tests := []struct {
		name       string
		remoteAddr string
		cidrs      string
		hdr        map[string][]string
		wantIP     string
		wantProto  string
		wantHost   string // empty = request Host
	}{
		{
			name:       "no proxies configured: headers ignored",
			remoteAddr: "203.0.113.7:4444",
			cidrs:      "",
			hdr:        map[string][]string{"X-Forwarded-For": {"198.51.100.1"}, "X-Forwarded-Proto": {"https"}},
			wantIP:     "203.0.113.7",
			wantProto:  "http",
		},
		{
			name:       "untrusted peer spoofing headers: ignored",
			remoteAddr: "203.0.113.7:4444",
			cidrs:      trusted,
			hdr:        map[string][]string{"X-Forwarded-For": {"198.51.100.1"}, "X-Forwarded-Proto": {"https"}, "X-Forwarded-Host": {"evil.example"}},
			wantIP:     "203.0.113.7",
			wantProto:  "http",
		},
		{
			name:       "trusted peer, single hop",
			remoteAddr: "172.18.0.2:1234",
			cidrs:      trusted,
			hdr:        map[string][]string{"X-Forwarded-For": {"198.51.100.1"}, "X-Forwarded-Proto": {"https"}, "X-Forwarded-Host": {"bounty.example.com"}},
			wantIP:     "198.51.100.1",
			wantProto:  "https",
			wantHost:   "bounty.example.com",
		},
		{
			name:       "client-spoofed XFF prefix is skipped (rightmost untrusted wins)",
			remoteAddr: "172.18.0.2:1234",
			cidrs:      trusted,
			hdr:        map[string][]string{"X-Forwarded-For": {"6.6.6.6, 198.51.100.1"}},
			wantIP:     "198.51.100.1",
			wantProto:  "http",
		},
		{
			name:       "chain of trusted proxies collapses to real client",
			remoteAddr: "172.18.0.2:1234",
			cidrs:      trusted,
			hdr:        map[string][]string{"X-Forwarded-For": {"198.51.100.1, 172.20.0.5, 172.20.0.9"}},
			wantIP:     "198.51.100.1",
			wantProto:  "http",
		},
		{
			name:       "all-trusted chain falls back to leftmost",
			remoteAddr: "127.0.0.1:9999",
			cidrs:      trusted,
			hdr:        map[string][]string{"X-Forwarded-For": {"172.20.0.5"}},
			wantIP:     "172.20.0.5",
			wantProto:  "http",
		},
		{
			name:       "trusted peer without forwarded headers",
			remoteAddr: "172.18.0.2:1234",
			cidrs:      trusted,
			hdr:        nil,
			wantIP:     "172.18.0.2",
			wantProto:  "http",
		},
		{
			name:       "malformed XFF entry stops the walk",
			remoteAddr: "172.18.0.2:1234",
			cidrs:      trusted,
			hdr:        map[string][]string{"X-Forwarded-For": {"garbage, 198.51.100.1"}},
			wantIP:     "198.51.100.1",
			wantProto:  "http",
		},
		{
			name:       "multiple XFF headers are concatenated in order",
			remoteAddr: "172.18.0.2:1234",
			cidrs:      trusted,
			hdr:        map[string][]string{"X-Forwarded-For": {"198.51.100.1", "172.20.0.5"}},
			wantIP:     "198.51.100.1",
			wantProto:  "http",
		},
		{
			name:       "invalid forwarded proto ignored",
			remoteAddr: "172.18.0.2:1234",
			cidrs:      trusted,
			hdr:        map[string][]string{"X-Forwarded-For": {"198.51.100.1"}, "X-Forwarded-Proto": {"gopher"}},
			wantIP:     "198.51.100.1",
			wantProto:  "http",
		},
		{
			name:       "ipv6 client through trusted proxy",
			remoteAddr: "172.18.0.2:1234",
			cidrs:      trusted,
			hdr:        map[string][]string{"X-Forwarded-For": {"2001:db8::1"}},
			wantIP:     "2001:db8::1",
			wantProto:  "http",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			r := httptest.NewRequest(http.MethodGet, "http://app.local/", nil)
			r.RemoteAddr = tt.remoteAddr
			for k, vs := range tt.hdr {
				for _, v := range vs {
					r.Header.Add(k, v)
				}
			}
			var cidrs []netip.Prefix
			if tt.cidrs != "" {
				cidrs = prefixes(t, "172.16.0.0/12", "127.0.0.0/8")
			}
			got := resolveClient(r, cidrs, "http")
			if got.IP.String() != tt.wantIP {
				t.Errorf("IP = %s, want %s", got.IP, tt.wantIP)
			}
			if got.Proto != tt.wantProto {
				t.Errorf("Proto = %s, want %s", got.Proto, tt.wantProto)
			}
			wantHost := tt.wantHost
			if wantHost == "" {
				wantHost = "app.local"
			}
			if got.Host != wantHost {
				t.Errorf("Host = %s, want %s", got.Host, wantHost)
			}
		})
	}
}