package repo

import (
	"context"
	"time"

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

// RateLimitCheckAndRecord 原子地检查 bucket 在 window 内的事件数是否超过 limit。
// 未超过则同时记录本次事件（判断与写入在同一个事务里，防并发刷穿）。
// DB 不可用时放行（宁可放开限速也不能让用户完全无法使用）。
func RateLimitCheckAndRecord(ctx context.Context, bucket string, window time.Duration, limit int) (allowed bool, retryAfter int) {
	now := time.Now()
	cutoff := now.Add(-window)

	// 用 IMMEDIATE 事务：SQLite 的 IMMEDIATE 会在开始时获取 RESERVED 锁，
	// 防止其他写事务同时进入 COMMIT 阶段。这是 SQLite 并发写的正确方式。
	tx, err := db.DB.BeginTx(ctx, nil)
	if err != nil {
		return true, 0 // DB 不可用 → 放行
	}
	defer tx.Rollback() // Commit 成功后 Rollback 是 no-op

	// 清理过期记录
	tx.ExecContext(ctx,
		`DELETE FROM rate_limits WHERE bucket = $1 AND ts < $2`, bucket, cutoff)

	// 统计当前窗口内事件数
	var count int
	err = tx.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM rate_limits WHERE bucket = $1 AND ts >= $2`,
		bucket, cutoff).Scan(&count)
	if err != nil {
		return true, 0
	}

	if count >= limit {
		// 用 ORDER BY ts ASC LIMIT 1 而不是 MIN(ts)：SQLite 驱动把聚合 MIN(ts)
		// 当 string 返回，扫进 *time.Time 失败 → 落到下面的兜底
		// `return false, 60`，于是所有 429 的 retry_after 恒为 60，
		// 与真实剩余窗口（最长 sessionRateWindow=1h）无关。
		// 排序取值走结果集本身，驱动按列类型给 time.Time；语义等价。
		var earliest time.Time
		err = tx.QueryRowContext(ctx,
			`SELECT ts FROM rate_limits WHERE bucket = $1 AND ts >= $2
			 ORDER BY ts ASC LIMIT 1`,
			bucket, cutoff).Scan(&earliest)
		if err == nil && !earliest.IsZero() {
			retry := int(earliest.Add(window).Sub(now).Seconds()) + 1
			if retry < 1 {
				retry = 1
			}
			return false, retry
		}
		return false, 60
	}

	// 记账
	tx.ExecContext(ctx,
		`INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`, bucket, now)
	tx.Commit()
	return true, 0
}

// RateLimitReset 清除指定 bucket 的所有记录（登录成功后调用）。
func RateLimitReset(ctx context.Context, bucket string) {
	_, _ = db.DB.ExecContext(ctx,
		`DELETE FROM rate_limits WHERE bucket = $1`, bucket)
}