package proxy

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

	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt/v5"
	"github.com/stretchr/testify/require"
)

func TestSubjectAndExpiry(t *testing.T) {
	const did = "did:plc:abcdefghijklmnopqrstuvwx"
	t.Setenv("MCP_ALLOWED_SUBJECT", did)
	key, pub, err := generateRSAKeyPair()
	require.NoError(t, err)
	p, err := NewProxyRouter("https://example.com", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(204) }), pub, nil, false, false, nil, "/userinfo")
	require.NoError(t, err)
	router := gin.New()
	p.SetupRoutes(router)
	for _, tc := range []struct {
		name, sub, aud string
		exp            any
		status         int
	}{
		{"owner", did, "https://example.com/mcp", time.Now().Add(time.Hour).Unix(), 204},
		{"other DID", "did:plc:other", "https://example.com/mcp", time.Now().Add(time.Hour).Unix(), 401},
		{"missing subject", "", "https://example.com/mcp", time.Now().Add(time.Hour).Unix(), 401},
		{"expired", did, "https://example.com/mcp", time.Now().Add(-time.Hour).Unix(), 401},
		{"no expiry", did, "https://example.com/mcp", nil, 401},
		{"issuer audience", did, "https://example.com", time.Now().Add(time.Hour).Unix(), 401},
	} {
		t.Run(tc.name, func(t *testing.T) {
			claims := jwt.MapClaims{"iss": "https://example.com", "aud": tc.aud, "sub": tc.sub}
			if tc.exp != nil {
				claims["exp"] = tc.exp
			}
			token, err := jwt.NewWithClaims(jwt.SigningMethodRS256, claims).SignedString(key)
			require.NoError(t, err)
			req := httptest.NewRequest("POST", "/mcp", nil)
			req.Header.Set("Authorization", "Bearer "+token)
			w := httptest.NewRecorder()
			router.ServeHTTP(w, req)
			require.Equal(t, tc.status, w.Code)
			if tc.status == 401 {
				require.Equal(t, `Bearer resource_metadata="https://example.com/.well-known/oauth-protected-resource"`, w.Header().Get("WWW-Authenticate"))
			}
		})
	}
	w := httptest.NewRecorder()
	router.ServeHTTP(w, httptest.NewRequest("POST", "/mcp", nil))
	require.Equal(t, 401, w.Code)
	require.NotEmpty(t, w.Header().Get("WWW-Authenticate"))
}