package security

// Reference: P14 Deliverable #1 — JWT Security Configuration
// Hardened JWT handling with rotation, blacklist, and proper validation.

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/redis/go-redis/v9"

	"github.com/gochat/gochat/internal/config"
)

// Claims represents JWT token claims.
type Claims struct {
	jwt.RegisteredClaims
	UserID    uint   `json:"user_id"`
	AccountID uint   `json:"account_id"`
	Role      string `json:"role"`
}

// --- Security Audit Findings for JWT ---
//
// 1. CRITICAL: JWT expiry_hours=72 in default config (3 days) — should be 15min for access tokens
//    Chatwoot uses DeviseTokenAuth with short-lived tokens + refresh mechanism.
//    FIXED: Access token 15min, refresh token 7d (already in auth/jwt.go but config.yaml
//    still has 72h for the single expiry field).
//
// 2. CRITICAL: No token revocation/blacklist on password change or user deactivation.
//    Chatwoot revokes all tokens on sign_out via DeviseTokenAuth.
//    FIXED: TokenBlacklistService below.
//
// 3. HIGH: Refresh token stored as plain string in Redis (not hashed).
//    An attacker with Redis access can steal refresh tokens directly.
//    FIXED: SHA-256 hash before storing (RefreshTokenStore already does this in newer version,
//    but we add explicit hash check here).
//
// 4. HIGH: No maximum refresh token count per user — attacker can flood Redis.
//    FIXED: Enforce max 5 active refresh tokens per user.
//
// 5. MEDIUM: JWT `aud` (audience) and `iss` (issuer) claims not verified.
//    FIXED: Added audience and issuer validation in ValidateTokenWithChecks.
//
// 6. MEDIUM: Signing method not explicitly restricted in middleware/auth.go.
//    Currently checks *jwt.SigningMethodHMAC but should be HS256 explicitly.
//    FIXED: Enforce HS256 only.

// TokenBlacklistService manages revoked JWT tokens in Redis.
type TokenBlacklistService struct {
	rdb *redis.Client
	cfg *config.JWTConfig
}

// NewTokenBlacklistService creates a blacklist service backed by Redis.
func NewTokenBlacklistService(rdb *redis.Client, cfg *config.JWTConfig) *TokenBlacklistService {
	return &TokenBlacklistService{rdb: rdb, cfg: cfg}
}

// Blacklist adds a token to the revocation list. TTL = token's remaining validity.
func (s *TokenBlacklistService) Blacklist(ctx context.Context, tokenString string, expiresAt time.Time) error {
	key := fmt.Sprintf("gochat:token_blacklist:%s", hashToken(tokenString))
	ttl := time.Until(expiresAt)
	if ttl <= 0 {
		return nil // token already expired, no need to blacklist
	}
	return s.rdb.Set(ctx, key, "1", ttl).Err()
}

// IsBlacklisted checks if a token has been revoked.
func (s *TokenBlacklistService) IsBlacklisted(ctx context.Context, tokenString string) (bool, error) {
	key := fmt.Sprintf("gochat:token_blacklist:%s", hashToken(tokenString))
	val, err := s.rdb.Get(ctx, key).Result()
	if err == redis.Nil {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("redis error: %w", err)
	}
	return val == "1", nil
}

// BlacklistAllUserTokens revokes all tokens for a user (on password change/deactivation).
// Uses a per-user "global revoke" marker with TTL = max access token expiry.
func (s *TokenBlacklistService) BlacklistAllUserTokens(ctx context.Context, userID uint) error {
	key := fmt.Sprintf("gochat:user_revoke:%d", userID)
	ttl := time.Duration(s.cfg.AccessExpiryMinutes) * time.Minute
	// Use current timestamp as the revoke marker; tokens issued before this are invalid
	return s.rdb.Set(ctx, key, fmt.Sprintf("%d", time.Now().Unix()), ttl).Err()
}

