package payment

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"math/big"
	"testing"
	"time"
)

func TestValidateAppleJWSChain(t *testing.T) {
	chain := testAppleCertificateChain(t)
	jws := testAppleJWS(t, "ES256", chain, time.Now())
	if err := validateAppleJWSChainWithRoots(jws, testAppleTrustedRoots(t, chain)); err != nil {
		t.Fatalf("validateAppleJWSChain() error = %v", err)
	}
}

func TestValidateAppleJWSChainRejectsUnrelatedLeaf(t *testing.T) {
	chain := testAppleCertificateChain(t)
	attackerChain := testAppleCertificateChain(t)
	chain[0] = attackerChain[0]
	if err := validateAppleJWSChainWithRoots(testAppleJWS(t, "ES256", chain, time.Now()), testAppleTrustedRoots(t, chain)); err == nil {
		t.Fatal("validateAppleJWSChain() accepted a leaf outside the declared chain")
	}
}

func TestValidateAppleJWSChainRejectsInvalidSignature(t *testing.T) {
	chain := testAppleCertificateChain(t)
	jws := testAppleJWS(t, "ES256", chain, time.Now())
	header, payload, _, err := splitAppleJWS(jws)
	if err != nil {
		t.Fatal(err)
	}
	tampered := header + "." + payload + "." + base64.RawURLEncoding.EncodeToString(make([]byte, 64))
	if err = validateAppleJWSChainWithRoots(tampered, testAppleTrustedRoots(t, chain)); err == nil {
		t.Fatal("validateAppleJWSChain accepted an invalid ES256 signature")
	}
}

func TestValidateAppleJWSChainRejectsUnexpectedAlgorithmAndShape(t *testing.T) {
	chain := testAppleCertificateChain(t)
	if err := validateAppleJWSChainWithRoots(testAppleJWS(t, "ES384", chain, time.Now()), testAppleTrustedRoots(t, chain)); err == nil {
		t.Fatal("validateAppleJWSChain() accepted a non-ES256 algorithm")
	}
	if err := validateAppleJWSChainWithRoots(testAppleJWS(t, "ES256", chain[:2], time.Now()), nil); err == nil {
		t.Fatal("validateAppleJWSChain() accepted an incomplete x5c chain")
	}
}

func TestValidateAppleJWSChainRejectsAdditionalCertificates(t *testing.T) {
	chain := testAppleCertificateChain(t)
	chain = append(chain, chain[2])
	if err := validateAppleJWSChainWithRoots(testAppleJWS(t, "ES256", chain, time.Now()), testAppleTrustedRoots(t, chain[:3])); err == nil {
		t.Fatal("validateAppleJWSChain() accepted additional x5c certificates")
	}
}

func TestValidateAppleJWSChainRejectsNonAppleSigningCertificates(t *testing.T) {
	chain := testAppleCertificateChainWithoutAppleExtensions(t)
	if err := validateAppleJWSChainWithRoots(testAppleJWS(t, "ES256", chain, time.Now()), testAppleTrustedRoots(t, chain)); err == nil {
		t.Fatal("validateAppleJWSChain() accepted certificates without Apple signing extensions")
	}
}

func TestValidateAppleJWSChainUsesSignedDateForHistoricalPayload(t *testing.T) {
	signedAt := time.Now().AddDate(-2, 0, 0)
	chain := testAppleCertificateChainAt(t, signedAt, true)
	if err := validateAppleJWSChainWithRoots(testAppleJWS(t, "ES256", chain, signedAt), testAppleTrustedRoots(t, chain)); err != nil {
		t.Fatalf("validateAppleJWSChain() rejected historical payload: %v", err)
	}
}

