package zitadel

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

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

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

func TestSignAssertionCarriesTheKeyIDAndIssuer(t *testing.T) {
	pemBytes, _ := newKey(t)
	now := time.Unix(1_700_000_000, 0)

	token, err := SignAssertion(pemBytes, "key-42", "client-1", "user-9", "https://zitadel.riotpiao.com", now)
	if err != nil {
		t.Fatalf("signing failed: %v", err)
	}
	if strings.Count(token, ".") != 2 {
		t.Fatalf("a JWT has three segments, got %q", token)
	}

	parts := strings.Split(token, ".")
	var header map[string]any
	decodeSegment(t, parts[0], &header)
	if header["kid"] != "key-42" {
		t.Fatalf("the header must name the key id so ZITADEL can pick the key, got %v", header["kid"])
	}
	if header["alg"] != "RS256" {
		t.Fatalf("expected RS256, got %v", header["alg"])
	}

	var claims map[string]any
	decodeSegment(t, parts[1], &claims)
	if claims["iss"] != "client-1" || claims["sub"] != "user-9" {
		t.Fatalf("iss and sub are wrong: %v / %v", claims["iss"], claims["sub"])
	}
	if claims["aud"] != "https://zitadel.riotpiao.com" {
		t.Fatalf("aud must equal the issuer, got %v", claims["aud"])
	}
	exp := int64(claims["exp"].(float64))
	iat := int64(claims["iat"].(float64))
	if exp-iat != 60 {
		t.Fatalf("the assertion must live 60 seconds, got %d", exp-iat)
	}
}

func TestSignAssertionRejectsGarbage(t *testing.T) {
	if _, err := SignAssertion([]byte("not a key"), "k", "c", "u", "https://i", time.Now()); err == nil {
		t.Fatalf("a non-PEM key must be rejected")
	}
}

func TestSignAssertionDefaultsSubjectToClientID(t *testing.T) {
	pemBytes, _ := newKey(t)
	token, err := SignAssertion(pemBytes, "k1", "client-1", "", "https://i", time.Unix(1, 0))
	if err != nil {
		t.Fatalf("signing failed: %v", err)
	}
	var claims map[string]any
	decodeSegment(t, strings.Split(token, ".")[1], &claims)
	if claims["sub"] != "client-1" {
		t.Fatalf("an empty user id must fall back to the client id, got %v", claims["sub"])
	}
}