package ratelimit

import (
	"context"
	"fmt"
	"time"

	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"
)

// LoginRateLimiter provides specialized rate limiting for login attempts
// with account lockout functionality
type LoginRateLimiter interface {
	// CheckAndRecordAttempt checks if login attempt is allowed and records it
	// Returns: allowed (bool), isLocked (bool), remainingAttempts (int), error
	CheckAndRecordAttempt(ctx context.Context, email string, clientIP string) (bool, bool, int, error)

	// RecordFailedAttempt records a failed login attempt
	RecordFailedAttempt(ctx context.Context, email string, clientIP string) error

	// RecordSuccessfulLogin records a successful login and resets counters
	RecordSuccessfulLogin(ctx context.Context, email string, clientIP string) error

	// IsAccountLocked checks if an account is locked due to too many failed attempts
	IsAccountLocked(ctx context.Context, email string) (bool, time.Duration, error)

	// UnlockAccount manually unlocks an account (admin function)
	UnlockAccount(ctx context.Context, email string) error

	// GetFailedAttempts returns the number of failed attempts for an email
	GetFailedAttempts(ctx context.Context, email string) (int, error)
}

// LoginRateLimiterConfig holds configuration for login rate limiting
type LoginRateLimiterConfig struct {
	// MaxAttemptsPerIP is the maximum login attempts per IP in the window
	MaxAttemptsPerIP int
	// IPWindow is the time window for IP-based rate limiting
	IPWindow time.Duration

	// MaxFailedAttemptsPerAccount is the maximum failed attempts before account lockout
	MaxFailedAttemptsPerAccount int
	// AccountLockoutDuration is how long to lock an account after too many failures
	AccountLockoutDuration time.Duration

	// KeyPrefix is the prefix for Redis keys
	KeyPrefix string
}

// DefaultLoginRateLimiterConfig returns recommended configuration
func DefaultLoginRateLimiterConfig() LoginRateLimiterConfig {
	return LoginRateLimiterConfig{
		MaxAttemptsPerIP:            10,               // 10 attempts per IP
		IPWindow:                    15 * time.Minute, // in 15 minutes
		MaxFailedAttemptsPerAccount: 10,               // 10 failed attempts per account
		AccountLockoutDuration:      30 * time.Minute, // lock for 30 minutes
		KeyPrefix:                   "login_rl",
	}
}

type loginRateLimiter struct {
	client *redis.Client
	config LoginRateLimiterConfig
	logger *zap.Logger
}

// NewLoginRateLimiter creates a new login rate limiter
func NewLoginRateLimiter(client *redis.Client, config LoginRateLimiterConfig, logger *zap.Logger) LoginRateLimiter {
	return &loginRateLimiter{
		client: client,
		config: config,
		logger: logger.Named("login-rate-limiter"),
	}
}

// CheckAndRecordAttempt checks if login attempt is allowed
// CWE-307: Implements protection against brute force attacks
func (r *loginRateLimiter) CheckAndRecordAttempt(ctx context.Context, email string, clientIP string) (bool, bool, int, error) {
	// Check account lockout first
	locked, remaining, err := r.IsAccountLocked(ctx, email)
	if err != nil {
		r.logger.Error("failed to check account lockout",
			zap.String("email_hash", hashEmail(email)),
			zap.Error(err))
		// Fail open on Redis error
		return true, false, 0, err
	}

	if locked {
		r.logger.Warn("login attempt on locked account",
			zap.String("email_hash", hashEmail(email)),
			zap.String("ip", clientIP),
			zap.Duration("remaining_lockout", remaining))
		return false, true, 0, nil
	}

	// Check IP-based rate limit
	ipKey := r.getIPKey(clientIP)
	allowed, err := r.checkIPRateLimit(ctx, ipKey)
	if err != nil {
		r.logger.Error("failed to check IP rate limit",
			zap.String("ip", clientIP),
			zap.Error(err))
		// Fail open on Redis error
		return true, false, 0, err
	}

	if !allowed {
		r.logger.Warn("IP rate limit exceeded",
			zap.String("ip", clientIP))
		return false, false, 0, nil
	}

	// Record the attempt for IP
	if err := r.recordIPAttempt(ctx, ipKey); err != nil {
		r.logger.Error("failed to record IP attempt",
			zap.String("ip", clientIP),
			zap.Error(err))
	}

	// Get remaining attempts for account
	failedAttempts, err := r.GetFailedAttempts(ctx, email)
	if err != nil {
		r.logger.Error("failed to get failed attempts",
			zap.String("email_hash", hashEmail(email)),
			zap.Error(err))
	}

	remainingAttempts := r.config.MaxFailedAttemptsPerAccount - failedAttempts
	if remainingAttempts < 0 {
		remainingAttempts = 0
	}

	r.logger.Debug("login attempt check passed",
		zap.String("email_hash", hashEmail(email)),
		zap.String("ip", clientIP),
		zap.Int("remaining_attempts", remainingAttempts))

	return true, false, remainingAttempts, nil
}