func TestValidateAppleJWSChainUsesNestedSignedDateForNotification(t *testing.T) {
	signedAt := time.Now().AddDate(-2, 0, 0)
	chain := testAppleCertificateChainAt(t, signedAt, true)
	inner := testAppleJWS(t, "ES256", chain, signedAt)
	header, err := json.Marshal(map[string]any{"alg": "ES256", "x5c": chain})
	if err != nil {
		t.Fatal(err)
	}
	payload, err := json.Marshal(map[string]any{
		"notificationType": "DID_RENEW",
		"data":             map[string]any{"signedTransactionInfo": inner},
	})
	if err != nil {
		t.Fatal(err)
	}
	outer := testAppleSignedJWS(t, header, payload, chain)
	if err := validateAppleJWSChainWithRoots(outer, testAppleTrustedRoots(t, chain)); err != nil {
		t.Fatalf("validateAppleJWSChain() rejected historical notification: %v", err)
	}
}

func TestValidateAppleJWSChainRejectsExpiredChainWithoutSignedDate(t *testing.T) {
	signedAt := time.Now().AddDate(-2, 0, 0)
	chain := testAppleCertificateChainAt(t, signedAt, true)
	if err := validateAppleJWSChainWithRoots(testAppleJWSWithoutDate(t, chain), testAppleTrustedRoots(t, chain)); err == nil {
		t.Fatal("validateAppleJWSChain() accepted an expired chain without signedDate")
	}
}

func TestValidateAppleJWSChainRejectsBrokenChainWithoutSignedDate(t *testing.T) {
	chain := testAppleCertificateChain(t)
	attackerChain := testAppleCertificateChain(t)
	chain[0] = attackerChain[0]
	if err := validateAppleJWSChainWithRoots(testAppleJWSWithoutDate(t, chain), testAppleTrustedRoots(t, chain)); err == nil {
		t.Fatal("validateAppleJWSChain() accepted a broken chain without signedDate")
	}
}

func testAppleCertificateChain(t *testing.T) []string {
	return testAppleCertificateChainWithAppleExtensions(t, true)
}

func testAppleCertificateChainWithoutAppleExtensions(t *testing.T) []string {
	return testAppleCertificateChainWithAppleExtensions(t, false)
}

func testAppleCertificateChainWithAppleExtensions(t *testing.T, includeAppleExtensions bool) []string {
	return testAppleCertificateChainAt(t, time.Now(), includeAppleExtensions)
}

func testAppleCertificateChainAt(t *testing.T, validAt time.Time, includeAppleExtensions bool) []string {
	t.Helper()
	rootKey := testAppleECDSAKey(t)
	rootTemplate := &x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{CommonName: "test Apple root"},
		NotBefore:             validAt.Add(-time.Hour),
		NotAfter:              validAt.Add(time.Hour),
		IsCA:                  true,
		BasicConstraintsValid: true,
		KeyUsage:              x509.KeyUsageCertSign,
	}
	rootDER := testAppleCertificate(t, rootTemplate, rootTemplate, &rootKey.PublicKey, rootKey)
	root, err := x509.ParseCertificate(rootDER)
	if err != nil {
		t.Fatal(err)
	}

	intermediateKey := testAppleECDSAKey(t)
	intermediateTemplate := &x509.Certificate{
		SerialNumber:          big.NewInt(2),
		Subject:               pkix.Name{CommonName: "test Apple intermediate"},
		NotBefore:             validAt.Add(-time.Hour),
		NotAfter:              validAt.Add(time.Hour),
		IsCA:                  true,
		BasicConstraintsValid: true,
		KeyUsage:              x509.KeyUsageCertSign,
	}
	if includeAppleExtensions {
		intermediateTemplate.ExtraExtensions = []pkix.Extension{{
			Id:    asn1.ObjectIdentifier{1, 2, 840, 113635, 100, 6, 2, 1},
			Value: []byte{0x05, 0x00},
		}}
	}
	intermediateDER := testAppleCertificate(t, intermediateTemplate, root, &intermediateKey.PublicKey, rootKey)
	intermediate, err := x509.ParseCertificate(intermediateDER)
	if err != nil {
		t.Fatal(err)
	}

	leafKey := testAppleECDSAKey(t)
	leafTemplate := &x509.Certificate{
		SerialNumber:          big.NewInt(3),
		Subject:               pkix.Name{CommonName: "test Apple signing leaf"},
		NotBefore:             validAt.Add(-time.Hour),
		NotAfter:              validAt.Add(time.Hour),
		BasicConstraintsValid: true,
		KeyUsage:              x509.KeyUsageDigitalSignature,
	}
	if includeAppleExtensions {
		leafTemplate.ExtraExtensions = []pkix.Extension{{
			Id:    asn1.ObjectIdentifier{1, 2, 840, 113635, 100, 6, 11, 1},
			Value: []byte{0x05, 0x00},
		}}
	}
	leafDER := testAppleCertificate(t, leafTemplate, intermediate, &leafKey.PublicKey, intermediateKey)
	chain := []string{
		base64.StdEncoding.EncodeToString(leafDER),
		base64.StdEncoding.EncodeToString(intermediateDER),
		base64.StdEncoding.EncodeToString(rootDER),
	}
	testAppleLeafKeys[chain[0]] = leafKey
	return chain
}

