package middleware

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

	"github.com/golang-jwt/jwt/v5"
	"github.com/rs/zerolog/log"
)

// JWKS representerar JSON Web Key Set
type JWKS struct {
	Keys []JWK `json:"keys"`
}

// JWK representerar en JSON Web Key
type JWK struct {
	Kty string `json:"kty"`
	Kid string `json:"kid"`
	Use string `json:"use,omitempty"`
	N   string `json:"n"`
	E   string `json:"e"`
	Alg string `json:"alg,omitempty"`
}

// JWTValidator hanterar RS256 JWT-validering med JWKS
type JWTValidator struct {
	jwksURL    string
	keys       map[string]*rsa.PublicKey
	mu         sync.RWMutex
	lastFetch  time.Time
	fetchMutex sync.Mutex
}

// NewJWTValidator skapar en ny validator med given JWKS-URL
func NewJWTValidator(jwksURL string) *JWTValidator {
	v := &JWTValidator{
		jwksURL: jwksURL,
		keys:    make(map[string]*rsa.PublicKey),
	}
	// Försök hämta keys direkt
	if err := v.fetchKeys(); err != nil {
		log.Warn().Err(err).Str("url", jwksURL).Msg("Failed to fetch JWKS initially")
	}
	return v
}

// fetchKeys hämtar och parsar JWKS från konfigurerad URL
func (v *JWTValidator) fetchKeys() error {
	v.fetchMutex.Lock()
	defer v.fetchMutex.Unlock()

	// Cache i 5 minuter
	if time.Since(v.lastFetch) < 5*time.Minute && len(v.keys) > 0 {
		return nil
	}

	resp, err := http.Get(v.jwksURL)
	if err != nil {
		return fmt.Errorf("failed to fetch JWKS: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("JWKS endpoint returned %d", resp.StatusCode)
	}

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

	newKeys := make(map[string]*rsa.PublicKey)
	for _, jwk := range jwks.Keys {
		if jwk.Kty != "RSA" {
			continue
		}
		pubKey, err := jwkToRSAPublicKey(jwk)
		if err != nil {
			log.Warn().Err(err).Str("kid", jwk.Kid).Msg("Failed to parse JWK")
			continue
		}
		newKeys[jwk.Kid] = pubKey
	}

	v.mu.Lock()
	v.keys = newKeys
	v.lastFetch = time.Now()
	v.mu.Unlock()

	log.Info().Int("keys", len(newKeys)).Str("url", v.jwksURL).Msg("JWKS fetched successfully")
	return nil
}

// jwkToRSAPublicKey konverterar en JWK till rsa.PublicKey
func jwkToRSAPublicKey(jwk JWK) (*rsa.PublicKey, error) {
	nBytes, err := base64.RawURLEncoding.DecodeString(jwk.N)
	if err != nil {
		return nil, fmt.Errorf("decode N: %w", err)
	}
	eBytes, err := base64.RawURLEncoding.DecodeString(jwk.E)
	if err != nil {
		return nil, fmt.Errorf("decode E: %w", err)
	}

	n := new(big.Int).SetBytes(nBytes)
	e := int(new(big.Int).SetBytes(eBytes).Int64())

	return &rsa.PublicKey{
		N: n,
		E: e,
	}, nil
}

// ValidateToken validerar en JWT token med RS256
func (v *JWTValidator) ValidateToken(tokenString string) (*jwt.Token, jwt.MapClaims, error) {
	// Hämta keys om nödvändigt
	if err := v.fetchKeys(); err != nil {
		return nil, nil, err
	}

	// Parse token utan validering först för att få kid
	token, _, err := new(jwt.Parser).ParseUnverified(tokenString, jwt.MapClaims{})
	if err != nil {
		return nil, nil, fmt.Errorf("parse token: %w", err)
	}

	// Avvisa HS256 tokens - endast RS256 tillåts
	if alg, ok := token.Header["alg"].(string); ok && alg == "HS256" {
		return nil, nil, fmt.Errorf("HS256 tokens not supported - use RS256")
	}

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

	v.mu.RLock()
	pubKey, ok := v.keys[kid]
	v.mu.RUnlock()

	if !ok {
		// Försök hämta keys igen (kan ha roterats)
		if err := v.fetchKeys(); err != nil {
			return nil, nil, err
		}
		v.mu.RLock()
		pubKey, ok = v.keys[kid]
		v.mu.RUnlock()
		if !ok {
			return nil, nil, fmt.Errorf("unknown key ID: %s", kid)
		}
	}

	// Validera token
	claims := jwt.MapClaims{}
	validatedToken, err := jwt.ParseWithClaims(tokenString, claims, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return pubKey, nil
	})

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

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

	return validatedToken, claims, nil
}

// Claims representerar standard JWT claims
type Claims struct {
	Sub   string   `json:"sub"`
	Email string   `json:"email"`
	Name  string   `json:"name"`
	Roles []string `json:"roles"`
}

// JWTAuth middleware som validerar RS256 tokens
func JWTAuth(jwksURL string) func(http.Handler) http.Handler {
	validator := NewJWTValidator(jwksURL)

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			authHeader := r.Header.Get("Authorization")
			if authHeader == "" {
				writeError(w, http.StatusUnauthorized, "missing authorization header")
				return
			}

			if !strings.HasPrefix(authHeader, "Bearer ") {
				writeError(w, http.StatusUnauthorized, "invalid authorization header format")
				return
			}

			tokenString := strings.TrimPrefix(authHeader, "Bearer ")
			
			// Validera med RS256 endast - ingen HS256 fallback
			_, claims, err := validator.ValidateToken(tokenString)
			if err != nil {
				log.Warn().Err(err).Str("path", r.URL.Path).Msg("JWT validation failed")
				writeError(w, http.StatusUnauthorized, "invalid token")
				return
			}

			// Extrahera claims
			userClaims := Claims{
				Sub:   getStringClaim(claims, "sub"),
				Email: getStringClaim(claims, "email"),
				Name:  getStringClaim(claims, "name"),
			}

			// Hantera roles som kan vara []interface{}
			if roles, ok := claims["roles"].([]interface{}); ok {
				for _, r := range roles {
					if s, ok := r.(string); ok {
						userClaims.Roles = append(userClaims.Roles, s)
					}
				}
			}

			ctx := WithContext(r.Context(), &userClaims)
			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}

func getStringClaim(claims jwt.MapClaims, key string) string {
	if val, ok := claims[key].(string); ok {
		return val
	}
	return ""
}