package signing

import (
	"context"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"strings"
	"testing"
	"time"

	"github.com/go-jose/go-jose/v4"
)

func TestIssuerProducesInteroperableRS256JWT(t *testing.T) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}
	signer := &localRS256Signer{keyID: "workload-sts-v1", key: privateKey}
	issuer, err := NewIssuer(signer)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}
	now := time.Unix(1_789_062_000, 0)
	issuer.now = func() time.Time { return now }

	token, err := issuer.Sign(context.Background(), Claims{
		Issuer:        "https://identity.ad.ddupan.top",
		Subject:       "01993f4d-5e1a-7000-8000-000000000001",
		PrincipalName: "ci/homelab-infra-plan",
		Audience:      []string{"https://bao.ad.ddupan.top:8200"},
		IssuedAt:      now.Unix(),
		NotBefore:     now.Unix(),
		ExpiresAt:     now.Add(5 * time.Minute).Unix(),
		JWTID:         "01993f4d-5e1a-7000-8000-000000000002",
		Scope:         "bao.login",
	})
	if err != nil {
		t.Fatalf("Sign() error = %v", err)
	}

	jws, err := jose.ParseSigned(token, []jose.SignatureAlgorithm{jose.RS256})
	if err != nil {
		t.Fatalf("go-jose rejected compact JWT: %v", err)
	}
	payload, err := jws.Verify(&privateKey.PublicKey)
	if err != nil {
		t.Fatalf("go-jose rejected signature: %v", err)
	}

	var got Claims
	if err := json.Unmarshal(payload, &got); err != nil {
		t.Fatalf("decode claims: %v", err)
	}
	if got.Subject != "01993f4d-5e1a-7000-8000-000000000001" {
		t.Fatalf("subject = %q", got.Subject)
	}

	parts := strings.Split(token, ".")
	headerJSON, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		t.Fatalf("decode header: %v", err)
	}
	var header map[string]string
	if err := json.Unmarshal(headerJSON, &header); err != nil {
		t.Fatalf("unmarshal header: %v", err)
	}
	if header["typ"] != "at+jwt" || header["kid"] != "workload-sts-v1" {
		t.Fatalf("header = %#v", header)
	}
}

func TestIssuerRejectsExcessiveLifetime(t *testing.T) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("generate key: %v", err)
	}
	issuer, err := NewIssuer(&localRS256Signer{keyID: "test", key: privateKey})
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}
	now := time.Unix(1_789_062_000, 0)
	issuer.now = func() time.Time { return now }

	_, err = issuer.Sign(context.Background(), Claims{
		Issuer:    "https://identity.ad.ddupan.top",
		Subject:   "principal",
		Audience:  []string{"audience"},
		IssuedAt:  now.Unix(),
		NotBefore: now.Unix(),
		ExpiresAt: now.Add(5*time.Minute + time.Second).Unix(),
		JWTID:     "jti",
		Scope:     "scope",
	})
	if err == nil {
		t.Fatal("Sign() error = nil, want excessive lifetime error")
	}
}

func TestIntegerRejectsNonIntegralJSONNumber(t *testing.T) {
	for _, value := range []any{1.5, json.Number("1.5")} {
		if _, err := integer(value); err == nil {
			t.Fatalf("integer(%v) error = nil, want invalid integer error", value)
		}
	}
}

type localRS256Signer struct {
	keyID string
	key   *rsa.PrivateKey
}

func (s *localRS256Signer) ActiveKey(context.Context) (SigningKey, error) {
	return SigningKey{ID: s.keyID, Version: 1}, nil
}

func (s *localRS256Signer) SignRS256(_ context.Context, _ SigningKey, input []byte) ([]byte, error) {
	digest := sha256.Sum256(input)
	return rsa.SignPKCS1v15(rand.Reader, s.key, crypto.SHA256, digest[:])
}