package workload

import (
	"crypto/hmac"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"time"
)

// SignedWorkloadToken is a control-plane-issued bearer credential for the
// WorkloadVerify boundary (see multiplayer-next.md 8.10). It exists because
// the obvious approach -- verifying a Kubernetes-projected service-account
// JWT via TokenReview/JWKS (see jwt.go, ParseAndValidate) -- needs a live
// cluster to validate against and so cannot be built or tested here.
//
// This sidesteps that requirement entirely: the control plane signs its own
// short-lived token over (allocation_id, expiry) with a secret only it
// holds, exactly the way domain.SessionStore already mints player session
// tokens elsewhere in this codebase. It needs no Kubernetes trust boundary
// to verify -- HMAC signature plus expiry is self-contained.
//
// The token deliberately binds ONLY allocation_id, not match_id/server_id
// too: it is meant to be requested as a GameServerAllocation annotation
// (see agones/allocation.go) in the SAME request that asks Agones to pick a
// server for this allocation -- so at mint time, the allocator knows
// allocation_id (it generates it) but not yet which server_id Agones will
// return. match_id and server_id are instead resolved durably at verify
// time from the allocations table, which the allocator records immediately
// after Agones responds (see store.AllocationBindingByAllocationID) -- so a
// token can never claim a match/server pairing that isn't what was actually,
// durably allocated.
//
// The delivery channel is what makes this safe despite not proving pod
// identity the way a Kubernetes-issued token would: the token reaches the
// allocated GameServer via the same annotation channel allocation.go
// already uses for match-id/allocation-id, which only the actually-
// allocated pod's local SDK sidecar can read. A caller who can present this
// token has already proven, via that channel, that it is the pod Agones
// allocated.
type SignedWorkloadToken struct {
	AllocationID string    `json:"a"`
	ExpiresAt    time.Time `json:"e"`
}

var (
	ErrEmptyWorkloadSecret = errors.New("workload token signing secret is empty")
	ErrMalformedToken      = errors.New("malformed signed workload token")
	ErrTokenSignature      = errors.New("signed workload token signature mismatch")
	ErrTokenExpired        = errors.New("signed workload token expired")
	ErrTokenClaims         = errors.New("signed workload token missing required claims")
)

// IssueSignedWorkloadToken produces a compact "payload.signature" token
// binding allocation_id, the one identifier known at mint time (see the
// type doc above for why match_id/server_id aren't embedded). now must be
// non-zero and ttl must be positive so a token is never silently issued
// already-expired.
func IssueSignedWorkloadToken(secret []byte, allocationID string, now time.Time, ttl time.Duration) (string, error) {
	if len(secret) == 0 {
		return "", ErrEmptyWorkloadSecret
	}
	if allocationID == "" {
		return "", ErrTokenClaims
	}
	if now.IsZero() || ttl <= 0 {
		return "", fmt.Errorf("issue signed workload token: now and ttl must be valid")
	}
	claims := SignedWorkloadToken{
		AllocationID: allocationID,
		ExpiresAt:    now.Add(ttl).UTC(),
	}
	payload, err := json.Marshal(claims)
	if err != nil {
		return "", fmt.Errorf("marshal signed workload token: %w", err)
	}
	payloadEnc := base64.RawURLEncoding.EncodeToString(payload)
	mac := hmac.New(sha256.New, secret)
	mac.Write([]byte(payloadEnc))
	sigEnc := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
	return payloadEnc + "." + sigEnc, nil
}

// ParseSignedWorkloadToken verifies the signature in constant time, checks
// expiry against now, and returns the claims. It never trusts the payload
// before the signature is verified.
func ParseSignedWorkloadToken(secret []byte, token string, now time.Time) (SignedWorkloadToken, error) {
	if len(secret) == 0 {
		return SignedWorkloadToken{}, ErrEmptyWorkloadSecret
	}
	dot := -1
	for i := 0; i < len(token); i++ {
		if token[i] == '.' {
			dot = i
			break
		}
	}
	if dot <= 0 || dot == len(token)-1 {
		return SignedWorkloadToken{}, ErrMalformedToken
	}
	payloadEnc, sigEnc := token[:dot], token[dot+1:]
	mac := hmac.New(sha256.New, secret)
	mac.Write([]byte(payloadEnc))
	expectedSig := mac.Sum(nil)
	gotSig, err := base64.RawURLEncoding.DecodeString(sigEnc)
	if err != nil {
		return SignedWorkloadToken{}, ErrMalformedToken
	}
	if subtle.ConstantTimeCompare(expectedSig, gotSig) != 1 {
		return SignedWorkloadToken{}, ErrTokenSignature
	}
	payload, err := base64.RawURLEncoding.DecodeString(payloadEnc)
	if err != nil {
		return SignedWorkloadToken{}, ErrMalformedToken
	}
	var claims SignedWorkloadToken
	if err := json.Unmarshal(payload, &claims); err != nil {
		return SignedWorkloadToken{}, ErrMalformedToken
	}
	if claims.AllocationID == "" || claims.ExpiresAt.IsZero() {
		return SignedWorkloadToken{}, ErrTokenClaims
	}
	if now.IsZero() {
		return SignedWorkloadToken{}, fmt.Errorf("parse signed workload token: now must be valid")
	}
	if !now.Before(claims.ExpiresAt) {
		return SignedWorkloadToken{}, ErrTokenExpired
	}
	return claims, nil
}