package idp

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/url"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestGeminiRegistersAllRelayCallbacks(t *testing.T) {
	hosts := []string{
		"oauth-redirect.googleusercontent.com",
		"oauth-redirect-sandbox.googleusercontent.com",
		"oauth-redirect-test.googleusercontent.com",
	}
	var prefixes, redirects []string
	for _, host := range hosts {
		prefixes = append(prefixes, "https://"+host+"/")
		for _, callback := range []string{"/r/user_bound_custom-mcp-test", "/other-callback/custom-mcp-test"} {
			redirects = append(redirects, "https://"+host+callback)
		}
	}
	policy := strings.Join(prefixes, ",")
	var oldPrefixes []string
	for _, prefix := range prefixes {
		oldPrefixes = append(oldPrefixes, prefix+"r/")
	}
	require.False(t, validRedirectURIs(redirects, strings.Join(oldPrefixes, ",")))
	require.False(t, validRedirectURIs(redirects, prefixes[0]))
	require.True(t, validRedirectURIs(redirects, policy))
	t.Setenv("MCP_REDIRECT_PREFIXES", policy)
	server, _, _ := setupTestServer(t)
	defer server.Close()

	register := func(t *testing.T, uris []string) *http.Response {
		t.Helper()
		body, err := json.Marshal(registrationRequest{
			RedirectURIs: uris, TokenEndpointAuthMethod: "client_secret_basic",
			GrantTypes:    []string{"authorization_code", "refresh_token"},
			ResponseTypes: []string{"code"},
		})
		require.NoError(t, err)
		response, err := http.Post(server.URL+RegistrationEndpoint, "application/json", bytes.NewReader(body))
		require.NoError(t, err)
		return response
	}
	response := register(t, redirects)
	defer response.Body.Close()
	require.Equal(t, http.StatusCreated, response.StatusCode)
	var client registrationResponse
	require.NoError(t, json.NewDecoder(response.Body).Decode(&client))
	require.Equal(t, redirects, client.RedirectURIs)
	require.NotEmpty(t, client.ClientID)
	require.NotEmpty(t, client.ClientSecret)
	require.Equal(t, "client_secret_basic", client.TokenEndpointAuthMethod)

	t.Run("authorize requires an exact registered callback", func(t *testing.T) {
		browser := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error {
			return http.ErrUseLastResponse
		}}
		query := url.Values{
			"client_id": {client.ClientID}, "response_type": {"code"},
			"redirect_uri": {"https://" + hosts[0] + "/other-callback/unregistered"},
			"state":        {"redirect-mismatch-test"}, "code_challenge": {strings.Repeat("a", 43)},
			"code_challenge_method": {"S256"},
		}
		response, err := browser.Get(server.URL + AuthorizationEndpoint + "?" + query.Encode())
		require.NoError(t, err)
		defer response.Body.Close()
		require.Equal(t, http.StatusBadRequest, response.StatusCode)
		require.Empty(t, response.Header.Get("Location"))
		var rejection map[string]any
		require.NoError(t, json.NewDecoder(response.Body).Decode(&rejection))
		require.Contains(t, rejection["error_description"], "redirect_uri")
	})

	for _, host := range hosts {
		for _, bad := range []string{
			"https://" + host + ".evil/r/steal",
			"http://" + host + "/r/steal",
			"https://" + host + ":444/r/steal",
			"https://" + strings.Replace(host, "oauth-redirect", "unknown-relay", 1) + "/r/steal",
			"https://attacker@" + host + "/r/steal",
			"https://" + host + "/r/steal?value=1",
			"https://" + host + "/r/steal#fragment",
			"https://" + host + "/other/%2e%2e/steal",
			"https://" + host + "/r/../steal",
		} {
			t.Run(bad, func(t *testing.T) {
				mixed := append([]string(nil), redirects...)
				mixed[len(mixed)-1] = bad
				response := register(t, mixed)
				defer response.Body.Close()
				require.Equal(t, http.StatusBadRequest, response.StatusCode)
			})
		}
	}
}