// Package token checks the hub's tokens: EdDSA (Ed25519) JWTs in compact form, signed with a key from the hub's
// JWKS (GET <hub>/.well-known/t5de-keys.json). See t5de-hub's docs/products.md, "Tokens".
package token

import (
	"context"
	"crypto/ed25519"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"
)

// Token types.
const (
	Upload   = "t5de-upload"
	Access   = "t5de-access"
	Revision = "t5de-revision"
)

// KeysPath is where the hub publishes its keys.
const KeysPath = "/.well-known/t5de-keys.json"

// Leeway allows for clocks that disagree a little.
const Leeway = time.Minute

// RefetchInterval is the least time between fetches of the hub's keys, however many unknown key IDs show up.
const RefetchInterval = 30 * time.Second

const maxToken = 16 << 10
const maxKeys = 1 << 20

// ErrInvalid wraps every reason a token is refused.
var ErrInvalid = errors.New("invalid token")

// Claims are the claims T5DE's tokens carry.
type Claims struct {
	Issuer   string `json:"iss"`
	Type     string `json:"typ"`
	Subject  string `json:"sub"`
	IssuedAt *int64 `json:"iat"`
	Expires  *int64 `json:"exp"`
	NotYet   *int64 `json:"nbf"`
	Catalog  string `json:"catalog"`  // t5de-upload
	Product  string `json:"product"`  // t5de-upload (optional), t5de-revision
	Adult    bool   `json:"adult"`    // t5de-access
	Revision int    `json:"revision"` // t5de-revision
	Rating   string `json:"rating"`   // t5de-revision
}

// Verifier checks tokens signed by one hub.
type Verifier struct {
	issuer string
	keyURL string
	client *http.Client
	now    func() time.Time

	fetch   sync.Mutex // held while fetching, so one fetch serves everyone waiting
	mu      sync.Mutex
	keys    map[string]ed25519.PublicKey
	fetched time.Time // when the keys were last fetched (or tried), zero if never
}

// NewVerifier returns a verifier for tokens from the hub at hubURL (its API's address, which is also the tokens'
// issuer). client may be nil.
func NewVerifier(hubURL string, client *http.Client) *Verifier {
	hubURL = strings.TrimRight(hubURL, "/")
	if client == nil {
		client = &http.Client{Timeout: 10 * time.Second}
	}
	return &Verifier{issuer: hubURL, keyURL: hubURL + KeysPath, client: client, now: time.Now}
}

// SetClock replaces the verifier's clock, for tests.
func (v *Verifier) SetClock(now func() time.Time) { v.now = now }

type header struct {
	Alg  string `json:"alg"`
	Kid  string `json:"kid"`
	Crit any    `json:"crit"`
}

// Verify checks a token of type typ: its signature, issuer, type and times. Tokens of other types than
// t5de-revision must have an expiry.
func (v *Verifier) Verify(ctx context.Context, token, typ string) (*Claims, error) {
	c, err := v.verify(ctx, token, typ)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrInvalid, err)
	}
	return c, nil
}

func (v *Verifier) verify(ctx context.Context, token, typ string) (*Claims, error) {
	if len(token) > maxToken {
		return nil, errors.New("too long")
	}
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return nil, errors.New("not a compact JWT")
	}
	var h header
	if err := decodePart(parts[0], &h); err != nil {
		return nil, fmt.Errorf("header: %v", err)
	}
	if h.Alg != "EdDSA" {
		return nil, fmt.Errorf("algorithm %q isn't EdDSA", h.Alg)
	}
	if h.Crit != nil {
		return nil, errors.New("unsupported critical header")
	}
	sig, err := base64.RawURLEncoding.Strict().DecodeString(parts[2])
	if err != nil || len(sig) != ed25519.SignatureSize {
		return nil, errors.New("bad signature encoding")
	}
	key, err := v.key(ctx, h.Kid)
	if err != nil {
		return nil, err
	}
	if !ed25519.Verify(key, []byte(parts[0]+"."+parts[1]), sig) {
		return nil, errors.New("bad signature")
	}

	var c Claims
	if err := decodePart(parts[1], &c); err != nil {
		return nil, fmt.Errorf("claims: %v", err)
	}
	if c.Issuer != v.issuer {
		return nil, fmt.Errorf("issuer %q isn't the hub", c.Issuer)
	}
	if c.Type != typ {
		return nil, fmt.Errorf("type %q isn't %q", c.Type, typ)
	}
	now := v.now()
	if c.Expires == nil {
		if typ != Revision {
			return nil, errors.New("no expiry")
		}
	} else if !now.Before(time.Unix(*c.Expires, 0).Add(Leeway)) {
		return nil, errors.New("expired")
	}
	if c.NotYet != nil && now.Add(Leeway).Before(time.Unix(*c.NotYet, 0)) {
		return nil, errors.New("not valid yet")
	}
	if c.IssuedAt != nil && now.Add(Leeway).Before(time.Unix(*c.IssuedAt, 0)) {
		return nil, errors.New("issued in the future")
	}
	return &c, nil
}