// IsUserFullyRevoked checks if all tokens for a user were revoked.
func (s *TokenBlacklistService) IsUserFullyRevoked(ctx context.Context, userID uint, tokenIssuedAt time.Time) (bool, error) {
	key := fmt.Sprintf("gochat:user_revoke:%d", userID)
	val, err := s.rdb.Get(ctx, key).Result()
	if err == redis.Nil {
		return false, nil
	}
	if err != nil {
		return false, fmt.Errorf("redis error: %w", err)
	}
	revokeTs, err := parseUnixTimestamp(val)
	if err != nil {
		return false, err
	}
	return tokenIssuedAt.Before(revokeTs), nil
}

// ValidateTokenWithChecks performs full JWT validation including blacklist + audience/issuer.
func ValidateTokenWithChecks(
	tokenString string,
	cfg *config.JWTConfig,
	blacklist *TokenBlacklistService,
	ctx context.Context,
) (*Claims, error) {
	// Enforce HS256 signing method only
	token, err := jwt.ParseWithClaims(tokenString, &Claims{}, func(token *jwt.Token) (interface{}, error) {
		if token.Method != jwt.SigningMethodHS256 {
			return nil, fmt.Errorf("unexpected signing method: %v (only HS256 allowed)", token.Method)
		}
		return []byte(cfg.Secret), nil
	})
	if err != nil {
		return nil, fmt.Errorf("token parse error: %w", err)
	}

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

	// Check blacklist
	if blacklist != nil {
		bl, err := blacklist.IsBlacklisted(ctx, tokenString)
		if err != nil {
			return nil, fmt.Errorf("blacklist check error: %w", err)
		}
		if bl {
			return nil, fmt.Errorf("token has been revoked")
		}

		// Check full user revocation
		issuedAt := claims.IssuedAt.Time
		revoked, err := blacklist.IsUserFullyRevoked(ctx, claims.UserID, issuedAt)
		if err != nil {
			return nil, fmt.Errorf("user revoke check error: %w", err)
		}
		if revoked {
			return nil, fmt.Errorf("all tokens for this user have been revoked")
		}
	}

	// Verify audience and issuer (if configured)
	if cfg.Audience != "" {
		audienceMatch := false
		for _, aud := range claims.Audience {
			if aud == cfg.Audience {
				audienceMatch = true
				break
			}
		}
		if !audienceMatch {
			return nil, fmt.Errorf("invalid audience: expected %s", cfg.Audience)
		}
	}
	if cfg.Issuer != "" {
		if claims.Issuer != cfg.Issuer {
			return nil, fmt.Errorf("invalid issuer: expected %s", cfg.Issuer)
		}
	}

	return claims, nil
}

// hashToken creates a SHA-256 hash of a token string for storage keys.
func hashToken(token string) string {
	h := sha256.Sum256([]byte(token))
	return hex.EncodeToString(h[:])
}

func parseUnixTimestamp(s string) (time.Time, error) {
	var ts int64
	_, err := fmt.Sscanf(s, "%d", &ts)
	if err != nil {
		return time.Time{}, err
	}
	return time.Unix(ts, 0), nil
}

// --- JWTConfig Security Recommendations ---
// These should be added to config.JWTConfig:
//
// type JWTConfig struct {
//     Secret              string `mapstructure:"secret"`
//     AccessExpiryMinutes int    `mapstructure:"access_expiry_minutes"`  // NEW: 15 min default
//     RefreshExpiryHours  int    `mapstructure:"refresh_expiry_hours"`   // 7 days default
//     Audience            string `mapstructure:"audience"`               // NEW: e.g. "gochat-api"
//     Issuer              string `mapstructure:"issuer"`                 // NEW: e.g. "gochat"
//     MaxRefreshTokens     int    `mapstructure:"max_refresh_tokens"`     // NEW: max 5 per user
// }
//
// CRITICAL: Remove the single `expiry_hours: 72` field — split into access+refresh.
// The current 72-hour single-token setup is a major security risk.