package payment

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/sha256"
	"crypto/x509"
	"encoding/asn1"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"strings"
	"time"
)

const (
	maxAppleJWSHeaderSize  = 256 << 10
	maxAppleJWSPayloadSize = 1 << 20
	maxAppleJWSTimeDepth   = 2
)

var (
	appleJWSLeafExtensionOID         = asn1.ObjectIdentifier{1, 2, 840, 113635, 100, 6, 11, 1}
	appleJWSIntermediateExtensionOID = asn1.ObjectIdentifier{1, 2, 840, 113635, 100, 6, 2, 1}
	appleTrustedRootFingerprints     = map[string]struct{}{
		"63343abfb89a6a03ebb57e9b3f5fa7be7c4f5c756f3017b3a8c488c3653e9179": {},
	}
)

func validateAppleJWSChain(signedPayload string) error {
	return validateAppleJWSChainWithRoots(signedPayload, appleTrustedRootFingerprints)
}

func validateAppleJWSChainWithRoots(signedPayload string, trustedRoots map[string]struct{}) error {
	headerSegment, payloadSegment, signatureSegment, err := splitAppleJWS(signedPayload)
	if err != nil {
		return err
	}
	verificationTime, hasVerificationTime, err := appleJWSVerificationTime(payloadSegment)
	if err != nil {
		return err
	}
	if signatureSegment == "" {
		return errors.New("Apple JWS 格式无效")
	}
	if len(headerSegment) > base64.RawURLEncoding.EncodedLen(maxAppleJWSHeaderSize) {
		return errors.New("Apple JWS header 过大")
	}
	headerBytes, err := base64.RawURLEncoding.DecodeString(headerSegment)
	if err != nil {
		return fmt.Errorf("解析 Apple JWS header: %w", err)
	}
	if len(headerBytes) > maxAppleJWSHeaderSize {
		return errors.New("Apple JWS header 过大")
	}
	var header struct {
		Algorithm string   `json:"alg"`
		Chain     []string `json:"x5c"`
	}
	if err = json.Unmarshal(headerBytes, &header); err != nil {
		return fmt.Errorf("解析 Apple JWS header: %w", err)
	}
	if header.Algorithm != "ES256" {
		return errors.New("Apple JWS 算法必须是 ES256")
	}
	if len(header.Chain) != 3 {
		return errors.New("Apple JWS x5c 必须包含叶子、中间和根证书")
	}

	certificates := make([]*x509.Certificate, 0, len(header.Chain))
	for index, encoded := range header.Chain {
		der, decodeErr := base64.StdEncoding.DecodeString(encoded)
		if decodeErr != nil {
			return fmt.Errorf("解析 Apple JWS x5c[%d]: %w", index, decodeErr)
		}
		certificate, parseErr := x509.ParseCertificate(der)
		if parseErr != nil {
			return fmt.Errorf("解析 Apple JWS x5c[%d] 证书: %w", index, parseErr)
		}
		certificates = append(certificates, certificate)
	}
	rootFingerprint := sha256.Sum256(certificates[2].Raw)
	if _, trusted := trustedRoots[hex.EncodeToString(rootFingerprint[:])]; !trusted {
		return errors.New("Apple JWS 根证书不受信任")
	}
	if err = certificates[0].CheckSignatureFrom(certificates[1]); err != nil {
		return fmt.Errorf("Apple JWS 叶子证书不属于声明的证书链: %w", err)
	}
	if err = certificates[1].CheckSignatureFrom(certificates[2]); err != nil {
		return fmt.Errorf("Apple JWS 中间证书不属于声明的证书链: %w", err)
	}
	if err = certificates[2].CheckSignatureFrom(certificates[2]); err != nil {
		return fmt.Errorf("Apple JWS 根证书不是自签名证书: %w", err)
	}
	if !certificates[1].IsCA || !certificates[1].BasicConstraintsValid {
		return errors.New("Apple JWS 中间证书不是有效 CA")
	}
	if !certificates[2].IsCA || !certificates[2].BasicConstraintsValid {
		return errors.New("Apple JWS 根证书不是有效 CA")
	}
	if certificates[1].KeyUsage != 0 && certificates[1].KeyUsage&x509.KeyUsageCertSign == 0 {
		return errors.New("Apple JWS 中间证书不允许签发证书")
	}
	if certificates[2].KeyUsage != 0 && certificates[2].KeyUsage&x509.KeyUsageCertSign == 0 {
		return errors.New("Apple JWS 根证书不允许签发证书")
	}
	if !hasAppleCertificateExtension(certificates[0], appleJWSLeafExtensionOID) {
		return errors.New("Apple JWS 叶子证书缺少 App Store 签名扩展")
	}
	if !hasAppleCertificateExtension(certificates[1], appleJWSIntermediateExtensionOID) {
		return errors.New("Apple JWS 中间证书缺少 Apple 签名扩展")
	}

	if !hasVerificationTime {
		verificationTime = time.Now()
	}
	roots := x509.NewCertPool()
	roots.AddCert(certificates[2])
	intermediates := x509.NewCertPool()
	intermediates.AddCert(certificates[1])
	if _, err = certificates[0].Verify(x509.VerifyOptions{
		Roots:         roots,
		Intermediates: intermediates,
		KeyUsages:     []x509.ExtKeyUsage{x509.ExtKeyUsageAny},
		CurrentTime:   verificationTime,
	}); err != nil {
		return fmt.Errorf("验证 Apple JWS 证书链: %w", err)
	}
	if certificates[0].IsCA {
		return errors.New("Apple JWS 叶子证书不能是 CA")
	}
	if certificates[0].KeyUsage != 0 && certificates[0].KeyUsage&x509.KeyUsageDigitalSignature == 0 {
		return errors.New("Apple JWS 叶子证书不允许数字签名")
	}
	publicKey, ok := certificates[0].PublicKey.(*ecdsa.PublicKey)
	if !ok || publicKey.Curve != elliptic.P256() {
		return errors.New("Apple JWS 叶子证书不是 P-256 ECDSA 公钥")
	}
	signature, err := base64.RawURLEncoding.DecodeString(signatureSegment)
	if err != nil || len(signature) != 64 {
		return errors.New("Apple JWS ES256 签名格式无效")
	}
	digest := sha256.Sum256([]byte(headerSegment + "." + payloadSegment))
	if !ecdsa.Verify(publicKey, digest[:], new(big.Int).SetBytes(signature[:32]), new(big.Int).SetBytes(signature[32:])) {
		return errors.New("Apple JWS 签名校验失败")
	}
	return nil
}