func testAppleECDSAKey(t *testing.T) *ecdsa.PrivateKey {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	return key
}

func testAppleCertificate(t *testing.T, template, parent *x509.Certificate, publicKey any, signer any) []byte {
	t.Helper()
	der, err := x509.CreateCertificate(rand.Reader, template, parent, publicKey, signer)
	if err != nil {
		t.Fatal(err)
	}
	return der
}

func testAppleJWS(t *testing.T, algorithm string, chain []string, signedAt time.Time) string {
	t.Helper()
	header, err := json.Marshal(map[string]any{"alg": algorithm, "x5c": chain})
	if err != nil {
		t.Fatal(err)
	}
	payload, err := json.Marshal(map[string]any{"signedDate": signedAt.UnixMilli()})
	if err != nil {
		t.Fatal(err)
	}
	return testAppleSignedJWS(t, header, payload, chain)
}

func testAppleJWSWithoutDate(t *testing.T, chain []string) string {
	t.Helper()
	header, err := json.Marshal(map[string]any{"alg": "ES256", "x5c": chain})
	if err != nil {
		t.Fatal(err)
	}
	payload, err := json.Marshal(map[string]any{"bundleId": "com.example"})
	if err != nil {
		t.Fatal(err)
	}
	return testAppleSignedJWS(t, header, payload, chain)
}

var testAppleLeafKeys = map[string]*ecdsa.PrivateKey{}

func testAppleSignedJWS(t *testing.T, header, payload []byte, chain []string) string {
	t.Helper()
	headerSegment := base64.RawURLEncoding.EncodeToString(header)
	payloadSegment := base64.RawURLEncoding.EncodeToString(payload)
	key := testAppleLeafKeys[chain[0]]
	if key == nil {
		t.Fatal("missing test Apple leaf key")
	}
	digest := sha256.Sum256([]byte(headerSegment + "." + payloadSegment))
	r, s, err := ecdsa.Sign(rand.Reader, key, digest[:])
	if err != nil {
		t.Fatal(err)
	}
	signature := make([]byte, 64)
	r.FillBytes(signature[:32])
	s.FillBytes(signature[32:])
	return headerSegment + "." + payloadSegment + "." + base64.RawURLEncoding.EncodeToString(signature)
}

func testAppleTrustedRoots(t *testing.T, chain []string) map[string]struct{} {
	t.Helper()
	rootDER, err := base64.StdEncoding.DecodeString(chain[2])
	if err != nil {
		t.Fatal(err)
	}
	fingerprint := sha256.Sum256(rootDER)
	return map[string]struct{}{hex.EncodeToString(fingerprint[:]): {}}
}