package jwtauth

import (
	"context"
	"crypto/rsa"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"github.com/go-jose/go-jose/v4"
)

// jwksCache fetches & caches the Casdoor JWKS.
//
// Behavior:
//   - successful fetch: refreshes the entire key set
//   - kid miss: forces a single refresh attempt
//   - stale-while-error: if cached entry exists, return it even when refresh fails
type jwksCache struct {
	url string
	ttl time.Duration

	mu        sync.RWMutex
	keys      map[string]*rsa.PublicKey
	fetchedAt time.Time

	httpClient *http.Client
}

func newJWKSCache(url string, ttl time.Duration) *jwksCache {
	return &jwksCache{
		url:        url,
		ttl:        ttl,
		keys:       map[string]*rsa.PublicKey{},
		httpClient: &http.Client{Timeout: 5 * time.Second},
	}
}

func (j *jwksCache) GetKey(ctx context.Context, kid string) (*rsa.PublicKey, error) {
	if kid == "" {
		return nil, errors.New("missing kid")
	}

	j.mu.RLock()
	cached, found := j.keys[kid]
	fresh := !j.fetchedAt.IsZero() && time.Since(j.fetchedAt) < j.ttl
	j.mu.RUnlock()

	if found && fresh {
		return cached, nil
	}

	if err := j.refresh(ctx); err != nil {
		if found {
			slog.Warn("jwks refresh failed; returning stale key",
				"err", err, "kid", kid)
			return cached, nil
		}
		return nil, fmt.Errorf("jwks refresh: %w", err)
	}

	j.mu.RLock()
	defer j.mu.RUnlock()
	if k, ok := j.keys[kid]; ok {
		return k, nil
	}
	return nil, fmt.Errorf("kid %q not found in JWKS", kid)
}

func (j *jwksCache) refresh(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, j.url, nil)
	if err != nil {
		return err
	}
	resp, err := j.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("status %d", resp.StatusCode)
	}
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	var set jose.JSONWebKeySet
	if err := json.Unmarshal(body, &set); err != nil {
		return fmt.Errorf("parse JWKS: %w", err)
	}

	next := map[string]*rsa.PublicKey{}
	for _, k := range set.Keys {
		pk, ok := k.Key.(*rsa.PublicKey)
		if !ok || k.KeyID == "" {
			continue
		}
		next[k.KeyID] = pk
	}

	j.mu.Lock()
	defer j.mu.Unlock()
	j.keys = next
	j.fetchedAt = time.Now()
	return nil
}