package payment

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

func TestValidateAppleJWSChain(t *testing.T) {
	chain := testAppleCertificateChain(t)
	if err := validateAppleJWSChain(testAppleJWS(t, "ES256", chain, time.Now())); 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 := validateAppleJWSChain(testAppleJWS(t, "ES256", chain, time.Now())); err == nil {
		t.Fatal("validateAppleJWSChain() accepted a leaf outside the declared chain")
	}
}

func TestValidateAppleJWSChainRejectsUnexpectedAlgorithmAndShape(t *testing.T) {
	chain := testAppleCertificateChain(t)
	if err := validateAppleJWSChain(testAppleJWS(t, "ES384", chain, time.Now())); err == nil {
		t.Fatal("validateAppleJWSChain() accepted a non-ES256 algorithm")
	}
	if err := validateAppleJWSChain(testAppleJWS(t, "ES256", chain[:2], time.Now())); 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 := validateAppleJWSChain(testAppleJWS(t, "ES256", chain, time.Now())); err == nil {
		t.Fatal("validateAppleJWSChain() accepted additional x5c certificates")
	}
}

func TestValidateAppleJWSChainRejectsNonAppleSigningCertificates(t *testing.T) {
	chain := testAppleCertificateChainWithoutAppleExtensions(t)
	if err := validateAppleJWSChain(testAppleJWS(t, "ES256", chain, time.Now())); 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 := validateAppleJWSChain(testAppleJWS(t, "ES256", chain, signedAt)); 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 := base64.RawURLEncoding.EncodeToString(header) + "." + base64.RawURLEncoding.EncodeToString(payload) + ".signature"
	if err := validateAppleJWSChain(outer); err != nil {
		t.Fatalf("validateAppleJWSChain() rejected historical notification: %v", err)
	}
}

func TestValidateAppleJWSChainAcceptsExpiredChainWithoutSignedDate(t *testing.T) {
	signedAt := time.Now().AddDate(-2, 0, 0)
	chain := testAppleCertificateChainAt(t, signedAt, true)
	if err := validateAppleJWSChain(testAppleJWSWithoutDate(t, chain)); err != nil {
		t.Fatalf("validateAppleJWSChain() rejected an expired chain without signedDate: %v", err)
	}
}

func TestValidateAppleJWSChainRejectsBrokenChainWithoutSignedDate(t *testing.T) {
	chain := testAppleCertificateChain(t)
	attackerChain := testAppleCertificateChain(t)
	chain[0] = attackerChain[0]
	if err := validateAppleJWSChain(testAppleJWSWithoutDate(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)
	return []string{
		base64.StdEncoding.EncodeToString(leafDER),
		base64.StdEncoding.EncodeToString(intermediateDER),
		base64.StdEncoding.EncodeToString(rootDER),
	}
}

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 base64.RawURLEncoding.EncodeToString(header) + "." + base64.RawURLEncoding.EncodeToString(payload) + ".signature"
}

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 base64.RawURLEncoding.EncodeToString(header) + "." + base64.RawURLEncoding.EncodeToString(payload) + ".signature"
}