func decodePart(s string, v any) error {
	data, err := base64.RawURLEncoding.Strict().DecodeString(s)
	if err != nil {
		return err
	}
	return json.Unmarshal(data, v)
}

// key returns the key with ID kid, fetching the hub's keys if it isn't known and they weren't fetched lately.
func (v *Verifier) key(ctx context.Context, kid string) (ed25519.PublicKey, error) {
	if k, ok := v.cached(kid); ok {
		return k, nil
	}
	v.fetch.Lock()
	defer v.fetch.Unlock()
	if k, ok := v.cached(kid); ok { // fetched while waiting
		return k, nil
	}
	v.mu.Lock()
	recent := !v.fetched.IsZero() && v.now().Sub(v.fetched) < RefetchInterval
	v.mu.Unlock()
	if recent {
		return nil, fmt.Errorf("unknown key %q", kid)
	}
	keys, err := v.fetchKeys(ctx)
	v.mu.Lock()
	v.fetched = v.now()
	if err == nil {
		v.keys = keys
	}
	v.mu.Unlock()
	if err != nil {
		return nil, fmt.Errorf("fetching the hub's keys: %v", err)
	}
	if k, ok := keys[kid]; ok {
		return k, nil
	}
	return nil, fmt.Errorf("unknown key %q", kid)
}

func (v *Verifier) cached(kid string) (ed25519.PublicKey, bool) {
	v.mu.Lock()
	defer v.mu.Unlock()
	k, ok := v.keys[kid]
	return k, ok
}

type jwk struct {
	Kty string `json:"kty"`
	Crv string `json:"crv"`
	X   string `json:"x"`
	Kid string `json:"kid"`
	Alg string `json:"alg"`
	Use string `json:"use"`
}

func (v *Verifier) fetchKeys(ctx context.Context) (map[string]ed25519.PublicKey, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, v.keyURL, nil)
	if err != nil {
		return nil, err
	}
	resp, err := v.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("%s: %s", v.keyURL, resp.Status)
	}
	data, err := io.ReadAll(io.LimitReader(resp.Body, maxKeys+1))
	if err != nil {
		return nil, err
	}
	if len(data) > maxKeys {
		return nil, errors.New("key set too large")
	}
	return ParseKeys(data)
}

// ParseKeys reads the Ed25519 signing keys from a JWKS. Other keys are skipped.
func ParseKeys(data []byte) (map[string]ed25519.PublicKey, error) {
	var set struct {
		Keys []jwk `json:"keys"`
	}
	if err := json.Unmarshal(data, &set); err != nil {
		return nil, err
	}
	keys := map[string]ed25519.PublicKey{}
	for _, k := range set.Keys {
		if k.Kty != "OKP" || k.Crv != "Ed25519" || k.Kid == "" {
			continue
		}
		if (k.Alg != "" && k.Alg != "EdDSA") || (k.Use != "" && k.Use != "sig") {
			continue
		}
		x, err := base64.RawURLEncoding.Strict().DecodeString(k.X)
		if err != nil || len(x) != ed25519.PublicKeySize {
			continue
		}
		keys[k.Kid] = ed25519.PublicKey(x)
	}
	return keys, nil
}