package appleiap

import (
    "crypto/ecdsa"
    "crypto/elliptic"
    "encoding/base64"
    "encoding/json"
    "errors"
    "math/big"
    "net/http"
    "sync"
    "time"
)

type jwk struct {
    Kty string `json:"kty"`
    Kid string `json:"kid"`
    Crv string `json:"crv"`
    X   string `json:"x"`
    Y   string `json:"y"`
}

type jwks struct {
    Keys []jwk `json:"keys"`
}

type cacheEntry struct {
    keys map[string]*ecdsa.PublicKey
    exp  time.Time
}

var (
    mu    sync.Mutex
    cache = map[string]*cacheEntry{}
)

func endpoint(env string) string {
    if env == "sandbox" {
        return "https://api.storekit-sandbox.itunes.apple.com/inApps/v1/keys"
    }
    return "https://api.storekit.itunes.apple.com/inApps/v1/keys"
}

func fetch(env string) (map[string]*ecdsa.PublicKey, error) {
    resp, err := http.Get(endpoint(env))
    if err != nil {
        return nil, err
    }
    defer resp.Body.Close()
    var set jwks
    if err := json.NewDecoder(resp.Body).Decode(&set); err != nil {
        return nil, err
    }
    m := make(map[string]*ecdsa.PublicKey)
    for _, k := range set.Keys {
        if k.Kty != "EC" || k.Crv != "P-256" || k.X == "" || k.Y == "" || k.Kid == "" {
            continue
        }
        xb, err := base64.RawURLEncoding.DecodeString(k.X)
        if err != nil { continue }
        yb, err := base64.RawURLEncoding.DecodeString(k.Y)
        if err != nil { continue }
        var x, y big.Int
        x.SetBytes(xb)
        y.SetBytes(yb)
        m[k.Kid] = &ecdsa.PublicKey{Curve: elliptic.P256(), X: &x, Y: &y}
    }
    if len(m) == 0 {
        return nil, errors.New("empty jwks")
    }
    return m, nil
}

func GetKey(env, kid string) (*ecdsa.PublicKey, error) {
    mu.Lock()
    defer mu.Unlock()
    c := cache[env]
    if c == nil || time.Now().After(c.exp) {
        keys, err := fetch(env)
        if err != nil { return nil, err }
        cache[env] = &cacheEntry{ keys: keys, exp: time.Now().Add(10 * time.Minute) }
        c = cache[env]
    }
    k := c.keys[kid]
    if k == nil { return nil, errors.New("key not found") }
    return k, nil
}