package auth

import (
	"context"
	"crypto/ed25519"
	"crypto/rsa"
	"crypto/x509"
	"errors"
	"sync"
	"time"

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

// JWTKey binds a trusted key ID to exactly one asymmetric signing algorithm.
// Keys come from server configuration, never a token's jku/x5u URL.
type JWTKey struct {
	Algorithm string
	PublicKey any
}
type JWTConfig struct {
	Issuer   string            `json:"issuer"`
	Audience string            `json:"audience"`
	Leeway   time.Duration     `json:"leeway"`
	Keys     map[string]JWTKey `json:"-"`
}
type JWT struct {
	config JWTConfig
	mu     sync.RWMutex
	keys   map[string]JWTKey
}
type JWTClaims struct {
	jwt.RegisteredClaims
	TenantID string   `json:"tenantId,omitempty"`
	Roles    []string `json:"roles,omitempty"`
}

func NewJWT(cfg JWTConfig) (*JWT, error) {
	if cfg.Issuer == "" || cfg.Audience == "" || cfg.Leeway < 0 || cfg.Leeway > 2*time.Minute {
		return nil, errors.New("auth: issuer/audience required, leeway must be 0..2m")
	}
	j := &JWT{config: cfg}
	j.config.Keys = nil
	if err := j.ReplaceKeys(cfg.Keys); err != nil {
		return nil, err
	}
	return j, nil
}

// ReplaceKeys atomically rotates the trusted key set. Retain old keys until
// their tokens expire if a rotation should preserve existing sessions.
func (j *JWT) ReplaceKeys(keys map[string]JWTKey) error {
	if len(keys) == 0 || len(keys) > 64 {
		return errors.New("auth: require 1..64 JWT keys")
	}
	copyKeys := make(map[string]JWTKey, len(keys))
	for id, key := range keys {
		if id == "" || len(id) > 128 {
			return errors.New("auth: invalid key ID")
		}
		switch key.Algorithm {
		case "RS256":
			k, ok := key.PublicKey.(*rsa.PublicKey)
			if !ok || k == nil || k.N == nil || k.N.BitLen() < 2048 || k.E < 3 {
				return errors.New("auth: RS256 requires a >=2048-bit RSA public key")
			}
		case "EdDSA":
			k, ok := key.PublicKey.(ed25519.PublicKey)
			if !ok || len(k) != ed25519.PublicKeySize {
				return errors.New("auth: EdDSA requires an Ed25519 public key")
			}
		default:
			return errors.New("auth: supported algorithms are RS256 and EdDSA")
		}
		encoded, err := x509.MarshalPKIXPublicKey(key.PublicKey)
		if err != nil {
			return err
		}
		key.PublicKey, err = x509.ParsePKIXPublicKey(encoded)
		if err != nil {
			return err
		}
		copyKeys[id] = key
	}
	j.mu.Lock()
	j.keys = copyKeys
	j.mu.Unlock()
	return nil
}

func (j *JWT) Authenticate(ctx context.Context, raw string) (Principal, error) {
	if ctx == nil {
		return Principal{}, ErrInvalidToken
	}
	if err := ctx.Err(); err != nil {
		return Principal{}, err
	}
	if len(raw) == 0 || len(raw) > 8192 {
		return Principal{}, ErrInvalidToken
	}
	claims := new(JWTClaims)
	token, err := jwt.ParseWithClaims(raw, claims, func(t *jwt.Token) (any, error) {
		id, ok := t.Header["kid"].(string)
		if !ok {
			return nil, ErrInvalidToken
		}
		j.mu.RLock()
		key, exists := j.keys[id]
		j.mu.RUnlock()
		if !exists || t.Method.Alg() != key.Algorithm {
			return nil, ErrInvalidToken
		}
		return key.PublicKey, nil
	}, jwt.WithValidMethods([]string{"RS256", "EdDSA"}), jwt.WithIssuer(j.config.Issuer), jwt.WithAudience(j.config.Audience), jwt.WithExpirationRequired(), jwt.WithIssuedAt(), jwt.WithLeeway(j.config.Leeway))
	if err != nil || !token.Valid || claims.Subject == "" || len(claims.Subject) > 256 || len(claims.TenantID) > 256 || len(claims.Roles) > 64 {
		return Principal{}, ErrInvalidToken
	}
	for _, role := range claims.Roles {
		if role == "" || len(role) > 128 {
			return Principal{}, ErrInvalidToken
		}
	}
	if err := ctx.Err(); err != nil {
		return Principal{}, err
	}
	return Principal{UserID: claims.Subject, TenantID: claims.TenantID, Roles: append([]string(nil), claims.Roles...)}, nil
}