package proxy

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/openteam/server/internal/pkg/apikey"
	"github.com/openteam/server/internal/pkg/crypto"
	"github.com/openteam/server/internal/pkg/ratelimit"
	"github.com/openteam/server/internal/store"
	"github.com/openteam/server/internal/usage"
	"gorm.io/gorm"
)

type harness struct {
	gw  *Gateway
	db  *gorm.DB
	rec *usage.Recorder
	key string
}

// newHarness 起一个内存 SQLite 网关：1 用户 + 1 密钥 + 1 模型 + 1 渠道（可选绑定）。
func newHarness(t *testing.T, baseURL string, timeoutMS int, bind bool, keyModels []string) *harness {
	t.Helper()
	dsn := "/tmp/proxy_fix_test_" + time.Now().Format("150405.000000000") + ".db"
	db, err := store.Open("sqlite", dsn)
	if err != nil {
		t.Fatalf("open db: %v", err)
	}
	t.Cleanup(func() {
		sqlDB, _ := db.DB()
		sqlDB.Close()
		os.Remove(dsn)
	})

	enc := crypto.NewEncryptor("test-master-key-1234567890")
	encKey, _ := enc.Encrypt("upstream-key")
	u := store.User{Username: "u1", Email: "u1@x.com", PasswordHash: "x", Role: store.RoleUser,
		Balance: 100, Status: store.UserStatusActive}
	if err := db.Create(&u).Error; err != nil {
		t.Fatal(err)
	}
	plain, hash, prefix, _ := apikey.Generate()
	k := store.APIKey{UserID: u.ID, Name: "t", KeyHash: hash, KeyPrefix: prefix,
		Status: store.KeyStatusActive, AllowedModels: keyModels}
	if err := db.Create(&k).Error; err != nil {
		t.Fatal(err)
	}
	m := store.Model{Name: "m1", InputPrice: 0.15, OutputPrice: 0.6, Enabled: true}
	if err := db.Create(&m).Error; err != nil {
		t.Fatal(err)
	}
	ch := store.Channel{Name: "c1", Provider: store.ChannelProviderOpenAI,
		Formats: []string{store.FormatChat}, BaseURL: baseURL, APIKeyEnc: encKey,
		Weight: 1, TimeoutMS: timeoutMS, MaxConcurrency: 4,
		HealthStatus: store.ChannelHealthHealthy, Enabled: true}
	if err := db.Create(&ch).Error; err != nil {
		t.Fatal(err)
	}
	if bind {
		if err := db.Create(&store.ChannelModelBinding{ChannelID: ch.ID, ModelID: m.ID, UpstreamModel: "m1"}).Error; err != nil {
			t.Fatal(err)
		}
	}

	rec := usage.NewRecorder(db)
	t.Cleanup(rec.Close)
	return &harness{
		gw: NewGateway(db, enc, rec, ratelimit.New(), 0, false, 300*time.Second, 32),
		db: db, rec: rec, key: plain,
	}
}

// run 走完整链路：Auth（含新的密钥白名单注入）→ Handle。
func (h *harness) run(t *testing.T, ctx context.Context, body string) *httptest.ResponseRecorder {
	t.Helper()
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
	req.Header.Set("Authorization", "Bearer "+h.key)
	if ctx != nil {
		req = req.WithContext(ctx)
	}
	c.Request = req
	h.gw.Auth(c)
	if !c.IsAborted() {
		h.gw.Handle(c)
	}
	return w
}

// 修复①：渠道 timeout_ms 真正生效（此前被 int64 溢出成"永不超时"）。
func TestChannelTimeoutIsEnforced(t *testing.T) {
	var hits atomic.Int64
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		select {
		case <-r.Context().Done():
		case <-time.After(1 * time.Second):
		}
		_, _ = io.WriteString(w, `{"id":"late"}`)
	}))
	defer srv.Close()

	h := newHarness(t, srv.URL, 200, true, nil)
	start := time.Now()
	w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
	elapsed := time.Since(start)

	if elapsed > 600*time.Millisecond {
		t.Fatalf("渠道 timeout_ms=200 没生效（上游 1s 后才返回），实际耗时 %s", elapsed)
	}
	if w.Code != http.StatusBadGateway {
		t.Fatalf("超时后应报 upstream_error，实际 %d: %s", w.Code, w.Body.String())
	}
	if hits.Load() == 0 {
		t.Fatal("上游没被调用")
	}
}

