package proxy

import (
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

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

func generateRSAKeyPair() (*rsa.PrivateKey, *rsa.PublicKey, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, err
	}
	return privateKey, &privateKey.PublicKey, nil
}

func createJWT(privateKey *rsa.PrivateKey, claims jwt.MapClaims) (string, error) {
	if _, ok := claims["iss"]; !ok {
		claims["iss"] = "https://example.com"
	}
	if _, ok := claims["aud"]; !ok {
		claims["aud"] = "https://example.com/mcp"
	}
	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	return token.SignedString(privateKey)
}

func createDummyBackendServer() *httptest.Server {
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, `{"message": "Hello from backend", "method": "%s", "path": "%s"}`, r.Method, r.URL.Path)
	}))
}

func TestProxyRouter_RejectsWrongIssuerOrAudience(t *testing.T) {
	privateKey, publicKey, err := generateRSAKeyPair()
	require.NoError(t, err)

	proxyRouter, err := NewProxyRouter("https://example.com", http.NotFoundHandler(), publicKey, http.Header{}, false, false, nil, "/userinfo")
	require.NoError(t, err)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	proxyRouter.SetupRoutes(router)

	cases := []struct {
		name   string
		claims jwt.MapClaims
	}{
		{
			name: "wrong issuer",
			claims: jwt.MapClaims{
				"sub": "test-user",
				"iss": "https://issuer.example.com",
				"aud": "https://example.com/mcp",
				"exp": time.Now().Add(time.Hour).Unix(),
			},
		},
		{
			name: "wrong audience",
			claims: jwt.MapClaims{
				"sub": "test-user",
				"iss": "https://example.com",
				"aud": "https://other.example.com",
				"exp": time.Now().Add(time.Hour).Unix(),
			},
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			token, err := createJWT(privateKey, tt.claims)
			require.NoError(t, err)

			req, err := http.NewRequest(http.MethodGet, "/test-endpoint", nil)
			require.NoError(t, err)
			req.Header.Set("Authorization", "Bearer "+token)

			w := httptest.NewRecorder()
			router.ServeHTTP(w, req)
			assert.Equal(t, http.StatusUnauthorized, w.Code)
		})
	}
}

func TestProxyRouter_HandleProxy_ValidToken(t *testing.T) {
	privateKey, publicKey, err := generateRSAKeyPair()
	require.NoError(t, err)

	backendServer := createDummyBackendServer()
	defer backendServer.Close()

	proxyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		resp, err := http.Get(backendServer.URL + r.URL.Path)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		defer resp.Body.Close()

		w.Header().Set("Content-Type", resp.Header.Get("Content-Type"))
		w.WriteHeader(resp.StatusCode)

		buf := make([]byte, 1024)
		for {
			n, err := resp.Body.Read(buf)
			if n > 0 {
				w.Write(buf[:n])
			}
			if err != nil {
				break
			}
		}
	})

	proxyHeaders := make(http.Header)
	proxyHeaders.Set("X-Forwarded-By", "mcp-auth-proxy")

	proxyRouter, err := NewProxyRouter("https://example.com", proxyHandler, publicKey, proxyHeaders, false, false, nil, "/userinfo")
	require.NoError(t, err)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	proxyRouter.SetupRoutes(router)

	claims := jwt.MapClaims{
		"sub": "test-user",
		"exp": time.Now().Add(time.Hour).Unix(),
		"iat": time.Now().Unix(),
	}

	token, err := createJWT(privateKey, claims)
	require.NoError(t, err)

	req, err := http.NewRequest("GET", "/test-endpoint", nil)
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer "+token)

	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	assert.Equal(t, http.StatusOK, w.Code)

	invalidToken := "invalid"
	req, err = http.NewRequest("GET", "/test-endpoint", nil)
	require.NoError(t, err)
	req.Header.Set("Authorization", "Bearer "+invalidToken)

	w = httptest.NewRecorder()
	router.ServeHTTP(w, req)
	assert.Equal(t, http.StatusUnauthorized, w.Code)
}

