// Shared test utilities. The package-level sharedTestJWT is bound during
// TestMain so every test can produce a signed JWT without re-spinning a
// JWKS server per test.
//
// This replaces the previous mintToken() unsigned-JWT helper. Post-patch
// the KMS authz layer requires real RS256 signatures; an unsigned token
// returns 401 (as Red F3 demands).
package kms

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

	"github.com/golang-jwt/jwt/v5"
)

// shared fixtures bound in TestMain.
var (
	sharedPriv     *rsa.PrivateKey
	sharedKid      = "kms-shared-test-kid"
	sharedJWKS     *httptest.Server
	sharedIssuer   = "https://iam.test.hanzo.id"
	sharedAudience = "kms"
)

func TestMain(m *testing.M) {
	// Generate a single RSA key for the whole test binary.
	priv, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic("testutil: keygen: " + err.Error())
	}
	sharedPriv = priv

	nB64 := base64.RawURLEncoding.EncodeToString(priv.N.Bytes())
	eB64 := base64.RawURLEncoding.EncodeToString(big.NewInt(int64(priv.E)).Bytes())
	jwksJSON, _ := json.Marshal(map[string]any{
		"keys": []map[string]string{
			{"kty": "RSA", "kid": sharedKid, "use": "sig", "alg": "RS256", "n": nB64, "e": eB64},
		},
	})
	sharedJWKS = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Write(jwksJSON)
	}))

	os.Setenv("IAM_ISSUER", sharedIssuer)
	os.Setenv("IAM_AUDIENCE", sharedAudience)
	os.Setenv("IAM_KEYS_URL", sharedJWKS.URL)
	os.Setenv("KMS_ENV", "dev")

	applyAuthConfig(loadAuthConfig())
	resetJWKSCacheForTest()

	code := m.Run()
	sharedJWKS.Close()
	os.Exit(code)
}

// probePath is the route the JWT suite authenticates against.
const probePath = "/v1/kms/probe"

// registerAuthProbe mounts a test-only route whose entire body is the
// production authentication seam: authorize() decides, and a verified
// bearer gets 200.
//
// The JWT suite used to aim at the secret CRUD routes. Those are gone —
// authorization for secrets belongs to cloud, which folds the caller's
// org into the storage key. The verification contract they were really
// exercising (signature, alg allowlist, iss, aud, exp, kid) is a
// property of authorize(), so the probe targets it directly instead of
// through a route that no longer exists. Same assertions, no production
// surface required to host them.
func registerAuthProbe(mux *http.ServeMux) {
	mux.HandleFunc("GET "+probePath, func(w http.ResponseWriter, r *http.Request) {
		claims, ok := authorize(w, r)
		if !ok {
			return // authorize() already wrote 401
		}
		writeJSON(w, http.StatusOK, map[string]any{"sub": claims.Sub})
	})
}

// mintTestJWTSigned produces an RS256-signed JWT using the shared test
// keypair + kid + JWKS. Defaults iss/aud/exp/iat to the shared test
// fixtures unless the caller overrides them.
func mintTestJWTSigned(t *testing.T, claims jwt.MapClaims) string {
	t.Helper()
	if _, ok := claims["iss"]; !ok {
		claims["iss"] = sharedIssuer
	}
	if _, ok := claims["aud"]; !ok {
		claims["aud"] = sharedAudience
	}
	if _, ok := claims["exp"]; !ok {
		claims["exp"] = time.Now().Add(10 * time.Minute).Unix()
	}
	if _, ok := claims["iat"]; !ok {
		claims["iat"] = time.Now().Unix()
	}
	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	token.Header["kid"] = sharedKid
	s, err := token.SignedString(sharedPriv)
	if err != nil {
		t.Fatalf("mintTestJWTSigned: %v", err)
	}
	return s
}