// 修复①延伸：流式不受整段超时限制，但首包超时由 transport 兜住。
func TestStreamHasNoWholeBodyDeadline(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(200)
		f := w.(http.Flusher)
		for i := 0; i < 3; i++ {
			_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"x"}}]}`+"\n\n")
			f.Flush()
			time.Sleep(150 * time.Millisecond) // 总计 450ms > 渠道 timeout_ms
		}
		_, _ = io.WriteString(w, "data: [DONE]\n\n")
	}))
	defer srv.Close()

	h := newHarness(t, srv.URL, 200, true, nil)
	w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK {
		t.Fatalf("stream code=%d body=%s", w.Code, w.Body.String())
	}
	if n := strings.Count(w.Body.String(), `"content":"x"`); n != 3 {
		t.Fatalf("流被超时截断，只收到 %d 块，期望 3", n)
	}
}

// 修复②：客户端断连必须把取消传到上游，两侧都停止，而不是让上游继续生成。
func TestClientDisconnectCancelsUpstream(t *testing.T) {
	started := make(chan struct{})
	upstreamStop := make(chan error, 1)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		f := w.(http.Flusher)
		close(started)
		// 模拟持续生成：客户端断开后写会失败、request context 会被取消
		for i := 0; i < 40; i++ {
			if _, err := io.WriteString(w, `data: {"choices":[{"delta":{"content":"x"}}]}`+"\n\n"); err != nil {
				upstreamStop <- err
				return
			}
			f.Flush()
			if err := r.Context().Err(); err != nil {
				upstreamStop <- err
				return
			}
			time.Sleep(50 * time.Millisecond)
		}
		upstreamStop <- errors.New("上游跑完了 2s 都没察觉客户端断开")
	}))
	defer srv.Close()

	h := newHarness(t, srv.URL, 0, true, nil)
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan *httptest.ResponseRecorder, 1)
	go func() {
		done <- h.run(t, ctx, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
	}()

	<-started
	time.Sleep(120 * time.Millisecond) // 先收几块，保证已进入流式转发
	cancel()

	select {
	case err := <-upstreamStop:
		t.Logf("上游因客户端断开而停止：%v", err)
	case <-time.After(2 * time.Second):
		t.Fatal("客户端断开后上游仍在生成：context 没有传播")
	}
	<-done
	h.rec.Close()

	var n int64
	h.db.Model(&store.UsageLog{}).Where("status = ? AND error_code = ?",
		store.UsageStatusCanceled, "client_disconnect").Count(&n)
	if n != 1 {
		t.Fatalf("canceled 记账条数 = %d，期望 1", n)
	}
}

// 修复③：APIKey.AllowedModels 现在真的是道闸门。
func TestKeyModelWhitelistIsEnforced(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
	}))
	defer srv.Close()

	h := newHarness(t, srv.URL, 0, true, []string{"gpt-4o"})
	w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusForbidden {
		t.Fatalf("密钥白名单外的模型应 403，实际 %d: %s", w.Code, w.Body.String())
	}
	var body map[string]any
	_ = json.Unmarshal(w.Body.Bytes(), &body)
	errObj, _ := body["error"].(map[string]any)
	if errObj == nil || errObj["code"] != "model_not_allowed" {
		t.Fatalf("错误体不符：%v", body)
	}
}

// 修复④（渠道侧）：运行期连续失败让坏渠道从候选里消失。
func TestRuntimeCooldownEvictsFailingChannel(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusTooManyRequests)
		_, _ = io.WriteString(w, `{"error":{"message":"rate limited"}}`)
	}))
	defer srv.Close()

	h := newHarness(t, srv.URL, 0, true, nil)
	var ch store.Channel
	h.db.First(&ch)

	// 3 次 429：每次换不到别的渠道，最终 502，同时累计运行期失败
	body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
	for i := 0; i < 3; i++ {
		if got := len(h.gw.candidateChannels("m1")); got != 1 {
			t.Fatalf("第 %d 次失败前候选应为 1，实际 %d", i+1, got)
		}
		if w := h.run(t, nil, body); w.Code != http.StatusTooManyRequests {
			t.Fatalf("429 应作为最后状态返回，实际 %d: %s", w.Code, w.Body.String())
		}
	}
	if !h.gw.ch.InCooldown(ch.ID) {
		t.Fatal("3 次连续 429 后渠道应进入运行期冷却")
	}
	if got := h.gw.candidateChannels("m1"); len(got) != 0 {
		t.Fatalf("冷却中不应出现在候选集，实际 %d", len(got))
	}
	if w := h.run(t, nil, body); w.Code != http.StatusServiceUnavailable {
		t.Fatalf("冷却后应 503 no_channel，实际 %d: %s", w.Code, w.Body.String())
	}
}

// 修复④延伸：上游恢复正常后清空失败计数，不会带着旧账被误踢。
func TestRuntimeFailureCounterResetsOnSuccess(t *testing.T) {
	fail := true
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if fail {
			w.WriteHeader(http.StatusBadGateway)
			_, _ = io.WriteString(w, `{"error":{"message":"boom"}}`)
			return
		}
		_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
	}))
	defer srv.Close()

	h := newHarness(t, srv.URL, 0, true, nil)
	body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
	h.run(t, nil, body)
	h.run(t, nil, body)
	fail = false
	if w := h.run(t, nil, body); w.Code != http.StatusOK {
		t.Fatalf("应成功，实际 %d: %s", w.Code, w.Body.String())
	}
	var ch store.Channel
	h.db.First(&ch)
	// 成功清零后再失败两次仍不该进冷却（阈值是"连续"3 次）
	fail = true
	h.run(t, nil, body)
	h.run(t, nil, body)
	if h.gw.ch.InCooldown(ch.ID) {
		t.Fatal("失败计数没有被成功请求清零")
	}
}

// 回归：全局兜底超时值不再被二次换算。
func TestGatewayTimeoutIsNotMultiplied(t *testing.T) {
	gw := NewGateway(nil, nil, nil, nil, 0, false, 30*time.Second, 8)
	if got := gw.channelTimeout(&store.Channel{TimeoutMS: 0}); got != 30*time.Second {
		t.Fatalf("兜底超时 = %s，期望 30s", got)
	}
	if got := gw.channelTimeout(&store.Channel{TimeoutMS: 1500}); got != 1500*time.Millisecond {
		t.Fatalf("渠道超时 = %s，期望 1.5s", got)
	}
	if c := gw.clientFor(30 * time.Second); c.Timeout != 0 {
		t.Fatalf("client.Timeout 应为 0（由 context 控制整段时长），实际 %s", c.Timeout)
	}
}

// 请求体超过上限时应按客户端协议返回 413（上限由 /v1 组的 BodyLimit 中间件把住）。
func TestRequestBodyTooLarge(t *testing.T) {
	h := newHarness(t, "http://127.0.0.1:1", 0, true, nil)

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions",
		strings.NewReader(`{"model":"m1","messages":[{"role":"user","content":"`+strings.Repeat("x", 4096)+`"}]}`))
	req.Header.Set("Authorization", "Bearer "+h.key)
	// 模拟 middleware.BodyLimit：超限后 ReadAll 拿到 *http.MaxBytesError
	req.Body = http.MaxBytesReader(w, req.Body, 64)
	c.Request = req

	h.gw.Auth(c)
	if !c.IsAborted() {
		h.gw.Handle(c)
	}
	if w.Code != http.StatusRequestEntityTooLarge {
		t.Fatalf("应返回 413，实际 %d: %s", w.Code, w.Body.String())
	}
	if !strings.Contains(w.Body.String(), "request_too_large") {
		t.Fatalf("错误码不符：%s", w.Body.String())
	}
	if !strings.Contains(w.Body.String(), "32 MB") {
		t.Fatalf("报错应点名上限，实际：%s", w.Body.String())
	}
}