func TestProxyRouter_HeaderMapping(t *testing.T) {
	privateKey, publicKey, err := generateRSAKeyPair()
	require.NoError(t, err)

	cases := []struct {
		name            string
		headerMapping   map[string]string
		userinfo        map[string]any
		expectedHeaders map[string]string
	}{
		{
			name:          "string field",
			headerMapping: map[string]string{"/email": "X-Forwarded-Email"},
			userinfo:      map[string]any{"email": "user@example.com"},
			expectedHeaders: map[string]string{
				"X-Forwarded-Email": "user@example.com",
			},
		},
		{
			name:          "array field joined with comma",
			headerMapping: map[string]string{"/groups": "X-Forwarded-Groups"},
			userinfo:      map[string]any{"groups": []any{"admin", "users"}},
			expectedHeaders: map[string]string{
				"X-Forwarded-Groups": "admin,users",
			},
		},
		{
			name:          "multiple mappings",
			headerMapping: map[string]string{"/email": "X-Forwarded-Email", "/preferred_username": "X-Forwarded-User"},
			userinfo:      map[string]any{"email": "user@example.com", "preferred_username": "john"},
			expectedHeaders: map[string]string{
				"X-Forwarded-Email": "user@example.com",
				"X-Forwarded-User":  "john",
			},
		},
		{
			name:            "missing field is skipped",
			headerMapping:   map[string]string{"/email": "X-Forwarded-Email", "/missing": "X-Missing"},
			userinfo:        map[string]any{"email": "user@example.com"},
			expectedHeaders: map[string]string{"X-Forwarded-Email": "user@example.com"},
		},
		{
			name:            "nil headerMapping",
			headerMapping:   nil,
			userinfo:        map[string]any{"email": "user@example.com"},
			expectedHeaders: map[string]string{},
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			receivedHeaders := http.Header{}
			proxyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				for k, v := range r.Header {
					receivedHeaders[k] = v
				}
				w.WriteHeader(http.StatusOK)
			})

			proxyRouter, err := NewProxyRouter("https://example.com", proxyHandler, publicKey, http.Header{}, false, false, tt.headerMapping, "/userinfo")
			require.NoError(t, err)

			gin.SetMode(gin.TestMode)
			router := gin.New()
			proxyRouter.SetupRoutes(router)

			claims := jwt.MapClaims{
				"sub": "test-user",
				"exp": time.Now().Add(time.Hour).Unix(),
				"iat": time.Now().Unix(),
			}
			if tt.userinfo != nil {
				claims["userinfo"] = tt.userinfo
			}

			token, err := createJWT(privateKey, claims)
			require.NoError(t, err)

			req, err := http.NewRequest("GET", "/test", nil)
			require.NoError(t, err)
			req.Header.Set("Authorization", "Bearer "+token)

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

			assert.Equal(t, http.StatusOK, w.Code)
			for header, expected := range tt.expectedHeaders {
				assert.Equal(t, expected, receivedHeaders.Get(header), "header %s mismatch", header)
			}
		})
	}
}

