package proxy

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

	"github.com/openteam/server/internal/channel"
	"gorm.io/gorm"
	gormlogger "gorm.io/gorm/logger"
)

// sqlTracer 记录一次请求路径上打过的 SQL。
type sqlTracer struct {
	gormlogger.Interface
	mu   sync.Mutex
	list []string
}

func (s *sqlTracer) Trace(_ context.Context, _ time.Time, fc func() (string, int64), _ error) {
	sql, _ := fc()
	s.mu.Lock()
	s.list = append(s.list, sql)
	s.mu.Unlock()
}

func (s *sqlTracer) take() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := s.list
	s.list = nil
	return out
}

// 每请求 SQL 预算：热路径上的模型/渠道目录必须走进程内缓存。
// 原本一次非流式请求要打 10 条（含同一模型查三次），缓存后只剩鉴权两条。
// 这个上限是防回归的闸门：往请求路径上新增未缓存的查询会立刻让它变红。
func TestProxyRequestSQLBudget(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = io.WriteString(w, `{"id":"x","model":"m1","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`)
	}))
	defer srv.Close()

	h := newHarness(t, srv.URL, 0, true, nil)
	tr := &sqlTracer{Interface: gormlogger.Discard}
	traced := h.db.Session(&gorm.Session{Logger: tr})
	h.gw.db = traced
	h.gw.ch = channel.NewService(traced, h.gw.enc)

	body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
	// 第一次：冷缓存 + last_used_at 首次写入，不计入预算
	if w := h.run(t, nil, body); w.Code != http.StatusOK {
		t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
	}
	tr.take()

	// 第二次：模型/渠道目录与系统配置都应命中缓存
	if w := h.run(t, nil, body); w.Code != http.StatusOK {
		t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
	}
	sqls := tr.take()
	const budget = 3 // api_keys + users（+ 偶发 last_used_at）
	if len(sqls) > budget {
		t.Fatalf("热路径 SQL 条数 = %d，超过预算 %d：\n%s", len(sqls), budget, joinSQL(sqls))
	}
	t.Logf("第二条请求 SQL = %d 条：\n%s", len(sqls), joinSQL(sqls))
}

func joinSQL(sqls []string) string {
	out := ""
	for i, q := range sqls {
		out += "  " + string(rune('1'+i)) + ". " + q + "\n"
	}
	return out
}

// 开启鉴权缓存后，命中期间请求路径上一条 SQL 都不该有（缓存的是密钥+用户行，
// 状态/过期仍在内存里按当前时间判定）。
func TestProxyRequestSQLBudgetWithAuthCache(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = io.WriteString(w, `{"id":"x","model":"m1","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5}}`)
	}))
	defer srv.Close()

	h := newHarness(t, srv.URL, 0, true, nil)
	h.gw.auth = newAuthCache(time.Minute)
	tr := &sqlTracer{Interface: gormlogger.Discard}
	traced := h.db.Session(&gorm.Session{Logger: tr})
	h.gw.db = traced
	h.gw.ch = channel.NewService(traced, h.gw.enc)

	body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
	if w := h.run(t, nil, body); w.Code != http.StatusOK {
		t.Fatalf("预热请求失败：%d", w.Code)
	}
	tr.take()
	if w := h.run(t, nil, body); w.Code != http.StatusOK {
		t.Fatalf("code=%d", w.Code)
	}
	if sqls := tr.take(); len(sqls) != 0 {
		t.Fatalf("鉴权缓存命中时不应有 SQL，实际 %d 条：\n%s", len(sqls), joinSQL(sqls))
	}
}