package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"strings"
	"testing"
	"time"
)

func key(t *testing.T) []byte {
	t.Helper()
	k, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("cannot generate a key: %v", err)
	}
	return pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(k)})
}

func seg(t *testing.T, s string, out any) {
	t.Helper()
	raw, err := base64.RawURLEncoding.DecodeString(s)
	if err != nil {
		t.Fatalf("not base64url: %v", err)
	}
	if err := json.Unmarshal(raw, out); err != nil {
		t.Fatalf("not JSON: %v", err)
	}
}

func TestAssertionCarriesTheKeyAndTheAudience(t *testing.T) {
	c := config{issuer: "https://zitadel.example", userID: "3970", keyID: "3971"}
	tok, err := assertion(c, time.Unix(1_700_000_000, 0), key(t))
	if err != nil {
		t.Fatalf("signing failed: %v", err)
	}
	parts := strings.Split(tok, ".")
	if len(parts) != 3 {
		t.Fatalf("a JWT has three segments, got %d", len(parts))
	}
	var h map[string]any
	seg(t, parts[0], &h)
	if h["kid"] != "3971" || h["alg"] != "RS256" {
		t.Fatalf("header must name the key and the algorithm, got %v", h)
	}
	var cl map[string]any
	seg(t, parts[1], &cl)
	if cl["iss"] != "3970" || cl["sub"] != "3970" || cl["aud"] != "https://zitadel.example" {
		t.Fatalf("claims are wrong: %v", cl)
	}
	if int64(cl["iat"].(float64)) != 1_699_999_995 {
		t.Fatalf("iat must be backdated five seconds, got %v", cl["iat"])
	}
}

func TestAssertionRejectsAKeyThatIsNotPEM(t *testing.T) {
	if _, err := assertion(config{userID: "u", issuer: "i"}, time.Now(), []byte("nope")); err == nil {
		t.Fatalf("a non-PEM key must be rejected")
	}
}