package signing

import (
	"context"
	"encoding/json"
	"os"
	"testing"
	"time"

	"github.com/go-jose/go-jose/v4"
	baoapi "github.com/openbao/openbao/api/v2"
)

func TestTransitSignerProducesJWTVerifiedByPublishedJWKS(t *testing.T) {
	address := os.Getenv("WORKLOAD_STS_TEST_BAO_ADDR")
	token := os.Getenv("WORKLOAD_STS_TEST_BAO_TOKEN")
	if address == "" || token == "" {
		t.Skip("set WORKLOAD_STS_TEST_BAO_ADDR and WORKLOAD_STS_TEST_BAO_TOKEN")
	}

	config := baoapi.DefaultConfig()
	config.Address = address
	client, err := baoapi.NewClient(config)
	if err != nil {
		t.Fatalf("create OpenBao client: %v", err)
	}
	client.SetToken(token)

	signer, err := NewTransitSigner(client, "transit", "workload-sts")
	if err != nil {
		t.Fatalf("NewTransitSigner() error = %v", err)
	}
	issuer, err := NewIssuer(signer)
	if err != nil {
		t.Fatalf("NewIssuer() error = %v", err)
	}
	now := time.Now().UTC().Truncate(time.Second)
	issuer.now = func() time.Time { return now }

	compact, 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)
	}

	set, err := signer.JWKS(context.Background())
	if err != nil {
		t.Fatalf("JWKS() error = %v", err)
	}
	encodedSet, err := json.Marshal(set)
	if err != nil {
		t.Fatalf("marshal JWKS: %v", err)
	}
	var decodedSet jose.JSONWebKeySet
	if err := json.Unmarshal(encodedSet, &decodedSet); err != nil {
		t.Fatalf("decode JWKS with go-jose: %v", err)
	}

	jws, err := jose.ParseSigned(compact, []jose.SignatureAlgorithm{jose.RS256})
	if err != nil {
		t.Fatalf("parse compact JWT: %v", err)
	}
	keyID := jws.Signatures[0].Header.KeyID
	keys := decodedSet.Key(keyID)
	if len(keys) != 1 {
		t.Fatalf("JWKS keys for %q = %d, want 1", keyID, len(keys))
	}
	if _, err := jws.Verify(keys[0].Key); err != nil {
		t.Fatalf("published JWKS rejected Transit signature: %v", err)
	}
}