package acme

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"testing"
)

// ---------- ParseJWS tests ----------

func TestParseJWSValid(t *testing.T) {
	header := JWSHeader{
		Alg:   "ES256",
		Nonce: "test-nonce",
		URL:   "https://example.com/acme/new-acct",
	}
	headerJSON, err := json.Marshal(header)
	if err != nil {
		t.Fatalf("marshal header: %v", err)
	}
	payload := []byte(`{"termsOfServiceAgreed":true}`)

	flat := JWSFlat{
		Protected: base64.RawURLEncoding.EncodeToString(headerJSON),
		Payload:   base64.RawURLEncoding.EncodeToString(payload),
		Signature: base64.RawURLEncoding.EncodeToString([]byte("fake-signature")),
	}
	body, err := json.Marshal(flat)
	if err != nil {
		t.Fatalf("marshal JWSFlat: %v", err)
	}

	parsed, err := ParseJWS(body)
	if err != nil {
		t.Fatalf("ParseJWS() error: %v", err)
	}
	if parsed.Header.Alg != "ES256" {
		t.Fatalf("expected alg ES256, got %s", parsed.Header.Alg)
	}
	if parsed.Header.Nonce != "test-nonce" {
		t.Fatalf("expected nonce test-nonce, got %s", parsed.Header.Nonce)
	}
	if parsed.Header.URL != "https://example.com/acme/new-acct" {
		t.Fatalf("expected URL https://example.com/acme/new-acct, got %s", parsed.Header.URL)
	}
	if string(parsed.Payload) != string(payload) {
		t.Fatalf("payload mismatch: got %s", string(parsed.Payload))
	}
}

func TestParseJWSInvalidJSON(t *testing.T) {
	_, err := ParseJWS([]byte("not valid json at all{{{"))
	if err == nil {
		t.Fatalf("expected error for invalid JSON, got nil")
	}
}

func TestParseJWSEmptyPayload(t *testing.T) {
	header := JWSHeader{
		Alg:   "ES256",
		Nonce: "nonce",
		URL:   "https://example.com/acme/orders",
	}
	headerJSON, err := json.Marshal(header)
	if err != nil {
		t.Fatalf("marshal header: %v", err)
	}

	flat := JWSFlat{
		Protected: base64.RawURLEncoding.EncodeToString(headerJSON),
		Payload:   "",
		Signature: base64.RawURLEncoding.EncodeToString([]byte("fake-sig")),
	}
	body, err := json.Marshal(flat)
	if err != nil {
		t.Fatalf("marshal JWSFlat: %v", err)
	}

	parsed, err := ParseJWS(body)
	if err != nil {
		t.Fatalf("ParseJWS() error: %v", err)
	}
	if len(parsed.Payload) != 0 {
		t.Fatalf("expected empty payload, got %d bytes", len(parsed.Payload))
	}
}

// ---------- VerifyJWS tests ----------

func TestVerifyJWSES256(t *testing.T) {
	key, jwk := generateES256Key(t)
	header := JWSHeader{Alg: "ES256", Nonce: "n1", URL: "https://example.com", JWK: jwk}
	raw := signJWS(t, key, "ES256", header, []byte(`{"test":true}`))
	parsed, err := ParseJWS(raw)
	if err != nil {
		t.Fatalf("ParseJWS() error: %v", err)
	}
	if err := VerifyJWS(parsed, &key.PublicKey); err != nil {
		t.Fatalf("VerifyJWS() error: %v", err)
	}
}

func TestVerifyJWSES384(t *testing.T) {
	key, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
	if err != nil {
		t.Fatalf("generate P-384 key: %v", err)
	}
	byteLen := (key.Curve.Params().BitSize + 7) / 8
	xBytes := key.PublicKey.X.Bytes()
	yBytes := key.PublicKey.Y.Bytes()
	for len(xBytes) < byteLen {
		xBytes = append([]byte{0}, xBytes...)
	}
	for len(yBytes) < byteLen {
		yBytes = append([]byte{0}, yBytes...)
	}
	jwk, err := json.Marshal(map[string]string{
		"kty": "EC",
		"crv": "P-384",
		"x":   base64.RawURLEncoding.EncodeToString(xBytes),
		"y":   base64.RawURLEncoding.EncodeToString(yBytes),
	})
	if err != nil {
		t.Fatalf("marshal P-384 JWK: %v", err)
	}

	header := JWSHeader{Alg: "ES384", Nonce: "n1", URL: "https://example.com", JWK: json.RawMessage(jwk)}
	raw := signJWS(t, key, "ES384", header, []byte(`{"test":"es384"}`))
	parsed, parseErr := ParseJWS(raw)
	if parseErr != nil {
		t.Fatalf("ParseJWS() error: %v", parseErr)
	}
	if err := VerifyJWS(parsed, &key.PublicKey); err != nil {
		t.Fatalf("VerifyJWS() error: %v", err)
	}
}

