package auth

import (
	"context"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/big"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

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

const maxOIDCIDTokenBytes = 64 * 1024

type OIDCProviderConfig struct {
	Audience       string
	TrustedIssuers []string
	JWKSURL        string
	HTTPClient     *http.Client
	Now            func() time.Time
	ClockSkew      time.Duration
	MinimumKeyTTL  time.Duration
	DefaultKeyTTL  time.Duration
	MaximumKeyTTL  time.Duration
}

type OIDCIdentity struct {
	Subject string
	Email   string
}

type OIDCInvalidTokenError struct{}

func (*OIDCInvalidTokenError) Error() string   { return "invalid federated identity token" }
func (*OIDCInvalidTokenError) HTTPStatus() int { return http.StatusUnauthorized }

type OIDCProviderUnavailableError struct{ cause error }

func (e *OIDCProviderUnavailableError) Error() string {
	if e == nil || e.cause == nil {
		return "federated identity provider unavailable"
	}
	return "federated identity provider unavailable: " + e.cause.Error()
}
func (e *OIDCProviderUnavailableError) Unwrap() error { return e.cause }
func (*OIDCProviderUnavailableError) HTTPStatus() int { return http.StatusServiceUnavailable }

type oidcIDTokenClaims struct {
	Email           string `json:"email"`
	EmailVerified   bool   `json:"email_verified"`
	AuthorizedParty string `json:"azp"`
	jwt.RegisteredClaims
}

type OIDCIDTokenVerifier struct {
	audience   string
	issuers    map[string]struct{}
	jwksURL    string
	client     *http.Client
	now        func() time.Time
	skew       time.Duration
	minTTL     time.Duration
	defaultTTL time.Duration
	maxTTL     time.Duration

	mu        sync.Mutex
	keys      map[string]*rsa.PublicKey
	expiresAt time.Time
}

func NewOIDCIDTokenVerifier(config OIDCProviderConfig) (*OIDCIDTokenVerifier, error) {
	audience := strings.TrimSpace(config.Audience)
	jwksURL := strings.TrimSpace(config.JWKSURL)
	if audience == "" || audience != config.Audience || jwksURL == "" || jwksURL != config.JWKSURL {
		return nil, errors.New("OIDC verifier: canonical audience and JWKS URL are required")
	}
	issuers := make(map[string]struct{}, len(config.TrustedIssuers))
	for _, raw := range config.TrustedIssuers {
		issuer := strings.TrimSpace(raw)
		if issuer == "" || issuer != raw {
			return nil, errors.New("OIDC verifier: trusted issuer is invalid")
		}
		issuers[issuer] = struct{}{}
	}
	if len(issuers) == 0 {
		return nil, errors.New("OIDC verifier: at least one trusted issuer is required")
	}
	if config.HTTPClient == nil {
		config.HTTPClient = &http.Client{Timeout: 10 * time.Second}
	}
	// The operator-selected JWKS URL is the trust anchor. Go's default client
	// follows redirects, including a possible HTTPS-to-HTTP or cross-origin
	// redirect, which would let an endpoint response replace that authority.
	// Clone the caller's client so the verifier can reject redirects without
	// mutating a shared client used by another owner.
	httpClient := *config.HTTPClient
	httpClient.CheckRedirect = func(_ *http.Request, _ []*http.Request) error {
		return http.ErrUseLastResponse
	}
	if config.Now == nil {
		config.Now = time.Now
	}
	if config.ClockSkew <= 0 {
		config.ClockSkew = 2 * time.Minute
	}
	if config.MinimumKeyTTL <= 0 {
		config.MinimumKeyTTL = time.Minute
	}
	if config.DefaultKeyTTL <= 0 {
		config.DefaultKeyTTL = time.Hour
	}
	if config.MaximumKeyTTL <= 0 {
		config.MaximumKeyTTL = 24 * time.Hour
	}
	if config.MinimumKeyTTL > config.DefaultKeyTTL || config.DefaultKeyTTL > config.MaximumKeyTTL {
		return nil, errors.New("OIDC verifier: key-cache TTL bounds are invalid")
	}
	return &OIDCIDTokenVerifier{
		audience: audience, issuers: issuers, jwksURL: jwksURL,
		client: &httpClient, now: config.Now, skew: config.ClockSkew,
		minTTL: config.MinimumKeyTTL, defaultTTL: config.DefaultKeyTTL,
		maxTTL: config.MaximumKeyTTL,
	}, nil
}

func (v *OIDCIDTokenVerifier) Verify(ctx context.Context, raw string) (OIDCIdentity, error) {
	if v == nil || ctx == nil {
		return OIDCIdentity{}, &OIDCProviderUnavailableError{cause: errors.New("verifier is not configured")}
	}
	if err := ctx.Err(); err != nil {
		return OIDCIdentity{}, err
	}
	if raw == "" || len(raw) > maxOIDCIDTokenBytes || strings.TrimSpace(raw) != raw {
		return OIDCIdentity{}, &OIDCInvalidTokenError{}
	}
	claims := &oidcIDTokenClaims{}
	parser := jwt.NewParser(
		jwt.WithValidMethods([]string{"RS256"}),
		jwt.WithAudience(v.audience),
		jwt.WithExpirationRequired(),
		jwt.WithLeeway(v.skew),
		jwt.WithTimeFunc(v.now),
	)
	token, err := parser.ParseWithClaims(raw, claims, func(token *jwt.Token) (any, error) {
		if token == nil || token.Method.Alg() != "RS256" {
			return nil, &OIDCInvalidTokenError{}
		}
		kid, ok := token.Header["kid"].(string)
		kid = strings.TrimSpace(kid)
		if !ok || kid == "" {
			return nil, &OIDCInvalidTokenError{}
		}
		return v.key(ctx, kid)
	})
	if err != nil {
		var unavailable *OIDCProviderUnavailableError
		if errors.As(err, &unavailable) {
			return OIDCIdentity{}, unavailable
		}
		if ctxErr := ctx.Err(); ctxErr != nil {
			return OIDCIdentity{}, ctxErr
		}
		return OIDCIdentity{}, &OIDCInvalidTokenError{}
	}
	if token == nil || !token.Valid || claims.Subject == "" ||
		strings.TrimSpace(claims.Subject) != claims.Subject ||
		!claims.EmailVerified || claims.Email == "" || strings.TrimSpace(claims.Email) != claims.Email {
		return OIDCIdentity{}, &OIDCInvalidTokenError{}
	}
	if _, trusted := v.issuers[claims.Issuer]; !trusted {
		return OIDCIdentity{}, &OIDCInvalidTokenError{}
	}
	if claims.AuthorizedParty != "" && claims.AuthorizedParty != v.audience {
		return OIDCIdentity{}, &OIDCInvalidTokenError{}
	}
	if len(claims.Audience) > 1 && claims.AuthorizedParty != v.audience {
		return OIDCIdentity{}, &OIDCInvalidTokenError{}
	}
	return OIDCIdentity{Subject: claims.Subject, Email: claims.Email}, nil
}

func (v *OIDCIDTokenVerifier) key(ctx context.Context, kid string) (*rsa.PublicKey, error) {
	v.mu.Lock()
	if key, ok := v.lookupKey(kid); ok {
		v.mu.Unlock()
		return key, nil
	}
	v.mu.Unlock()

	// Never hold the cache lock over provider I/O. A slow JWKS refresh must not
	// prevent another request's cancellation or deadline from propagating.
	keys, expiresAt, err := v.fetchKeys(ctx)
	if err != nil {
		return nil, err
	}
	v.mu.Lock()
	v.keys = keys
	v.expiresAt = expiresAt
	if key, ok := v.lookupKey(kid); ok {
		v.mu.Unlock()
		return key, nil
	}
	v.mu.Unlock()
	return nil, &OIDCInvalidTokenError{}
}

func (v *OIDCIDTokenVerifier) lookupKey(kid string) (*rsa.PublicKey, bool) {
	if len(v.keys) == 0 || !v.now().Before(v.expiresAt) {
		return nil, false
	}
	key, ok := v.keys[kid]
	return key, ok
}

func (v *OIDCIDTokenVerifier) fetchKeys(ctx context.Context) (map[string]*rsa.PublicKey, time.Time, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, v.jwksURL, nil)
	if err != nil {
		return nil, time.Time{}, &OIDCProviderUnavailableError{cause: err}
	}
	response, err := v.client.Do(req)
	if err != nil {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return nil, time.Time{}, ctxErr
		}
		return nil, time.Time{}, &OIDCProviderUnavailableError{cause: err}
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return nil, time.Time{}, &OIDCProviderUnavailableError{cause: fmt.Errorf("JWKS status %d", response.StatusCode)}
	}
	body, err := io.ReadAll(io.LimitReader(response.Body, 1<<20))
	if err != nil {
		return nil, time.Time{}, &OIDCProviderUnavailableError{cause: err}
	}
	var document struct {
		Keys []struct {
			Kty string `json:"kty"`
			Kid string `json:"kid"`
			Alg string `json:"alg"`
			Use string `json:"use"`
			N   string `json:"n"`
			E   string `json:"e"`
		} `json:"keys"`
	}
	if err := json.Unmarshal(body, &document); err != nil {
		return nil, time.Time{}, &OIDCProviderUnavailableError{cause: err}
	}
	keys := make(map[string]*rsa.PublicKey, len(document.Keys))
	for _, candidate := range document.Keys {
		if candidate.Kid == "" || !strings.EqualFold(candidate.Kty, "RSA") ||
			(candidate.Alg != "" && !strings.EqualFold(candidate.Alg, "RS256")) ||
			(candidate.Use != "" && !strings.EqualFold(candidate.Use, "sig")) {
			continue
		}
		key, keyErr := oidcRSAKey(candidate.N, candidate.E)
		if keyErr == nil {
			keys[candidate.Kid] = key
		}
	}
	if len(keys) == 0 {
		return nil, time.Time{}, &OIDCProviderUnavailableError{cause: errors.New("JWKS has no usable signing key")}
	}
	return keys, v.now().Add(v.cacheTTL(response.Header.Get("Cache-Control"))), nil
}

