// Package token handles JWT issuance, validation, and revocation for MCIAS.
//
// Security design:
//   - Algorithm header is checked FIRST, before any signature verification.
//     This prevents algorithm-confusion attacks (CVE-2022-21449 class).
//   - Only "EdDSA" is accepted; "none", HS*, RS*, ES* are all rejected.
//   - The signing key is taken from the server's keystore, never from the token.
//   - All standard claims (exp, iat, iss, jti) are required and validated.
//   - JTIs are UUIDs generated from crypto/rand (via google/uuid).
//   - Token values are never stored; only JTIs are recorded for revocation.
package token

import (
	"crypto/ed25519"
	"errors"
	"fmt"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
)

const (
	// requiredAlg is the only JWT algorithm accepted by MCIAS.
	// Security: Hard-coding this as a constant rather than a variable ensures
	// it cannot be changed at runtime and cannot be confused by token headers.
	requiredAlg = "EdDSA"
)

// Claims holds the MCIAS-specific JWT claims.
type Claims struct {
	// Standard registered claims.
	Issuer    string    `json:"iss"`
	Subject   string    `json:"sub"` // account UUID
	IssuedAt  time.Time `json:"iat"`
	ExpiresAt time.Time `json:"exp"`
	JTI       string    `json:"jti"`

	// MCIAS-specific claims.
	Roles []string `json:"roles"`
}

// jwtClaims adapts Claims to the golang-jwt MapClaims interface.
type jwtClaims struct {
	jwt.RegisteredClaims
	Roles []string `json:"roles"`
}

// ErrExpiredToken is returned when the token's exp claim is in the past.
var ErrExpiredToken = errors.New("token: expired")

// ErrInvalidSignature is returned when Ed25519 signature verification fails.
var ErrInvalidSignature = errors.New("token: invalid signature")

// ErrWrongAlgorithm is returned when the alg header is not EdDSA.
var ErrWrongAlgorithm = errors.New("token: algorithm must be EdDSA")

// ErrMissingClaim is returned when a required claim is absent or empty.
var ErrMissingClaim = errors.New("token: missing required claim")

// IssueToken creates and signs a new JWT with the given claims.
// The jti is generated automatically using crypto/rand via uuid.New().
// Returns the signed token string.
//
// Security: The signing key is provided by the caller from the server's
// keystore. The alg header is set explicitly to "EdDSA" by the jwt library
// when an ed25519.PrivateKey is passed to SignedString.
func IssueToken(key ed25519.PrivateKey, issuer, subject string, roles []string, expiry time.Duration) (string, *Claims, error) {
	now := time.Now().UTC()
	exp := now.Add(expiry)
	jti := uuid.New().String()

	jc := jwtClaims{
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    issuer,
			Subject:   subject,
			IssuedAt:  jwt.NewNumericDate(now),
			ExpiresAt: jwt.NewNumericDate(exp),
			ID:        jti,
		},
		Roles: roles,
	}

	t := jwt.NewWithClaims(jwt.SigningMethodEdDSA, jc)
	signed, err := t.SignedString(key)
	if err != nil {
		return "", nil, fmt.Errorf("token: sign JWT: %w", err)
	}

	claims := &Claims{
		Issuer:    issuer,
		Subject:   subject,
		IssuedAt:  now,
		ExpiresAt: exp,
		JTI:       jti,
		Roles:     roles,
	}
	return signed, claims, nil
}

// ValidateToken parses and validates a JWT string.
//
// Security order of operations (all must pass):
// 1. Parse the token header and extract the alg field.
// 2. Reject immediately if alg != "EdDSA" (before any signature check).
// 3. Verify Ed25519 signature.
// 4. Validate exp, iat, iss, jti claims.
//
// Returns Claims on success, or a typed error on any failure.
// The caller is responsible for checking revocation status via the DB.
func ValidateToken(key ed25519.PublicKey, tokenString, expectedIssuer string) (*Claims, error) {
	// Step 1+2: Parse the header to check alg BEFORE any crypto.
	// Security: We use jwt.ParseWithClaims with an explicit key function that
	// enforces the algorithm. The key function is called by the library after
	// parsing the header but before verifying the signature, which is the
	// correct point to enforce algorithm constraints.
	var jc jwtClaims
	t, err := jwt.ParseWithClaims(tokenString, &jc, func(t *jwt.Token) (interface{}, error) {
		// Security: Check alg header first. This must happen in the key
		// function — it is the only place where the parsed (but unverified)
		// header is available before signature validation.
		if t.Method.Alg() != requiredAlg {
			return nil, fmt.Errorf("%w: got %q, want %q", ErrWrongAlgorithm, t.Method.Alg(), requiredAlg)
		}
		return key, nil
	},
		jwt.WithIssuedAt(),
		jwt.WithIssuer(expectedIssuer),
		jwt.WithExpirationRequired(),
	)
	if err != nil {
		// Map library errors to our typed errors for consistent handling.
		if errors.Is(err, ErrWrongAlgorithm) {
			return nil, ErrWrongAlgorithm
		}
		if errors.Is(err, jwt.ErrTokenExpired) {
			return nil, ErrExpiredToken
		}
		if errors.Is(err, jwt.ErrSignatureInvalid) {
			return nil, ErrInvalidSignature
		}
		return nil, fmt.Errorf("token: parse: %w", err)
	}
	if !t.Valid {
		return nil, ErrInvalidSignature
	}

	// Step 4: Validate required custom claims.
	if jc.ID == "" {
		return nil, fmt.Errorf("%w: jti", ErrMissingClaim)
	}
	if jc.Subject == "" {
		return nil, fmt.Errorf("%w: sub", ErrMissingClaim)
	}
	if jc.ExpiresAt == nil {
		return nil, fmt.Errorf("%w: exp", ErrMissingClaim)
	}
	if jc.IssuedAt == nil {
		return nil, fmt.Errorf("%w: iat", ErrMissingClaim)
	}

	claims := &Claims{
		Issuer:    jc.Issuer,
		Subject:   jc.Subject,
		IssuedAt:  jc.IssuedAt.Time,
		ExpiresAt: jc.ExpiresAt.Time,
		JTI:       jc.ID,
		Roles:     jc.Roles,
	}
	return claims, nil
}

// HasRole reports whether the claims include the given role.
func (c *Claims) HasRole(role string) bool {
	for _, r := range c.Roles {
		if r == role {
			return true
		}
	}
	return false
}