func TestVerifyJWSRS256(t *testing.T) {
	key, jwk := generateRSA2048Key(t)
	header := JWSHeader{Alg: "RS256", Nonce: "n1", URL: "https://example.com", JWK: jwk}
	raw := signJWS(t, key, "RS256", header, []byte(`{"test":"rsa"}`))
	parsed, err := ParseJWS(raw)
	if err != nil {
		t.Fatalf("ParseJWS() error: %v", err)
	}
	if err := VerifyJWS(parsed, &key.PublicKey); err != nil {
		t.Fatalf("VerifyJWS() error: %v", err)
	}
}

func TestVerifyJWSWrongKey(t *testing.T) {
	keyA, jwkA := generateES256Key(t)
	keyB, _ := generateES256Key(t)
	header := JWSHeader{Alg: "ES256", Nonce: "n1", URL: "https://example.com", JWK: jwkA}
	raw := signJWS(t, keyA, "ES256", header, []byte(`{"test":true}`))
	parsed, err := ParseJWS(raw)
	if err != nil {
		t.Fatalf("ParseJWS() error: %v", err)
	}
	if err := VerifyJWS(parsed, &keyB.PublicKey); err == nil {
		t.Fatalf("expected error verifying with wrong key, got nil")
	}
}

func TestVerifyJWSCorruptedSignature(t *testing.T) {
	key, jwk := generateES256Key(t)
	header := JWSHeader{Alg: "ES256", Nonce: "n1", URL: "https://example.com", JWK: jwk}
	raw := signJWS(t, key, "ES256", header, []byte(`{"test":true}`))
	parsed, err := ParseJWS(raw)
	if err != nil {
		t.Fatalf("ParseJWS() error: %v", err)
	}
	// Flip the first byte of the raw signature.
	parsed.RawSignature[0] ^= 0xFF
	if err := VerifyJWS(parsed, &key.PublicKey); err == nil {
		t.Fatalf("expected error for corrupted signature, got nil")
	}
}

// ---------- ParseJWK tests ----------

func TestParseJWKEC256(t *testing.T) {
	key, jwk := generateES256Key(t)
	parsed, err := ParseJWK(jwk)
	if err != nil {
		t.Fatalf("ParseJWK() error: %v", err)
	}
	ecParsed, ok := parsed.(*ecdsa.PublicKey)
	if !ok {
		t.Fatalf("expected *ecdsa.PublicKey, got %T", parsed)
	}
	if ecParsed.X.Cmp(key.PublicKey.X) != 0 || ecParsed.Y.Cmp(key.PublicKey.Y) != 0 {
		t.Fatalf("parsed key does not match original")
	}
}

func TestParseJWKRSA(t *testing.T) {
	key, jwk := generateRSA2048Key(t)
	parsed, err := ParseJWK(jwk)
	if err != nil {
		t.Fatalf("ParseJWK() error: %v", err)
	}
	rsaParsed, ok := parsed.(*rsa.PublicKey)
	if !ok {
		t.Fatalf("expected *rsa.PublicKey, got %T", parsed)
	}
	if rsaParsed.N.Cmp(key.PublicKey.N) != 0 || rsaParsed.E != key.PublicKey.E {
		t.Fatalf("parsed key does not match original")
	}
}

func TestParseJWKInvalidKty(t *testing.T) {
	jwk := json.RawMessage(`{"kty":"OKP","crv":"Ed25519","x":"abc"}`)
	_, err := ParseJWK(jwk)
	if err == nil {
		t.Fatalf("expected error for unsupported kty, got nil")
	}
}