// RecordFailedAttempt records a failed login attempt
// CWE-307: Tracks failed attempts to enable account lockout
func (r *loginRateLimiter) RecordFailedAttempt(ctx context.Context, email string, clientIP string) error {
	accountKey := r.getAccountKey(email)

	// Increment failed attempt counter
	count, err := r.client.Incr(ctx, accountKey).Result()
	if err != nil {
		r.logger.Error("failed to increment failed attempts",
			zap.String("email_hash", hashEmail(email)),
			zap.Error(err))
		return err
	}

	// Set expiration on first failed attempt
	if count == 1 {
		r.client.Expire(ctx, accountKey, r.config.AccountLockoutDuration)
	}

	// Check if account should be locked
	if count >= int64(r.config.MaxFailedAttemptsPerAccount) {
		lockKey := r.getLockKey(email)
		err := r.client.Set(ctx, lockKey, "locked", r.config.AccountLockoutDuration).Err()
		if err != nil {
			r.logger.Error("failed to lock account",
				zap.String("email_hash", hashEmail(email)),
				zap.Error(err))
			return err
		}

		r.logger.Warn("account locked due to too many failed attempts",
			zap.String("email_hash", hashEmail(email)),
			zap.String("ip", clientIP),
			zap.Int64("failed_attempts", count),
			zap.Duration("lockout_duration", r.config.AccountLockoutDuration))
	}

	r.logger.Info("failed login attempt recorded",
		zap.String("email_hash", hashEmail(email)),
		zap.String("ip", clientIP),
		zap.Int64("total_failed_attempts", count))

	return nil
}

// RecordSuccessfulLogin records a successful login and resets counters
func (r *loginRateLimiter) RecordSuccessfulLogin(ctx context.Context, email string, clientIP string) error {
	accountKey := r.getAccountKey(email)
	lockKey := r.getLockKey(email)

	// Delete failed attempt counter
	pipe := r.client.Pipeline()
	pipe.Del(ctx, accountKey)
	pipe.Del(ctx, lockKey)
	_, err := pipe.Exec(ctx)

	if err != nil {
		r.logger.Error("failed to reset login counters",
			zap.String("email_hash", hashEmail(email)),
			zap.Error(err))
		return err
	}

	r.logger.Info("successful login recorded, counters reset",
		zap.String("email_hash", hashEmail(email)),
		zap.String("ip", clientIP))

	return nil
}

// IsAccountLocked checks if an account is locked
func (r *loginRateLimiter) IsAccountLocked(ctx context.Context, email string) (bool, time.Duration, error) {
	lockKey := r.getLockKey(email)

	ttl, err := r.client.TTL(ctx, lockKey).Result()
	if err != nil {
		return false, 0, err
	}

	// TTL returns -2 if key doesn't exist, -1 if no expiration
	if ttl < 0 {
		return false, 0, nil
	}

	return true, ttl, nil
}

// UnlockAccount manually unlocks an account
func (r *loginRateLimiter) UnlockAccount(ctx context.Context, email string) error {
	accountKey := r.getAccountKey(email)
	lockKey := r.getLockKey(email)

	pipe := r.client.Pipeline()
	pipe.Del(ctx, accountKey)
	pipe.Del(ctx, lockKey)
	_, err := pipe.Exec(ctx)

	if err != nil {
		r.logger.Error("failed to unlock account",
			zap.String("email_hash", hashEmail(email)),
			zap.Error(err))
		return err
	}

	r.logger.Info("account unlocked",
		zap.String("email_hash", hashEmail(email)))

	return nil
}

// GetFailedAttempts returns the number of failed attempts
func (r *loginRateLimiter) GetFailedAttempts(ctx context.Context, email string) (int, error) {
	accountKey := r.getAccountKey(email)

	count, err := r.client.Get(ctx, accountKey).Int()
	if err == redis.Nil {
		return 0, nil
	}
	if err != nil {
		return 0, err
	}

	return count, nil
}

// checkIPRateLimit checks if IP has exceeded rate limit
func (r *loginRateLimiter) checkIPRateLimit(ctx context.Context, ipKey string) (bool, error) {
	now := time.Now()
	windowStart := now.Add(-r.config.IPWindow)

	// Remove old entries
	r.client.ZRemRangeByScore(ctx, ipKey, "0", fmt.Sprintf("%d", windowStart.UnixNano()))

	// Count current attempts
	count, err := r.client.ZCount(ctx, ipKey,
		fmt.Sprintf("%d", windowStart.UnixNano()),
		"+inf").Result()

	if err != nil && err != redis.Nil {
		return false, err
	}

	return count < int64(r.config.MaxAttemptsPerIP), nil
}

// recordIPAttempt records an IP attempt
func (r *loginRateLimiter) recordIPAttempt(ctx context.Context, ipKey string) error {
	now := time.Now()
	timestamp := now.UnixNano()

	pipe := r.client.Pipeline()
	pipe.ZAdd(ctx, ipKey, redis.Z{
		Score:  float64(timestamp),
		Member: fmt.Sprintf("%d", timestamp),
	})
	pipe.Expire(ctx, ipKey, r.config.IPWindow+time.Minute)
	_, err := pipe.Exec(ctx)

	return err
}

// Key generation helpers
func (r *loginRateLimiter) getIPKey(ip string) string {
	return fmt.Sprintf("%s:ip:%s", r.config.KeyPrefix, ip)
}

func (r *loginRateLimiter) getAccountKey(email string) string {
	return fmt.Sprintf("%s:account:%s:attempts", r.config.KeyPrefix, hashEmail(email))
}

func (r *loginRateLimiter) getLockKey(email string) string {
	return fmt.Sprintf("%s:account:%s:locked", r.config.KeyPrefix, hashEmail(email))
}

// hashEmail creates a consistent hash of an email for use as a key
// CWE-532: Prevents PII in Redis keys
func hashEmail(email string) string {
	// Use a simple hash for key generation (not for security)
	// In production, consider using SHA-256
	hash := 0
	for _, c := range email {
		hash = (hash * 31) + int(c)
	}
	return fmt.Sprintf("%x", hash)
}
