package main

import (
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"math/big"
	"net/http"
	"net/http/httptest"
	"strconv"
	"sync/atomic"
	"testing"
	"time"

	pdktest "github.com/Kong/go-pdk/test"
	"github.com/golang-jwt/jwt/v5"
	"github.com/stretchr/testify/require"
	authconsts "github.com/tidbcloud/aisupport/packages/sdk-go/infra/auth/consts"
)

const (
	testAppID          = "ari:cloud:ecosystem::app/test-app"
	testInstallationID = "ari:cloud:ecosystem::installation/test-installation"
	testPrincipal      = "atlassian-account-id"
	testKeyID          = "test-key"
)

func TestAccessAcceptsValidFITAndSetsIdentityHeaders(t *testing.T) {
	t.Parallel()

	privateKey := newPrivateKey(t)
	server, requests := newJWKSServer(t, &privateKey.PublicKey)
	token := signFIT(t, privateKey, fitOptions{})

	env := runForgeAccess(t, &Config{
		JWKSURL:                server.URL,
		AllowedAppIDs:          []string{testAppID},
		AllowedInstallationIDs: []string{testInstallationID},
		RequirePrincipal:       true,
	}, pdktest.Request{
		Method: http.MethodGet,
		Url:    "https://example.com/api/v1/jsm-syncer/forge/orgs",
		Headers: http.Header{
			authconsts.HeaderAuthorization:      {"Bearer " + token},
			authconsts.HeaderForgeAppID:         {"spoofed-app"},
			authconsts.HeaderAtlassianAccountID: {"spoofed-user"},
			"X-Forge-OAuth-User":                {"oauth-user-token"},
		},
	})

	require.Zero(t, env.ClientRes.Status)
	require.Equal(t, testAppID, env.ServiceReq.Headers.Get(authconsts.HeaderForgeAppID))
	require.Equal(t, testInstallationID, env.ServiceReq.Headers.Get(authconsts.HeaderForgeInstallationID))
	require.Equal(t, testPrincipal, env.ServiceReq.Headers.Get(authconsts.HeaderAtlassianAccountID))
	require.Equal(t, "invocation-id", env.ServiceReq.Headers.Get(authconsts.HeaderForgeInvocationID))
	require.Empty(t, env.ServiceReq.Headers.Get(authconsts.HeaderAuthorization))
	require.Empty(t, env.ServiceReq.Headers.Get("X-Forge-OAuth-User"))
	require.Equal(t, int32(1), requests.Load())
}

func TestAccessCachesJWKS(t *testing.T) {
	t.Parallel()

	privateKey := newPrivateKey(t)
	server, requests := newJWKSServer(t, &privateKey.PublicKey)
	config := &Config{
		JWKSURL:                server.URL,
		AllowedAppIDs:          []string{testAppID},
		AllowedInstallationIDs: []string{testInstallationID},
		RequirePrincipal:       true,
	}

	for range 2 {
		env := runForgeAccess(t, config, pdktest.Request{
			Method: http.MethodGet,
			Url:    "https://example.com/api/v1/jsm-syncer/forge/orgs",
			Headers: http.Header{
				authconsts.HeaderAuthorization: {"Bearer " + signFIT(t, privateKey, fitOptions{})},
			},
		})
		require.Zero(t, env.ClientRes.Status)
	}
	require.Equal(t, int32(1), requests.Load())
}

func TestAccessAllowsInstallationInExplicitBootstrapMode(t *testing.T) {
	t.Parallel()

	privateKey := newPrivateKey(t)
	server, _ := newJWKSServer(t, &privateKey.PublicKey)
	env := runForgeAccess(t, &Config{
		JWKSURL:                        server.URL,
		AllowedAppIDs:                  []string{testAppID},
		AllowUnregisteredInstallations: true,
		RequirePrincipal:               true,
	}, pdktest.Request{
		Method: http.MethodGet,
		Url:    "https://example.com/api/v1/jsm-syncer/forge/orgs",
		Headers: http.Header{
			authconsts.HeaderAuthorization: {"Bearer " + signFIT(t, privateKey, fitOptions{})},
		},
	})

	require.Zero(t, env.ClientRes.Status)
	require.Equal(t, testInstallationID, env.ServiceReq.Headers.Get(authconsts.HeaderForgeInstallationID))
}

func TestAccessRejectsInvalidFITContext(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name          string
		options       fitOptions
		authorization string
	}{
		{
			name:          "missing token",
			authorization: "",
		},
		{
			name:    "wrong audience",
			options: fitOptions{audience: "ari:cloud:ecosystem::app/other"},
		},
		{
			name:    "wrong issuer",
			options: fitOptions{issuer: "other-issuer"},
		},
		{
			name:    "expired",
			options: fitOptions{expiresAt: time.Now().Add(-time.Minute)},
		},
		{
			name:    "installation not allowed",
			options: fitOptions{installationID: "ari:cloud:ecosystem::installation/other"},
		},
		{
			name:    "principal required",
			options: fitOptions{principal: stringPointer("")},
		},
	}

	for _, test := range tests {
		test := test
		t.Run(test.name, func(t *testing.T) {
			t.Parallel()

			privateKey := newPrivateKey(t)
			server, _ := newJWKSServer(t, &privateKey.PublicKey)
			config := Config{
				JWKSURL:                server.URL,
				AllowedAppIDs:          []string{testAppID},
				AllowedInstallationIDs: []string{testInstallationID},
				RequirePrincipal:       true,
			}

			authorization := test.authorization
			if authorization == "" && test.name != "missing token" {
				authorization = "Bearer " + signFIT(t, privateKey, test.options)
			}
			env := runForgeAccess(t, &config, pdktest.Request{
				Method: http.MethodGet,
				Url:    "https://example.com/api/v1/jsm-syncer/forge/orgs",
				Headers: http.Header{
					authconsts.HeaderAuthorization: {authorization},
				},
			})

			require.Equal(t, http.StatusUnauthorized, env.ClientRes.Status)
			var response map[string]any
			require.NoError(t, json.Unmarshal(env.ClientRes.Body, &response))
			require.Equal(t, "unauthorized", response["errmsg"])
		})
	}
}

