package main

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"testing"
)

func generateTestPEM(t *testing.T) string {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("failed to generate test key: %v", err)
	}
	der, err := x509.MarshalPKCS8PrivateKey(key)
	if err != nil {
		t.Fatalf("failed to marshal test key: %v", err)
	}
	block := &pem.Block{Type: "PRIVATE KEY", Bytes: der}
	return string(pem.EncodeToMemory(block))
}

func TestParsePrivateKey(t *testing.T) {
	t.Run("valid P-256 PEM", func(t *testing.T) {
		pemData := generateTestPEM(t)
		key, err := ParsePrivateKey(pemData)
		if err != nil {
			t.Fatalf("unexpected error: %v", err)
		}
		if key.Curve != elliptic.P256() {
			t.Fatalf("expected P-256 curve, got %s", key.Curve.Params().Name)
		}
	})

	t.Run("invalid PEM", func(t *testing.T) {
		_, err := ParsePrivateKey("not a pem")
		if err == nil {
			t.Fatal("expected error for invalid PEM")
		}
	})

	t.Run("empty PEM", func(t *testing.T) {
		_, err := ParsePrivateKey("")
		if err == nil {
			t.Fatal("expected error for empty PEM")
		}
	})

	t.Run("wrong key type", func(t *testing.T) {
		// Generate a P-384 key (wrong curve)
		key, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
		if err != nil {
			t.Fatalf("failed to generate P-384 key: %v", err)
		}
		der, err := x509.MarshalPKCS8PrivateKey(key)
		if err != nil {
			t.Fatalf("failed to marshal key: %v", err)
		}
		block := &pem.Block{Type: "PRIVATE KEY", Bytes: der}
		pemData := string(pem.EncodeToMemory(block))

		_, err = ParsePrivateKey(pemData)
		if err == nil {
			t.Fatal("expected error for non-P-256 key")
		}
	})
}

func TestBuildJWKS(t *testing.T) {
	pemData := generateTestPEM(t)
	key, err := ParsePrivateKey(pemData)
	if err != nil {
		t.Fatalf("failed to parse key: %v", err)
	}

	kid := "test-key-1"
	jwksBytes, err := BuildJWKS([]keyEntry{{privateKey: key, kid: kid}})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	var jwks struct {
		Keys []struct {
			Kty string `json:"kty"`
			Crv string `json:"crv"`
			Kid string `json:"kid"`
			Use string `json:"use"`
			Alg string `json:"alg"`
			X   string `json:"x"`
			Y   string `json:"y"`
		} `json:"keys"`
	}
	if err := json.Unmarshal(jwksBytes, &jwks); err != nil {
		t.Fatalf("failed to unmarshal JWKS: %v", err)
	}

	if len(jwks.Keys) != 1 {
		t.Fatalf("expected 1 key, got %d", len(jwks.Keys))
	}

	k := jwks.Keys[0]
	if k.Kty != "EC" {
		t.Errorf("expected kty=EC, got %s", k.Kty)
	}
	if k.Crv != "P-256" {
		t.Errorf("expected crv=P-256, got %s", k.Crv)
	}
	if k.Kid != kid {
		t.Errorf("expected kid=%s, got %s", kid, k.Kid)
	}
	if k.Use != "sig" {
		t.Errorf("expected use=sig, got %s", k.Use)
	}
	if k.Alg != "ES256" {
		t.Errorf("expected alg=ES256, got %s", k.Alg)
	}
	if k.X == "" || k.Y == "" {
		t.Error("expected non-empty x and y coordinates")
	}
}

func TestBuildJWKS_MultipleKeys(t *testing.T) {
	pem1 := generateTestPEM(t)
	key1, err := ParsePrivateKey(pem1)
	if err != nil {
		t.Fatalf("failed to parse key1: %v", err)
	}

	pem2 := generateTestPEM(t)
	key2, err := ParsePrivateKey(pem2)
	if err != nil {
		t.Fatalf("failed to parse key2: %v", err)
	}

	jwksBytes, err := BuildJWKS([]keyEntry{
		{privateKey: key1, kid: "new-key"},
		{privateKey: key2, kid: "old-key"},
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	var jwks struct {
		Keys []struct {
			Kid string `json:"kid"`
			Kty string `json:"kty"`
		} `json:"keys"`
	}
	if err := json.Unmarshal(jwksBytes, &jwks); err != nil {
		t.Fatalf("failed to unmarshal JWKS: %v", err)
	}

	if len(jwks.Keys) != 2 {
		t.Fatalf("expected 2 keys, got %d", len(jwks.Keys))
	}

	kids := map[string]bool{}
	for _, k := range jwks.Keys {
		kids[k.Kid] = true
	}
	if !kids["new-key"] {
		t.Error("expected kid=new-key in JWKS")
	}
	if !kids["old-key"] {
		t.Error("expected kid=old-key in JWKS")
	}
}

func TestBuildJWKS_NoPrivateKey(t *testing.T) {
	pemData := generateTestPEM(t)
	key, err := ParsePrivateKey(pemData)
	if err != nil {
		t.Fatalf("failed to parse key: %v", err)
	}

	jwksBytes, err := BuildJWKS([]keyEntry{{privateKey: key, kid: "test-key"}})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	// Verify no private key material (d) is present
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(jwksBytes, &raw); err != nil {
		t.Fatalf("failed to unmarshal: %v", err)
	}

	var keys []map[string]json.RawMessage
	if err := json.Unmarshal(raw["keys"], &keys); err != nil {
		t.Fatalf("failed to unmarshal keys: %v", err)
	}

	if _, hasD := keys[0]["d"]; hasD {
		t.Fatal("JWKS must not contain private key material (d)")
	}
}

func TestNewSigner(t *testing.T) {
	pemData := generateTestPEM(t)
	key, err := ParsePrivateKey(pemData)
	if err != nil {
		t.Fatalf("failed to parse key: %v", err)
	}

	kid := "signer-key-1"
	signer, err := NewSigner(key, kid)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if signer.KeyID() != kid {
		t.Errorf("expected kid=%s, got %s", kid, signer.KeyID())
	}
	if signer.Algorithm().String() != "ES256" {
		t.Errorf("expected alg=ES256, got %s", signer.Algorithm())
	}
}