func splitAppleJWS(signedPayload string) (header, payload, signature string, err error) {
	header, remainder, ok := strings.Cut(strings.TrimSpace(signedPayload), ".")
	if !ok || header == "" {
		return "", "", "", errors.New("Apple JWS 格式无效")
	}
	payload, signature, ok = strings.Cut(remainder, ".")
	if !ok || payload == "" || signature == "" || strings.Contains(signature, ".") {
		return "", "", "", errors.New("Apple JWS 格式无效")
	}
	return header, payload, signature, nil
}

func appleJWSVerificationTime(payloadSegment string) (time.Time, bool, error) {
	return appleJWSVerificationTimeAtDepth(payloadSegment, 0)
}

// appleJWSVerificationTimeAtDepth extracts a signing time from the current
// payload, then follows the nested signed transaction/renewal JWS used by
// App Store Server Notifications. The outer notification often has no date of
// its own, while its nested transaction does.
func appleJWSVerificationTimeAtDepth(payloadSegment string, depth int) (time.Time, bool, error) {
	if len(payloadSegment) > base64.RawURLEncoding.EncodedLen(maxAppleJWSPayloadSize) {
		return time.Time{}, false, errors.New("Apple JWS payload 过大")
	}
	payloadBytes, err := base64.RawURLEncoding.DecodeString(payloadSegment)
	if err != nil {
		return time.Time{}, false, fmt.Errorf("解析 Apple JWS payload: %w", err)
	}
	if len(payloadBytes) > maxAppleJWSPayloadSize {
		return time.Time{}, false, errors.New("Apple JWS payload 过大")
	}
	var payload struct {
		SignedDate          json.Number `json:"signedDate"`
		ReceiptCreationDate json.Number `json:"receiptCreationDate"`
		Data                struct {
			SignedTransactionInfo string `json:"signedTransactionInfo"`
			SignedRenewalInfo     string `json:"signedRenewalInfo"`
		} `json:"data"`
	}
	if err = json.Unmarshal(payloadBytes, &payload); err != nil {
		return time.Time{}, false, fmt.Errorf("解析 Apple JWS payload: %w", err)
	}
	date := payload.SignedDate
	if date == "" {
		date = payload.ReceiptCreationDate
	}
	if date != "" {
		milliseconds, parseErr := date.Int64()
		if parseErr != nil || milliseconds <= 0 {
			return time.Time{}, false, errors.New("Apple JWS signedDate 无效")
		}
		return time.UnixMilli(milliseconds), true, nil
	}
	if depth >= maxAppleJWSTimeDepth {
		return time.Time{}, false, nil
	}
	for _, nested := range []string{payload.Data.SignedTransactionInfo, payload.Data.SignedRenewalInfo} {
		if strings.TrimSpace(nested) == "" {
			continue
		}
		_, nestedPayload, _, splitErr := splitAppleJWS(nested)
		if splitErr != nil {
			continue
		}
		if nestedTime, found, nestedErr := appleJWSVerificationTimeAtDepth(nestedPayload, depth+1); nestedErr != nil {
			return time.Time{}, false, nestedErr
		} else if found {
			return nestedTime, true, nil
		}
	}
	return time.Time{}, false, nil
}

func hasAppleCertificateExtension(certificate *x509.Certificate, oid asn1.ObjectIdentifier) bool {
	if certificate == nil {
		return false
	}
	for _, extension := range certificate.Extensions {
		if extension.Id.Equal(oid) {
			return true
		}
	}
	return false
}