func TestProxyRouter_HeaderMappingBase(t *testing.T) {
	privateKey, publicKey, err := generateRSAKeyPair()
	require.NoError(t, err)

	cases := []struct {
		name              string
		headerMapping     map[string]string
		headerMappingBase string
		claims            jwt.MapClaims
		expectedHeaders   map[string]string
		missingHeaders    []string
	}{
		{
			name:              "base=/ reads top-level claims",
			headerMapping:     map[string]string{"/email": "X-Forwarded-Email"},
			headerMappingBase: "/",
			claims: jwt.MapClaims{
				"sub":   "test-user",
				"email": "user@example.com",
				"exp":   time.Now().Add(time.Hour).Unix(),
				"iat":   time.Now().Unix(),
			},
			expectedHeaders: map[string]string{
				"X-Forwarded-Email": "user@example.com",
			},
		},
		{
			name:              "base=/userinfo reads userinfo claims",
			headerMapping:     map[string]string{"/email": "X-Forwarded-Email"},
			headerMappingBase: "/userinfo",
			claims: jwt.MapClaims{
				"sub":      "test-user",
				"email":    "toplevel@example.com",
				"userinfo": map[string]any{"email": "userinfo@example.com"},
				"exp":      time.Now().Add(time.Hour).Unix(),
				"iat":      time.Now().Unix(),
			},
			expectedHeaders: map[string]string{
				"X-Forwarded-Email": "userinfo@example.com",
			},
		},
		{
			name:              "base=/ with multiple claims",
			headerMapping:     map[string]string{"/email": "X-Forwarded-Email", "/name": "X-Forwarded-Name"},
			headerMappingBase: "/",
			claims: jwt.MapClaims{
				"sub":   "test-user",
				"email": "user@example.com",
				"name":  "John Doe",
				"exp":   time.Now().Add(time.Hour).Unix(),
				"iat":   time.Now().Unix(),
			},
			expectedHeaders: map[string]string{
				"X-Forwarded-Email": "user@example.com",
				"X-Forwarded-Name":  "John Doe",
			},
		},
		{
			name:              "base=/userinfo skips when userinfo is absent",
			headerMapping:     map[string]string{"/email": "X-Forwarded-Email"},
			headerMappingBase: "/userinfo",
			claims: jwt.MapClaims{
				"sub":   "test-user",
				"email": "user@example.com",
				"exp":   time.Now().Add(time.Hour).Unix(),
				"iat":   time.Now().Unix(),
			},
			missingHeaders: []string{"X-Forwarded-Email"},
		},
		{
			name:              "base=/ missing claim is skipped",
			headerMapping:     map[string]string{"/email": "X-Forwarded-Email", "/missing": "X-Missing"},
			headerMappingBase: "/",
			claims: jwt.MapClaims{
				"sub":   "test-user",
				"email": "user@example.com",
				"exp":   time.Now().Add(time.Hour).Unix(),
				"iat":   time.Now().Unix(),
			},
			expectedHeaders: map[string]string{
				"X-Forwarded-Email": "user@example.com",
			},
			missingHeaders: []string{"X-Missing"},
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			receivedHeaders := http.Header{}
			proxyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				for k, v := range r.Header {
					receivedHeaders[k] = v
				}
				w.WriteHeader(http.StatusOK)
			})

			proxyRouter, err := NewProxyRouter("https://example.com", proxyHandler, publicKey, http.Header{}, false, false, tt.headerMapping, tt.headerMappingBase)
			require.NoError(t, err)

			gin.SetMode(gin.TestMode)
			router := gin.New()
			proxyRouter.SetupRoutes(router)

			token, err := createJWT(privateKey, tt.claims)
			require.NoError(t, err)

			req, err := http.NewRequest("GET", "/test", nil)
			require.NoError(t, err)
			req.Header.Set("Authorization", "Bearer "+token)

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

			assert.Equal(t, http.StatusOK, w.Code)
			for header, expected := range tt.expectedHeaders {
				assert.Equal(t, expected, receivedHeaders.Get(header), "header %s mismatch", header)
			}
			for _, header := range tt.missingHeaders {
				assert.Empty(t, receivedHeaders.Get(header), "header %s should not be set", header)
			}
		})
	}
}

func TestProxyRouter_AuthorizationHeaderDefaultBehavior(t *testing.T) {
	privateKey, publicKey, err := generateRSAKeyPair()
	require.NoError(t, err)

	t.Run("strips authorization header by default", func(t *testing.T) {
		var backendAuthorization string
		proxyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			backendAuthorization = r.Header.Get("Authorization")
			w.WriteHeader(http.StatusOK)
		})

		proxyRouter, err := NewProxyRouter("https://example.com", proxyHandler, publicKey, http.Header{}, false, false, nil, "/userinfo")
		require.NoError(t, err)

		gin.SetMode(gin.TestMode)
		router := gin.New()
		proxyRouter.SetupRoutes(router)

		token, err := createJWT(privateKey, jwt.MapClaims{
			"sub": "user",
			"exp": time.Now().Add(time.Hour).Unix(),
			"iat": time.Now().Unix(),
		})
		require.NoError(t, err)

		req, err := http.NewRequest(http.MethodGet, "/mcp", nil)
		require.NoError(t, err)
		req.Header.Set("Authorization", "Bearer "+token)

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

		assert.Equal(t, http.StatusOK, w.Code)
		assert.Empty(t, backendAuthorization)
	})

	t.Run("forwards authorization header when enabled", func(t *testing.T) {
		var backendAuthorization string
		proxyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			backendAuthorization = r.Header.Get("Authorization")
			w.WriteHeader(http.StatusOK)
		})

		proxyRouter, err := NewProxyRouter("https://example.com", proxyHandler, publicKey, http.Header{}, false, true, nil, "/userinfo")
		require.NoError(t, err)

		gin.SetMode(gin.TestMode)
		router := gin.New()
		proxyRouter.SetupRoutes(router)

		token, err := createJWT(privateKey, jwt.MapClaims{
			"sub": "user",
			"exp": time.Now().Add(time.Hour).Unix(),
			"iat": time.Now().Unix(),
		})
		require.NoError(t, err)

		req, err := http.NewRequest(http.MethodGet, "/mcp", nil)
		require.NoError(t, err)
		req.Header.Set("Authorization", "Bearer "+token)

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

		assert.Equal(t, http.StatusOK, w.Code)
		assert.Equal(t, "Bearer "+token, backendAuthorization)
	})
}

