package apple

import (
	"crypto/ecdsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"math/big"
	"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 decodeB64URL(s string) ([]byte, error) {
	s = cleanB64(s)
	if b, err := base64.RawURLEncoding.DecodeString(s); err == nil {
		return b, nil
	}
	switch len(s) % 4 {
	case 2:
		s += "=="
	case 3:
		s += "="
	}
	return base64.URLEncoding.DecodeString(s)
}

func ParseTransactionJWS(jws string) (*TransactionPayload, error) {
	parts := strings.Split(strings.TrimSpace(jws), ".")
	if len(parts) != 3 {
		return nil, ErrInvalidJWS
	}
	data, err := decodeB64URL(parts[1])
	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["type"].(string); ok {
		resp.Type = 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
	}
	if v, ok := raw["appAccountToken"].(string); ok {
		resp.AppAccountToken = v
	}
	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
	}
	var hdr jwsHeader
	hdrBytes, err := decodeB64URL(parts[0])
	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])
	sigBytes, err := decodeB64URL(parts[2])
	if err != nil {
		return nil, err
	}
	d := sha256.Sum256([]byte(signingInput))
	// Try ASN.1 signature first
	if ecdsa.VerifyASN1(pub, d[:], sigBytes) {
		return ParseTransactionJWS(jws)
	}
	// Fallback: raw R||S (JWS ES256 uses raw signature)
	if len(sigBytes) == 64 {
		r := new(big.Int).SetBytes(sigBytes[:32])
		s := new(big.Int).SetBytes(sigBytes[32:])
		if ecdsa.Verify(pub, d[:], r, s) {
			return ParseTransactionJWS(jws)
		}
	}
	return nil, ErrInvalidJWS
}