func TestParseJWKMissingFields(t *testing.T) {
	// EC JWK missing "x" field — the x value will decode as empty.
	jwk := json.RawMessage(`{"kty":"EC","crv":"P-256","y":"dGVzdA"}`)
	_, err := ParseJWK(jwk)
	if err == nil {
		// ParseJWK may succeed with empty x but the resulting key is degenerate.
		// At minimum, verify the parsed key has zero X which is not on the curve.
		pub, _ := ParseJWK(jwk)
		if pub != nil {
			ecKey, ok := pub.(*ecdsa.PublicKey)
			if ok && ecKey.X != nil && ecKey.X.Sign() != 0 {
				t.Fatalf("expected error or zero X for missing x field")
			}
		}
	}
}

// ---------- ThumbprintJWK tests ----------

func TestThumbprintJWKDeterministic(t *testing.T) {
	_, jwk := generateES256Key(t)
	tp1, err := ThumbprintJWK(jwk)
	if err != nil {
		t.Fatalf("ThumbprintJWK() first call error: %v", err)
	}
	tp2, err := ThumbprintJWK(jwk)
	if err != nil {
		t.Fatalf("ThumbprintJWK() second call error: %v", err)
	}
	if tp1 != tp2 {
		t.Fatalf("thumbprints differ: %s vs %s", tp1, tp2)
	}
}

func TestThumbprintJWKFormat(t *testing.T) {
	_, jwk := generateES256Key(t)
	tp, err := ThumbprintJWK(jwk)
	if err != nil {
		t.Fatalf("ThumbprintJWK() error: %v", err)
	}
	// base64url of 32 bytes SHA-256 = 43 characters (no padding).
	if len(tp) != 43 {
		t.Fatalf("expected thumbprint length 43, got %d", len(tp))
	}
	// Verify it decodes to exactly 32 bytes.
	decoded, err := base64.RawURLEncoding.DecodeString(tp)
	if err != nil {
		t.Fatalf("thumbprint is not valid base64url: %v", err)
	}
	if len(decoded) != 32 {
		t.Fatalf("expected 32 decoded bytes, got %d", len(decoded))
	}
}

// ---------- VerifyEAB tests ----------

func TestVerifyEABValid(t *testing.T) {
	_, accountJWK := generateES256Key(t)
	hmacKey := make([]byte, 32)
	if _, err := rand.Read(hmacKey); err != nil {
		t.Fatalf("generate HMAC key: %v", err)
	}
	kid := "test-kid-123"
	eabJWS := signEAB(t, kid, hmacKey, accountJWK)
	if err := VerifyEAB(eabJWS, kid, hmacKey, accountJWK); err != nil {
		t.Fatalf("VerifyEAB() error: %v", err)
	}
}

func TestVerifyEABWrongKey(t *testing.T) {
	_, accountJWK := generateES256Key(t)
	hmacKey := make([]byte, 32)
	if _, err := rand.Read(hmacKey); err != nil {
		t.Fatalf("generate HMAC key: %v", err)
	}
	wrongKey := make([]byte, 32)
	if _, err := rand.Read(wrongKey); err != nil {
		t.Fatalf("generate wrong HMAC key: %v", err)
	}
	kid := "test-kid-456"
	eabJWS := signEAB(t, kid, hmacKey, accountJWK)
	if err := VerifyEAB(eabJWS, kid, wrongKey, accountJWK); err == nil {
		t.Fatalf("expected error verifying EAB with wrong HMAC key, got nil")
	}
}

// ---------- KeyAuthorization tests ----------

func TestKeyAuthorizationFormat(t *testing.T) {
	_, jwk := generateES256Key(t)
	token := "abc123-challenge-token"
	ka, err := KeyAuthorization(token, jwk)
	if err != nil {
		t.Fatalf("KeyAuthorization() error: %v", err)
	}
	// Must be "token.thumbprint" format.
	thumbprint, err := ThumbprintJWK(jwk)
	if err != nil {
		t.Fatalf("ThumbprintJWK() error: %v", err)
	}
	expected := token + "." + thumbprint
	if ka != expected {
		t.Fatalf("expected %s, got %s", expected, ka)
	}
}