package middleware

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/niels/banking-app/backend/internal/platform/httputil"
	"github.com/niels/banking-app/backend/internal/respond"
)

type RateLimitOptions struct {
	Enabled     bool
	Backend     string
	Window      time.Duration
	MaxRequests int
	FailOpen    bool
}

type RateLimiter struct {
	options RateLimitOptions
	store   rateLimitStore
	now     func() time.Time
}

type rateLimitStore interface {
	Allow(ctx context.Context, key string, now time.Time, options RateLimitOptions) (bool, time.Duration, error)
}

type memoryRateLimitStore struct {
	mu      sync.Mutex
	buckets map[string]rateBucket
}

type rateBucket struct {
	WindowStart time.Time
	Count       int
}

type postgresRateLimitStore struct {
	db          *pgxpool.Pool
	mu          sync.Mutex
	lastCleanup time.Time
}

func NewRateLimiter(options RateLimitOptions) *RateLimiter {
	options = normalizeRateLimitOptions(options)
	return &RateLimiter{
		options: options,
		store:   &memoryRateLimitStore{buckets: map[string]rateBucket{}},
		now:     time.Now,
	}
}

func NewPostgresRateLimiter(db *pgxpool.Pool, options RateLimitOptions) *RateLimiter {
	options = normalizeRateLimitOptions(options)
	options.Backend = "postgres"
	return &RateLimiter{
		options: options,
		store:   &postgresRateLimitStore{db: db},
		now:     time.Now,
	}
}

func RateLimit(limiter *RateLimiter) Middleware {
	return func(next http.Handler) http.Handler {
		if limiter == nil || !limiter.options.Enabled {
			return next
		}
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			allowed, retryAfter, err := limiter.allow(r.Context(), keyForRequest(r))
			if err != nil {
				if limiter.options.FailOpen {
					next.ServeHTTP(w, r)
					return
				}
				respond.Problem(w, http.StatusServiceUnavailable, "rate_limiter_unavailable", "rate limiter is unavailable")
				return
			}
			if !allowed {
				w.Header().Set("Retry-After", fmt.Sprintf("%.0f", retryAfter.Seconds()))
				respond.Problem(w, http.StatusTooManyRequests, "rate_limited", "too many requests, slow down and retry later")
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

func (l *RateLimiter) allow(ctx context.Context, key string) (bool, time.Duration, error) {
	now := l.now().UTC()
	return l.store.Allow(ctx, key, now, l.options)
}

func (s *memoryRateLimitStore) Allow(_ context.Context, key string, now time.Time, options RateLimitOptions) (bool, time.Duration, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	bucket := s.buckets[key]
	if bucket.WindowStart.IsZero() || now.Sub(bucket.WindowStart) >= options.Window {
		s.buckets[key] = rateBucket{WindowStart: now, Count: 1}
		s.cleanupLocked(now, options.Window)
		return true, 0, nil
	}
	if bucket.Count >= options.MaxRequests {
		return false, options.Window - now.Sub(bucket.WindowStart), nil
	}
	bucket.Count++
	s.buckets[key] = bucket
	return true, 0, nil
}

func (s *memoryRateLimitStore) cleanupLocked(now time.Time, window time.Duration) {
	staleBefore := now.Add(-2 * window)
	for key, bucket := range s.buckets {
		if bucket.WindowStart.Before(staleBefore) {
			delete(s.buckets, key)
		}
	}
}

func (s *postgresRateLimitStore) Allow(ctx context.Context, key string, now time.Time, options RateLimitOptions) (bool, time.Duration, error) {
	if s.db == nil {
		return false, 0, fmt.Errorf("postgres rate limiter requires database pool")
	}
	var windowStart time.Time
	var count int
	err := s.db.QueryRow(ctx, `
		INSERT INTO rate_limit_buckets (bucket_key, window_start, request_count, updated_at)
		VALUES ($1, $2, 1, now())
		ON CONFLICT (bucket_key) DO UPDATE
		SET window_start = CASE
				WHEN rate_limit_buckets.window_start <= $2::timestamptz - make_interval(secs => $3)
					THEN $2::timestamptz
				ELSE rate_limit_buckets.window_start
			END,
			request_count = CASE
				WHEN rate_limit_buckets.window_start <= $2::timestamptz - make_interval(secs => $3)
					THEN 1
				ELSE rate_limit_buckets.request_count + 1
			END,
			updated_at = now()
		RETURNING window_start, request_count
	`, key, now, int(options.Window.Seconds())).Scan(&windowStart, &count)
	if err != nil {
		return false, 0, err
	}

	s.cleanup(ctx, now, options.Window)
	if count > options.MaxRequests {
		retryAfter := options.Window - now.Sub(windowStart)
		if retryAfter < time.Second {
			retryAfter = time.Second
		}
		return false, retryAfter, nil
	}
	return true, 0, nil
}

func (s *postgresRateLimitStore) cleanup(ctx context.Context, now time.Time, window time.Duration) {
	s.mu.Lock()
	if !s.lastCleanup.IsZero() && now.Sub(s.lastCleanup) < window {
		s.mu.Unlock()
		return
	}
	s.lastCleanup = now
	s.mu.Unlock()

	_, _ = s.db.Exec(ctx, `
		DELETE FROM rate_limit_buckets
		WHERE updated_at < now() - make_interval(secs => $1)
	`, int((2 * window).Seconds()))
}

func normalizeRateLimitOptions(options RateLimitOptions) RateLimitOptions {
	if options.Backend == "" {
		options.Backend = "memory"
	}
	if options.Window <= 0 {
		options.Window = time.Minute
	}
	if options.MaxRequests <= 0 {
		options.MaxRequests = 300
	}
	return options
}

func keyForRequest(r *http.Request) string {
	return "ip:" + httputil.RemoteIP(r)
}