func TestProxyRouter_HeaderMappingStripsClientHeaders(t *testing.T) {
	privateKey, publicKey, err := generateRSAKeyPair()
	require.NoError(t, err)

	cases := []struct {
		name              string
		headerMapping     map[string]string
		headerMappingBase string
		claims            jwt.MapClaims
		clientHeaders     map[string]string
		expectedHeaders   map[string]string
		missingHeaders    []string
	}{
		{
			name:              "claim missing strips client header",
			headerMapping:     map[string]string{"/groups": "X-Forwarded-Groups"},
			headerMappingBase: "/userinfo",
			claims: jwt.MapClaims{
				"sub":      "test-user",
				"userinfo": map[string]any{"email": "user@example.com"},
				"exp":      time.Now().Add(time.Hour).Unix(),
				"iat":      time.Now().Unix(),
			},
			clientHeaders:  map[string]string{"X-Forwarded-Groups": "admin"},
			missingHeaders: []string{"X-Forwarded-Groups"},
		},
		{
			name:              "claim present overwrites client header",
			headerMapping:     map[string]string{"/email": "X-Forwarded-Email"},
			headerMappingBase: "/userinfo",
			claims: jwt.MapClaims{
				"sub":      "test-user",
				"userinfo": map[string]any{"email": "user@example.com"},
				"exp":      time.Now().Add(time.Hour).Unix(),
				"iat":      time.Now().Unix(),
			},
			clientHeaders: map[string]string{"X-Forwarded-Email": "attacker@example.com"},
			expectedHeaders: map[string]string{
				"X-Forwarded-Email": "user@example.com",
			},
		},
		{
			name:              "base path absent strips client header",
			headerMapping:     map[string]string{"/email": "X-Forwarded-Email"},
			headerMappingBase: "/userinfo",
			claims: jwt.MapClaims{
				"sub": "test-user",
				"exp": time.Now().Add(time.Hour).Unix(),
				"iat": time.Now().Unix(),
			},
			clientHeaders:  map[string]string{"X-Forwarded-Email": "attacker@example.com"},
			missingHeaders: []string{"X-Forwarded-Email"},
		},
		{
			name:              "mixed mappings only set claims that exist",
			headerMapping:     map[string]string{"/email": "X-Forwarded-Email", "/groups": "X-Forwarded-Groups"},
			headerMappingBase: "/",
			claims: jwt.MapClaims{
				"sub":   "test-user",
				"email": "user@example.com",
				"exp":   time.Now().Add(time.Hour).Unix(),
				"iat":   time.Now().Unix(),
			},
			clientHeaders: map[string]string{
				"X-Forwarded-Email":  "attacker@example.com",
				"X-Forwarded-Groups": "admin",
			},
			expectedHeaders: map[string]string{
				"X-Forwarded-Email": "user@example.com",
			},
			missingHeaders: []string{"X-Forwarded-Groups"},
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			receivedHeaders := http.Header{}
			proxyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				for k, v := range r.Header {
					receivedHeaders[k] = v
				}
				w.WriteHeader(http.StatusOK)
			})

			proxyRouter, err := NewProxyRouter("https://example.com", proxyHandler, publicKey, http.Header{}, false, false, tt.headerMapping, tt.headerMappingBase)
			require.NoError(t, err)

			gin.SetMode(gin.TestMode)
			router := gin.New()
			proxyRouter.SetupRoutes(router)

			token, err := createJWT(privateKey, tt.claims)
			require.NoError(t, err)

			req, err := http.NewRequest("GET", "/test", nil)
			require.NoError(t, err)
			req.Header.Set("Authorization", "Bearer "+token)
			for k, v := range tt.clientHeaders {
				req.Header.Set(k, v)
			}

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

			assert.Equal(t, http.StatusOK, w.Code)
			for header, expected := range tt.expectedHeaders {
				assert.Equal(t, expected, receivedHeaders.Get(header), "header %s mismatch", header)
			}
			for _, header := range tt.missingHeaders {
				assert.Empty(t, receivedHeaders.Get(header), "header %s should be stripped", header)
			}
		})
	}
}