func (v *OIDCIDTokenVerifier) cacheTTL(header string) time.Duration {
	for _, directive := range strings.Split(header, ",") {
		directive = strings.TrimSpace(strings.ToLower(directive))
		if !strings.HasPrefix(directive, "max-age=") {
			continue
		}
		seconds, err := strconv.ParseInt(strings.TrimPrefix(directive, "max-age="), 10, 64)
		if err != nil || seconds <= 0 {
			return v.defaultTTL
		}
		if seconds >= int64(v.maxTTL/time.Second) {
			return v.maxTTL
		}
		ttl := time.Duration(seconds) * time.Second
		if ttl < v.minTTL {
			return v.minTTL
		}
		if ttl > v.maxTTL {
			return v.maxTTL
		}
		return ttl
	}
	return v.defaultTTL
}

func oidcRSAKey(modulus, exponent string) (*rsa.PublicKey, error) {
	nBytes, err := base64.RawURLEncoding.DecodeString(strings.TrimSpace(modulus))
	if err != nil {
		return nil, err
	}
	eBytes, err := base64.RawURLEncoding.DecodeString(strings.TrimSpace(exponent))
	if err != nil {
		return nil, err
	}
	if len(nBytes) == 0 || len(eBytes) == 0 {
		return nil, errors.New("empty RSA key")
	}
	e := new(big.Int).SetBytes(eBytes)
	if !e.IsInt64() || e.Int64() < 2 {
		return nil, errors.New("invalid RSA exponent")
	}
	return &rsa.PublicKey{N: new(big.Int).SetBytes(nBytes), E: int(e.Int64())}, nil
}