func TestAccessRejectsIncompleteConfiguration(t *testing.T) {
	t.Parallel()

	env := runForgeAccess(t, &Config{}, pdktest.Request{
		Method: http.MethodGet,
		Url:    "https://example.com/api/v1/jsm-syncer/forge/orgs",
		Headers: http.Header{
			authconsts.HeaderAuthorization: {"Bearer token"},
		},
	})

	require.Equal(t, http.StatusInternalServerError, env.ClientRes.Status)
}

func TestAccessRejectsInvalidSignature(t *testing.T) {
	t.Parallel()

	trustedKey := newPrivateKey(t)
	untrustedKey := newPrivateKey(t)
	server, _ := newJWKSServer(t, &trustedKey.PublicKey)
	env := runForgeAccess(t, &Config{
		JWKSURL:                server.URL,
		AllowedAppIDs:          []string{testAppID},
		AllowedInstallationIDs: []string{testInstallationID},
		RequirePrincipal:       true,
	}, pdktest.Request{
		Method: http.MethodGet,
		Url:    "https://example.com/api/v1/jsm-syncer/forge/orgs",
		Headers: http.Header{
			authconsts.HeaderAuthorization: {"Bearer " + signFIT(t, untrustedKey, fitOptions{})},
		},
	})

	require.Equal(t, http.StatusUnauthorized, env.ClientRes.Status)
}

type fitOptions struct {
	appID          string
	audience       string
	expiresAt      time.Time
	installationID string
	issuer         string
	principal      *string
}

func signFIT(t *testing.T, privateKey *rsa.PrivateKey, options fitOptions) string {
	t.Helper()

	appID := valueOrDefault(options.appID, testAppID)
	audience := valueOrDefault(options.audience, appID)
	installationID := valueOrDefault(options.installationID, testInstallationID)
	issuer := valueOrDefault(options.issuer, forgeInvocationIssuer)
	expiresAt := options.expiresAt
	if expiresAt.IsZero() {
		expiresAt = time.Now().Add(5 * time.Minute)
	}
	principal := testPrincipal
	if options.principal != nil {
		principal = *options.principal
	}

	claims := forgeInvocationClaims{
		Principal: principal,
		RegisteredClaims: jwt.RegisteredClaims{
			Audience:  jwt.ClaimStrings{audience},
			ExpiresAt: jwt.NewNumericDate(expiresAt),
			ID:        "invocation-id",
			IssuedAt:  jwt.NewNumericDate(time.Now().Add(-time.Second)),
			Issuer:    issuer,
			NotBefore: jwt.NewNumericDate(time.Now().Add(-time.Second)),
		},
	}
	claims.App.ID = appID
	claims.App.InstallationID = installationID

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	token.Header["kid"] = testKeyID
	signed, err := token.SignedString(privateKey)
	require.NoError(t, err)
	return signed
}

func newPrivateKey(t *testing.T) *rsa.PrivateKey {
	t.Helper()
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	return key
}

func newJWKSServer(t *testing.T, publicKey *rsa.PublicKey) (*httptest.Server, *atomic.Int32) {
	t.Helper()

	var requests atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		requests.Add(1)
		w.Header().Set("Content-Type", "application/json")
		require.NoError(t, json.NewEncoder(w).Encode(jwksDocument{
			Keys: []jsonWebKey{{
				Algorithm: "RS256",
				Exponent:  encodeExponent(publicKey.E),
				KeyID:     testKeyID,
				KeyType:   "RSA",
				Modulus:   base64.RawURLEncoding.EncodeToString(publicKey.N.Bytes()),
				Use:       "sig",
			}},
		}))
	}))
	t.Cleanup(server.Close)
	return server, &requests
}

func encodeExponent(exponent int) string {
	value := big.NewInt(int64(exponent)).Bytes()
	return base64.RawURLEncoding.EncodeToString(value)
}

func runForgeAccess(t *testing.T, config *Config, request pdktest.Request) *pdktest.TestEnv {
	t.Helper()
	env, err := pdktest.New(t, request)
	require.NoError(t, err)
	env.DoAccess(config)
	return env
}

func valueOrDefault(value, fallback string) string {
	if value != "" {
		return value
	}
	return fallback
}

func stringPointer(value string) *string {
	return &value
}

func TestEncodeExponent(t *testing.T) {
	t.Parallel()
	decoded, err := base64.RawURLEncoding.DecodeString(encodeExponent(65537))
	require.NoError(t, err)
	require.Equal(t, "65537", strconv.FormatInt(new(big.Int).SetBytes(decoded).Int64(), 10))
}