func TestProxyRouter_ProtectedResourceTrailingSlash(t *testing.T) {
	_, publicKey, err := generateRSAKeyPair()
	require.NoError(t, err)

	proxyRouter, err := NewProxyRouter("https://example.com/", http.NotFoundHandler(), publicKey, http.Header{}, false, false, nil, "/userinfo")
	require.NoError(t, err)

	gin.SetMode(gin.TestMode)
	router := gin.New()
	proxyRouter.SetupRoutes(router)

	req, err := http.NewRequest("GET", OauthProtectedResourceEndpoint, nil)
	require.NoError(t, err)

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

	assert.Equal(t, http.StatusOK, w.Code)

	var resp protectedResourceResponse
	err = json.NewDecoder(w.Body).Decode(&resp)
	require.NoError(t, err)
	assert.Equal(t, "https://example.com/mcp", resp.Resource)
	assert.Equal(t, []string{"https://example.com/"}, resp.AuthorizationServers)
}

func TestProxyRouter_HTTPStreamingOnlyRejectsSSE(t *testing.T) {
	privateKey, publicKey, err := generateRSAKeyPair()
	require.NoError(t, err)

	cases := []struct {
		name          string
		method        string
		acceptHeader  string
		wantStatus    int
		expectBackend bool
		streamingOnly bool
	}{
		{
			name:          "plain text/event-stream",
			method:        http.MethodGet,
			acceptHeader:  "text/event-stream",
			wantStatus:    http.StatusMethodNotAllowed,
			streamingOnly: true,
		},
		{
			name:          "event-stream with params",
			method:        http.MethodGet,
			acceptHeader:  "text/event-stream; charset=utf-8",
			wantStatus:    http.StatusMethodNotAllowed,
			streamingOnly: true,
		},
		{
			name:          "multiple values",
			method:        http.MethodGet,
			acceptHeader:  "application/json, text/event-stream",
			wantStatus:    http.StatusMethodNotAllowed,
			streamingOnly: true,
		},
		{
			name:          "quality value",
			method:        http.MethodGet,
			acceptHeader:  "text/event-stream;q=0.9",
			wantStatus:    http.StatusMethodNotAllowed,
			streamingOnly: true,
		},
		{
			name:          "post should pass through",
			method:        http.MethodPost,
			acceptHeader:  "text/event-stream",
			wantStatus:    http.StatusOK,
			expectBackend: true,
			streamingOnly: true,
		},
		{
			name:          "get without accept header",
			method:        http.MethodGet,
			acceptHeader:  "",
			wantStatus:    http.StatusOK,
			expectBackend: true,
			streamingOnly: true,
		},
		{
			name:          "get with non-sse accept",
			method:        http.MethodGet,
			acceptHeader:  "application/json",
			wantStatus:    http.StatusOK,
			expectBackend: true,
			streamingOnly: true,
		},
		{
			name:          "sse allowed when streamingOnly disabled",
			method:        http.MethodGet,
			acceptHeader:  "text/event-stream",
			wantStatus:    http.StatusOK,
			expectBackend: true,
			streamingOnly: false,
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			backendCalled := false
			proxyHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				backendCalled = true
				w.WriteHeader(http.StatusOK)
			})

			proxyRouter, err := NewProxyRouter("https://example.com", proxyHandler, publicKey, http.Header{}, tt.streamingOnly, false, nil, "/userinfo")
			require.NoError(t, err)

			gin.SetMode(gin.TestMode)
			router := gin.New()
			proxyRouter.SetupRoutes(router)

			token, err := createJWT(privateKey, jwt.MapClaims{
				"sub": "user",
				"exp": time.Now().Add(time.Hour).Unix(),
				"iat": time.Now().Unix(),
			})
			require.NoError(t, err)

			req, err := http.NewRequest(tt.method, "/mcp", nil)
			require.NoError(t, err)
			req.Header.Set("Authorization", "Bearer "+token)
			req.Header.Set("Accept", tt.acceptHeader)

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

			assert.Equal(t, tt.wantStatus, w.Code)
			assert.Equal(t, tt.expectBackend, backendCalled, "backend call mismatch")
		})
	}
}