package security

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

func TestEnrichClaims(t *testing.T) {
	r := httptest.NewRequest("POST", "/login", nil)
	r.RemoteAddr = "10.0.0.1:5555"
	r.Header.Set("X-Forwarded-For", "1.2.3.4, 10.0.0.1")
	r.Header.Set("Authorization", "Bearer old-token")
	r.Header.Set("User-Agent", "ua/1")
	r.Header.Set("Origin", "https://app.example")
	r.Host = "api.example"
	r = r.WithContext(context.WithValue(r.Context(), http.LocalAddrContextKey, &net.TCPAddr{IP: net.ParseIP("192.168.1.5"), Port: 8080}))
	ctx := WithRequestInfo(r.Context(), r)

	in := map[string]any{ClaimUserAgent: "x"}
	got := enrichClaims(ctx, in)
	if got[ClaimIPAddress] != "1.2.3.4" || got[ClaimPreviousSession] != "old-token" || got[ClaimUserAgent] != "x" ||
		got[ClaimServerIPAddress] != "192.168.1.5" || got[ClaimHost] != "api.example" || got[ClaimOrigin] != "https://app.example" {
		t.Fatalf("unexpected claims: %v", got)
	}
	if _, ok := in[ClaimIPAddress]; ok {
		t.Fatal("input map mutated")
	}
	if got := enrichClaims(ctx, map[string]any{ClaimIPAddress: "9.9.9.9"}); got[ClaimIPAddress] != "9.9.9.9" {
		t.Fatalf("caller value overridden: %v", got)
	}
	if got := enrichClaims(r.Context(), nil); got != nil {
		t.Fatalf("expected nil without info, got %v", got)
	}
}

func TestClientIPTrustedProxies(t *testing.T) {
	mk := func(peer, xff string) *http.Request {
		r := httptest.NewRequest("GET", "/", nil)
		r.RemoteAddr = peer
		if xff != "" {
			r.Header.Set("X-Forwarded-For", xff)
		}
		return r
	}
	o := RequestInfoOptions{TrustedProxies: []string{"10.0.0.0/8", "172.16.0.1"}}
	if got := clientIP(mk("10.0.0.1:1", "6.6.6.6, 1.2.3.4, 10.0.0.2"), o); got != "1.2.3.4" {
		t.Fatalf("spoofed prefix not skipped: %s", got)
	}
	if got := clientIP(mk("8.8.8.8:1", "1.2.3.4"), o); got != "8.8.8.8" {
		t.Fatalf("untrusted peer header honoured: %s", got)
	}
	if got := clientIP(mk("10.0.0.1:1", "1.2.3.4"), RequestInfoOptions{IgnoreProxyHeaders: true}); got != "10.0.0.1" {
		t.Fatalf("IgnoreProxyHeaders: %s", got)
	}
}

func TestRequestInfoExtras(t *testing.T) {
	r := httptest.NewRequest("POST", "/login", nil)
	r.RemoteAddr = "192.168.0.7:4242"
	r.Header.Set("X-Device-ID", "dev-1")
	r.Header.Set("X-Client-ID", "app-9")
	r = r.WithContext(context.WithValue(r.Context(), http.LocalAddrContextKey, &net.TCPAddr{IP: net.ParseIP("10.1.1.1"), Port: 8443}))
	got := enrichClaims(WithRequestInfo(r.Context(), r), nil)
	if got[ClaimClientPort] != "4242" || got[ClaimServerPort] != "8443" || got[ClaimDeviceID] != "dev-1" || got[ClaimClientID] != "app-9" || got[ClaimIsPrivateIP] != true {
		t.Fatalf("unexpected claims: %v", got)
	}
}