package middleware

import (
	"sync"
	"time"

	"github.com/gofiber/fiber/v2"
)

type loginAttempt struct {
	Count    int
	LockedAt time.Time
}

// BruteForceGuard tracks login attempts per IP and enforces exponential backoff.
type BruteForceGuard struct {
	mu           sync.Mutex
	attempts     map[string]*loginAttempt
	maxAttempts  int
	lockoutSecs  int
}

func NewBruteForceGuard(maxAttempts, lockoutSecs int) *BruteForceGuard {
	guard := &BruteForceGuard{
		attempts:    make(map[string]*loginAttempt),
		maxAttempts: maxAttempts,
		lockoutSecs: lockoutSecs,
	}

	// Clean up stale entries every 10 minutes
	go func() {
		ticker := time.NewTicker(10 * time.Minute)
		defer ticker.Stop()
		for range ticker.C {
			guard.cleanup()
		}
	}()

	return guard
}

// Check returns a Fiber handler that blocks requests from locked-out IPs.
func (g *BruteForceGuard) Check() fiber.Handler {
	return func(c *fiber.Ctx) error {
		ip := c.IP()

		g.mu.Lock()
		attempt, exists := g.attempts[ip]

		if exists && attempt.Count >= g.maxAttempts {
			lockoutDuration := time.Duration(g.lockoutSecs) * time.Second
			if time.Since(attempt.LockedAt) < lockoutDuration {
				remaining := lockoutDuration - time.Since(attempt.LockedAt)
				g.mu.Unlock()
				return c.Status(fiber.StatusTooManyRequests).JSON(fiber.Map{
					"error":       "too many login attempts, account temporarily locked",
					"retry_after": int(remaining.Seconds()),
				})
			}
			// Lockout expired, reset
			delete(g.attempts, ip)
		}
		g.mu.Unlock()

		return c.Next()
	}
}

// RecordFailure increments the failure count for an IP.
func (g *BruteForceGuard) RecordFailure(ip string) {
	g.mu.Lock()
	defer g.mu.Unlock()

	attempt, exists := g.attempts[ip]
	if !exists {
		g.attempts[ip] = &loginAttempt{Count: 1, LockedAt: time.Now()}
		return
	}

	attempt.Count++
	if attempt.Count >= g.maxAttempts {
		attempt.LockedAt = time.Now()
	}
}

// RecordSuccess resets the failure count for an IP after a successful login.
func (g *BruteForceGuard) RecordSuccess(ip string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	delete(g.attempts, ip)
}

func (g *BruteForceGuard) cleanup() {
	g.mu.Lock()
	defer g.mu.Unlock()

	cutoff := time.Duration(g.lockoutSecs*2) * time.Second
	for ip, attempt := range g.attempts {
		if time.Since(attempt.LockedAt) > cutoff {
			delete(g.attempts, ip)
		}
	}
}