package auth

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/MicahParks/keyfunc/v2"
	"github.com/golang-jwt/jwt/v5"
)

// Validator validates JWTs against Authentik JWKS.
type Validator struct {
	issuer   string
	audience string
	jwksURL  string
	jwks     *keyfunc.JWKS
	mu       sync.Mutex
}

// NewValidator creates a new JWT validator for a service.
// JWKS fetching is lazy (deferred until first validation).
func NewValidator(issuer, audience, jwksURL string) *Validator {
	return &Validator{
		issuer:   issuer,
		audience: audience,
		jwksURL:  jwksURL,
		jwks:     nil, // Lazy-loaded on first use
	}
}

// ensureJWKS fetches JWKS on first use (lazy initialization, thread-safe).
func (v *Validator) ensureJWKS() error {
	v.mu.Lock()
	defer v.mu.Unlock()

	if v.jwks != nil {
		return nil
	}

	options := keyfunc.Options{
		Ctx:              context.Background(),
		RefreshInterval:  15 * time.Minute,
		RefreshRateLimit: 5 * time.Minute,
		RefreshTimeout:   10 * time.Second,
		RefreshErrorHandler: func(err error) {
			fmt.Printf("JWKS refresh error for %s: %v\n", v.issuer, err)
		},
	}

	jwks, err := keyfunc.Get(v.jwksURL, options)
	if err != nil {
		return fmt.Errorf("failed to fetch JWKS from %s: %v", v.jwksURL, err)
	}

	v.jwks = jwks
	return nil
}

// ValidateBearerToken extracts and validates the Bearer token from Authorization header.
// Returns claims on success, error message on failure.
func (v *Validator) ValidateBearerToken(authHeader string) (jwt.MapClaims, error) {
	if authHeader == "" {
		return nil, fmt.Errorf("missing Authorization header")
	}

	// Extract token from "Bearer <token>"
	tokenString := ""
	if len(authHeader) > 7 && authHeader[:7] == "Bearer " {
		tokenString = authHeader[7:]
	} else {
		return nil, fmt.Errorf("invalid Authorization header format")
	}

	// Ensure JWKS is loaded (lazy)
	if err := v.ensureJWKS(); err != nil {
		return nil, err
	}

	// Parse and validate
	claims := jwt.MapClaims{}
	token, err := jwt.ParseWithClaims(tokenString, claims, v.jwks.Keyfunc)
	if err != nil {
		return nil, fmt.Errorf("token validation failed: %v", err)
	}

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

	// Verify required claims
	now := time.Now()
	const skew = 60 * time.Second

	// Check exp
	if exp, ok := claims["exp"].(float64); ok {
		if time.Now().After(time.Unix(int64(exp), 0).Add(skew)) {
			return nil, fmt.Errorf("token expired")
		}
	}

	// Check nbf (not before)
	if nbf, ok := claims["nbf"].(float64); ok {
		if now.Before(time.Unix(int64(nbf), 0).Add(-skew)) {
			return nil, fmt.Errorf("token not yet valid")
		}
	}

	// Check iss (issuer) - if configured, must match exactly or be from Authentik
	// Service accounts have per-provider issuers, so we allow any authentik.riotpiao.com issuer
	if v.issuer != "" {
		iss, ok := claims["iss"].(string)
		if !ok {
			return nil, fmt.Errorf("missing issuer claim")
		}
		// Allow exact match or any Authentik issuer
		if iss != v.issuer && !isAuthentikIssuer(iss) {
			return nil, fmt.Errorf("invalid issuer: expected %s or Authentik, got %s", v.issuer, iss)
		}
	}

	// Check aud (audience) - if configured, must match
	if v.audience != "" {
		if aud, ok := claims["aud"].(string); !ok || aud != v.audience {
			return nil, fmt.Errorf("invalid audience: expected %s, got %s", v.audience, aud)
		}
	}

	return claims, nil
}

// CheckPermissions checks if claims contain required permission(s).
// Checks both "permissions" (user tokens) and "roles" (service account tokens).
// Returns true if any required permission is found or wildcard "*" exists.
func (v *Validator) CheckPermissions(claims jwt.MapClaims, required ...string) bool {
	// Check both claims - users have "permissions", service accounts have "roles"
	for _, claimKey := range []string{"permissions", "roles"} {
		if v.checkClaimList(claims, claimKey, required...) {
			return true
		}
	}
	return false
}

// checkClaimList checks if a specific claim contains any required value.
func (v *Validator) checkClaimList(claims jwt.MapClaims, claimKey string, required ...string) bool {
	valuesIface, ok := claims[claimKey]
	if !ok {
		return false
	}

	values, ok := valuesIface.([]interface{})
	if !ok {
		return false
	}

	for _, val := range values {
		valStr, ok := val.(string)
		if !ok {
			continue
		}
		if valStr == "*" {
			return true
		}
		for _, req := range required {
			if valStr == req {
				return true
			}
		}
	}

	return false
}

// isAuthentikIssuer checks if issuer is from our Authentik instance.
func isAuthentikIssuer(iss string) bool {
	return len(iss) > 0 && 
		(iss == "https://authentik.riotpiao.com" ||
		 strings.HasPrefix(iss, "https://authentik.riotpiao.com/"))
}

// DecodeToken decodes JWT payload without verification (for debugging/testing).
func DecodeToken(tokenString string) (jwt.MapClaims, error) {
	claims := jwt.MapClaims{}
	_, _, err := new(jwt.Parser).ParseUnverified(tokenString, claims)
	if err != nil {
		return nil, err
	}
	return claims, nil
}