package apple

import (
	"crypto/ecdsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"strings"
	"time"
	"unicode"
)

func cleanB64(s string) string {
	trimmed := strings.TrimSpace(s)
	return strings.Map(func(r rune) rune {
		if unicode.IsSpace(r) {
			return -1
		}
		return r
	}, trimmed)
}

func ParseTransactionJWS(jws string) (*TransactionPayload, error) {
	parts := strings.Split(strings.TrimSpace(jws), ".")
	if len(parts) != 3 {
		return nil, ErrInvalidJWS
	}
	payloadB64 := cleanB64(parts[1])
	// add padding if required
	switch len(payloadB64) % 4 {
	case 2:
		payloadB64 += "=="
	case 3:
		payloadB64 += "="
	}
	data, err := base64.RawURLEncoding.DecodeString(payloadB64)
	if err != nil {
		return nil, err
	}
	var raw map[string]interface{}
	if err = json.Unmarshal(data, &raw); err != nil {
		return nil, err
	}
	var resp TransactionPayload
	if v, ok := raw["bundleId"].(string); ok {
		resp.BundleId = v
	}
	if v, ok := raw["productId"].(string); ok {
		resp.ProductId = v
	}
	if v, ok := raw["transactionId"].(string); ok {
		resp.TransactionId = v
	}
	if v, ok := raw["originalTransactionId"].(string); ok {
		resp.OriginalTransactionId = v
	}
	if v, ok := raw["purchaseDate"].(float64); ok {
		resp.PurchaseDate = time.UnixMilli(int64(v))
	} else if v, ok := raw["purchaseDate"].(int64); ok {
		resp.PurchaseDate = time.UnixMilli(v)
	}
	if v, ok := raw["revocationDate"].(float64); ok {
		t := time.UnixMilli(int64(v))
		resp.RevocationDate = &t
	}
	return &resp, nil
}

type jwsHeader struct {
	Alg string   `json:"alg"`
	X5c []string `json:"x5c"`
	Kid string   `json:"kid"`
}

func VerifyTransactionJWS(jws string) (*TransactionPayload, error) {
	parts := strings.Split(strings.TrimSpace(jws), ".")
	if len(parts) != 3 {
		return nil, ErrInvalidJWS
	}
	hdrB64 := cleanB64(parts[0])
	switch len(hdrB64) % 4 {
	case 2:
		hdrB64 += "=="
	case 3:
		hdrB64 += "="
	}
	var hdr jwsHeader
	hdrBytes, err := base64.RawURLEncoding.DecodeString(hdrB64)
	if err != nil {
		return nil, err
	}
	if err = json.Unmarshal(hdrBytes, &hdr); err != nil {
		return nil, err
	}
	if len(hdr.X5c) == 0 {
		return nil, ErrInvalidJWS
	}
	certDer, err := base64.StdEncoding.DecodeString(hdr.X5c[0])
	if err != nil {
		return nil, err
	}
	cert, err := x509.ParseCertificate(certDer)
	if err != nil {
		return nil, err
	}
	pub, ok := cert.PublicKey.(*ecdsa.PublicKey)
	if !ok {
		return nil, ErrInvalidJWS
	}
	signingInput := cleanB64(parts[0]) + "." + cleanB64(parts[1])
	sig := cleanB64(parts[2])
	sigBytes, err := base64.RawURLEncoding.DecodeString(sig)
	if err != nil {
		return nil, err
	}
	d := sha256.Sum256([]byte(signingInput))
	if !ecdsa.VerifyASN1(pub, d[:], sigBytes) {
		return nil, ErrInvalidJWS
	}
	return ParseTransactionJWS(jws)
}