package auth

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

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

// Validator validates JWT tokens from Authentik.
type Validator struct {
	issuer   string
	audience string
	jwksURL  string
	client   *http.Client

	mu        sync.RWMutex
	keyset    *keySet
	lastFetch time.Time
	cacheTTL  time.Duration
}

type keySet struct {
	Keys map[string]interface{} `json:"keys"`
}

// Claims holds Authentik JWT claims.
type Claims struct {
	jwt.RegisteredClaims
	Permissions []string `json:"permissions"`
	Groups      []string `json:"groups"`
}

// NewValidator creates a validator for an Authentik app.
// appSlug is the OAuth2 provider slug (e.g., "gateway", "sqs", "memory").
func NewValidator(appSlug string) *Validator {
	issuer := fmt.Sprintf("https://authentik.riotpiao.com/application/o/%s/", appSlug)
	return &Validator{
		issuer:   issuer,
		audience: appSlug,
		jwksURL:  issuer + "jwks/",
		client:   &http.Client{Timeout: 10 * time.Second},
		cacheTTL: 15 * time.Minute,
	}
}

// ValidateToken extracts and validates a JWT from the Authorization header.
// Returns the claims if valid, or an error if invalid/missing.
func (v *Validator) ValidateToken(authHeader string) (*Claims, error) {
	// Extract token from "Bearer <token>"
	if authHeader == "" {
		return nil, fmt.Errorf("authorization header missing")
	}

	parts := strings.SplitN(authHeader, " ", 2)
	if len(parts) != 2 || parts[0] != "Bearer" {
		return nil, fmt.Errorf("invalid authorization header format")
	}

	tokenString := parts[1]

	// Ensure JWKS is fresh
	if err := v.ensureKeys(); err != nil {
		return nil, fmt.Errorf("failed to fetch JWKS: %w", err)
	}

	// Parse JWT with custom key func
	token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
		// Verify algorithm is RS256 only
		if token.Method.Alg() != "RS256" {
			return nil, fmt.Errorf("unexpected algorithm: %v", token.Header["alg"])
		}

		kid, ok := token.Header["kid"].(string)
		if !ok {
			return nil, fmt.Errorf("kid not in token header")
		}

		// Get key from cache
		v.mu.RLock()
		keys := v.keyset.Keys
		v.mu.RUnlock()

		rawKey, exists := keys[kid]
		if !exists {
			// Try refreshing JWKS (key rotation)
			if err := v.fetchKeys(); err == nil {
				v.mu.RLock()
				rawKey, exists = v.keyset.Keys[kid]
				v.mu.RUnlock()
			}
			if !exists {
				return nil, fmt.Errorf("key %s not found", kid)
			}
		}

		return rawKey, nil
	})

	if err != nil {
		return nil, fmt.Errorf("token validation failed: %w", err)
	}

	if !token.Valid {
		return nil, fmt.Errorf("token invalid")
	}

	claims, ok := token.Claims.(*Claims)
	if !ok {
		return nil, fmt.Errorf("invalid claims")
	}

	// Verify issuer
	if claims.Issuer != v.issuer {
		return nil, fmt.Errorf("issuer mismatch")
	}

	// Verify audience (check if audience is in the claims)
	if len(claims.Audience) == 0 {
		return nil, fmt.Errorf("no audience in token")
	}
	audFound := false
	for _, aud := range claims.Audience {
		if aud == v.audience {
			audFound = true
			break
		}
	}
	if !audFound {
		return nil, fmt.Errorf("audience mismatch: expected %s", v.audience)
	}

	// Verify exp, nbf, iat
	now := time.Now().Unix()
	if claims.ExpiresAt != nil && claims.ExpiresAt.Unix() < now {
		return nil, fmt.Errorf("token expired")
	}
	if claims.NotBefore != nil && claims.NotBefore.Unix() > now+60 {
		return nil, fmt.Errorf("token not yet valid")
	}

	return claims, nil
}

// HasPermission checks if claims contain the required permission.
// Wildcard "*" grants all permissions.
func (v *Validator) HasPermission(claims *Claims, permission string) bool {
	for _, p := range claims.Permissions {
		if p == "*" || p == permission {
			return true
		}
	}
	return false
}

// ensureKeys refreshes JWKS if cache is stale.
func (v *Validator) ensureKeys() error {
	v.mu.RLock()
	cacheValid := v.keyset != nil && time.Since(v.lastFetch) < v.cacheTTL
	v.mu.RUnlock()

	if cacheValid {
		return nil
	}

	return v.fetchKeys()
}

// fetchKeys fetches JWKS from Authentik.
func (v *Validator) fetchKeys() error {
	resp, err := v.client.Get(v.jwksURL)
	if err != nil {
		return fmt.Errorf("GET %s failed: %w", v.jwksURL, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("JWKS endpoint returned %d: %s", resp.StatusCode, string(body))
	}

	var ks keySet
	if err := json.NewDecoder(resp.Body).Decode(&ks); err != nil {
		return fmt.Errorf("decode JWKS failed: %w", err)
	}

	v.mu.Lock()
	v.keyset = &ks
	v.lastFetch = time.Now()
	v.mu.Unlock()

	return nil
}