package repo

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/agentmail/gateway/internal/db"
)

// 新建会话速率限制：DB 版
//
// 这些测试用真实的 SQLite（setupTestDB），验证速率限制的原子性与窗口滑动。
// 原来的内存版测试依赖 sessionRateLimiter 结构体，替换为 DB 版后重写。

func TestSessionRateAllowsUpToLimit(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()

	for i := 1; i <= sessionRateLimit; i++ {
		if ok, _ := AllowNewSession(ctx, "bot"); !ok {
			t.Fatalf("第 %d 次应放行（上限 %d）", i, sessionRateLimit)
		}
	}
	ok, retry := AllowNewSession(ctx, "bot")
	if ok {
		t.Fatal("超过上限应拦下")
	}
	if retry < 1 {
		t.Fatalf("应给出正的重试等待秒数，实际 %d", retry)
	}
}

func TestSessionRateIsPerAgent(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()

	for i := 0; i < sessionRateLimit; i++ {
		AllowNewSession(ctx, "busy")
	}
	if ok, _ := AllowNewSession(ctx, "busy"); ok {
		t.Fatal("busy 应已被拦")
	}
	if ok, _ := AllowNewSession(ctx, "idle"); !ok {
		t.Fatal("另一个 Agent 不该被牵连")
	}
}

// 并发请求不能把上限刷穿（判断与记账必须原子）
func TestSessionRateConcurrentDoesNotOverrun(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()

	var wg sync.WaitGroup
	var mu sync.Mutex
	passed := 0

	for i := 0; i < sessionRateLimit*4; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if ok, _ := AllowNewSession(ctx, "bot"); ok {
				mu.Lock()
				passed++
				mu.Unlock()
			}
		}()
	}
	wg.Wait()

	if passed != sessionRateLimit {
		t.Fatalf("%d 并发下放行 %d 次，期望恰好 %d 次",
			sessionRateLimit*4, passed, sessionRateLimit)
	}
}

// 建会话失败时要还名额
func TestSessionRateRelease(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()

	for i := 0; i < sessionRateLimit; i++ {
		AllowNewSession(ctx, "bot")
	}
	if ok, _ := AllowNewSession(ctx, "bot"); ok {
		t.Fatal("应已刷满")
	}
	ReleaseNewSession(ctx, "bot")
	if ok, _ := AllowNewSession(ctx, "bot"); !ok {
		t.Fatal("归还名额后应能再开一条")
	}
}

// 人类不走限速（空 agentName）
func TestAllowNewSessionSkipsHumans(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()

	for i := 0; i < sessionRateLimit*3; i++ {
		if ok, _ := AllowNewSession(ctx, ""); !ok {
			t.Fatal("人类不该被限速")
		}
	}
}

// 窗口滑过后自动恢复（过期记录自动清理）
func TestSessionRateWindowSlides(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()

	// 插入过期记录（1小时前）
	cutoff := time.Now().Add(-sessionRateWindow - time.Minute)
	for i := 0; i < sessionRateLimit; i++ {
		_, err := db.DB.ExecContext(ctx,
			`INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`,
			"session:bot", cutoff)
		if err != nil {
			t.Fatalf("插入过期记录失败: %v", err)
		}
	}

	// 窗口外的记录应被清理，此次应放行
	if ok, _ := AllowNewSession(ctx, "bot"); !ok {
		t.Fatal("窗口外的记录应被清掉，此次应放行")
	}
}