package jose_test

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"
	"strings"
	"testing"

	"keycape/internal/jose"
)

// internal/jose is the single signature verifier behind both the caller CLI and
// upstream provider verification, so it is tested directly rather than only
// through its callers.

var testKey = func() *rsa.PrivateKey {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic(err)
	}
	return key
}()

func jwks(kid string, key *rsa.PrivateKey) []byte {
	return []byte(fmt.Sprintf(`{"keys":[{"kty":"RSA","use":"sig","alg":"RS256","kid":%q,"n":%q,"e":%q}]}`,
		kid,
		base64.RawURLEncoding.EncodeToString(key.PublicKey.N.Bytes()),
		base64.RawURLEncoding.EncodeToString(big.NewInt(int64(key.PublicKey.E)).Bytes())))
}

func sign(t *testing.T, header, claims string, key *rsa.PrivateKey) string {
	t.Helper()
	input := base64.RawURLEncoding.EncodeToString([]byte(header)) + "." +
		base64.RawURLEncoding.EncodeToString([]byte(claims))
	digest := sha256.Sum256([]byte(input))
	signature, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, digest[:])
	if err != nil {
		t.Fatal(err)
	}
	return input + "." + base64.RawURLEncoding.EncodeToString(signature)
}

func TestVerifyAcceptsGenuineTokenAndReturnsClaims(t *testing.T) {
	keys, err := jose.ParseJWKS(jwks("k1", testKey))
	if err != nil {
		t.Fatal(err)
	}
	token := sign(t, `{"alg":"RS256","kid":"k1"}`, `{"sub":"alice","n":1}`, testKey)
	claims, err := jose.Verify(token, keys)
	if err != nil {
		t.Fatalf("genuine token rejected: %v", err)
	}
	if claims["sub"] != "alice" {
		t.Fatalf("claims not returned: %v", claims)
	}
}

func TestVerifyRejections(t *testing.T) {
	keys, err := jose.ParseJWKS(jwks("k1", testKey))
	if err != nil {
		t.Fatal(err)
	}
	other, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	genuine := sign(t, `{"alg":"RS256","kid":"k1"}`, `{"sub":"alice"}`, testKey)
	parts := strings.Split(genuine, ".")

	cases := map[string]string{
		"wrong signing key": sign(t, `{"alg":"RS256","kid":"k1"}`, `{"sub":"alice"}`, other),
		"unknown key id":    sign(t, `{"alg":"RS256","kid":"k2"}`, `{"sub":"alice"}`, testKey),
		"missing key id":    sign(t, `{"alg":"RS256"}`, `{"sub":"alice"}`, testKey),
		"alg none":          sign(t, `{"alg":"none","kid":"k1"}`, `{"sub":"alice"}`, testKey),
		"alg HS256":         sign(t, `{"alg":"HS256","kid":"k1"}`, `{"sub":"alice"}`, testKey),
		// A critical extension we do not understand must not be ignored.
		"critical extension": sign(t, `{"alg":"RS256","kid":"k1","crit":["exp"]}`, `{"sub":"alice"}`, testKey),
		"tampered payload":   parts[0] + "." + base64.RawURLEncoding.EncodeToString([]byte(`{"sub":"mallory"}`)) + "." + parts[2],
		"two segments":       parts[0] + "." + parts[1],
		"payload not json":   sign(t, `{"alg":"RS256","kid":"k1"}`, `not-json`, testKey),
		"empty":              "",
	}

	for name, token := range cases {
		t.Run(name, func(t *testing.T) {
			if _, err := jose.Verify(token, keys); err == nil {
				t.Fatal("accepted")
			}
		})
	}
}

func TestParseJWKSRejectsUnusableKeySets(t *testing.T) {
	small, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		t.Fatal(err)
	}
	cases := map[string][]byte{
		"undersized modulus": jwks("k1", small),
		"empty set":          []byte(`{"keys":[]}`),
		"no rsa signing key": []byte(`{"keys":[{"kty":"EC","kid":"k1","crv":"P-256"}]}`),
		"missing key id":     []byte(`{"keys":[{"kty":"RSA","n":"AQAB","e":"AQAB"}]}`),
		"bad base64":         []byte(`{"keys":[{"kty":"RSA","kid":"k1","n":"!!!","e":"AQAB"}]}`),
		"even exponent":      []byte(fmt.Sprintf(`{"keys":[{"kty":"RSA","kid":"k1","n":%q,"e":"AQAA"}]}`, base64.RawURLEncoding.EncodeToString(testKey.PublicKey.N.Bytes()))),
		"not json":           []byte(`nonsense`),
	}
	for name, raw := range cases {
		t.Run(name, func(t *testing.T) {
			if _, err := jose.ParseJWKS(raw); err == nil {
				t.Fatal("accepted")
			}
		})
	}
}

// Two entries under one key id make key selection ambiguous, so the set is
// refused rather than resolved by order.
func TestParseJWKSRejectsDuplicateKeyIDs(t *testing.T) {
	var first, second map[string]interface{}
	if err := json.Unmarshal(jwks("k1", testKey), &first); err != nil {
		t.Fatal(err)
	}
	if err := json.Unmarshal(jwks("k1", testKey), &second); err != nil {
		t.Fatal(err)
	}
	merged, err := json.Marshal(map[string]interface{}{
		"keys": []interface{}{
			first["keys"].([]interface{})[0],
			second["keys"].([]interface{})[0],
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := jose.ParseJWKS(merged); err == nil {
		t.Fatal("duplicate key id accepted")
	}
}

// Keys for other algorithms alongside a usable RS256 key are ignored, not fatal.
func TestParseJWKSIgnoresIrrelevantKeys(t *testing.T) {
	var usable map[string]interface{}
	if err := json.Unmarshal(jwks("k1", testKey), &usable); err != nil {
		t.Fatal(err)
	}
	mixed, err := json.Marshal(map[string]interface{}{
		"keys": []interface{}{
			map[string]interface{}{"kty": "EC", "kid": "ec", "crv": "P-256"},
			usable["keys"].([]interface{})[0],
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	keys, err := jose.ParseJWKS(mixed)
	if err != nil {
		t.Fatalf("usable key set rejected: %v", err)
	}
	if _, ok := keys["k1"]; !ok || len(keys) != 1 {
		t.Fatalf("unexpected key set: %v", keys)
	}
}