package proxy

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/openteam/server/internal/channel"
	"github.com/openteam/server/internal/store"
)

// 会话亲和：同一会话（同密钥+模型+同首条消息）粘住上次应答的渠道——
// 即使加权随机此时会选别的渠道；不同会话不受影响，照常负载均衡。
func TestAffinitySticksToAnsweringChannel(t *testing.T) {
	var ch1Hits, ch2Hits atomic.Int64
	ch1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ch1Hits.Add(1)
		_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"from-1"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
	}))
	defer ch1.Close()
	ch2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ch2Hits.Add(1)
		_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"from-2"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
	}))
	defer ch2.Close()

	h := newHarness(t, ch1.URL, 0, true, nil)
	h.gw.affinityTTL = 5 * time.Minute
	addChannel(t, h, "c2", ch2.URL, 1)
	pushWeightTo(t, h, "c1", 1000)

	sameConv := `{"model":"m1","messages":[{"role":"user","content":"会话的第一条消息"}]}`
	// 第一轮：权重压在渠道1，应答来自渠道1，亲和绑定建立
	if w := h.run(t, nil, sameConv); w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "from-1") {
		t.Fatalf("第一轮应来自渠道1：%d %s", w.Code, w.Body.String())
	}
	// 反转权重：加权随机此刻会选渠道2
	pushWeightTo(t, h, "c1", 1)
	pushWeightTo(t, h, "c2", 1000)

	// 同一会话第二轮：仍应粘在渠道1（缓存值得保留）
	if w := h.run(t, nil, sameConv); !strings.Contains(w.Body.String(), "from-1") {
		t.Fatalf("同一会话应粘住渠道1：%s", w.Body.String())
	}
	// 不同会话（不同首条消息）：走加权随机，落到渠道2
	otherConv := `{"model":"m1","messages":[{"role":"user","content":"另一个会话的开场"}]}`
	if w := h.run(t, nil, otherConv); !strings.Contains(w.Body.String(), "from-2") {
		t.Fatalf("不同会话应走加权随机落到渠道2：%s", w.Body.String())
	}
	if ch1Hits.Load() != 2 || ch2Hits.Load() != 1 {
		t.Fatalf("命中次数不符：ch1=%d ch2=%d", ch1Hits.Load(), ch2Hits.Load())
	}
}

// 亲和渠道进入运行期冷却后自动失效，回退加权随机；不会把请求发给冷却中的渠道。
func TestAffinityRespectsCooldown(t *testing.T) {
	ch1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("冷却中的渠道1 不应收到请求")
		_, _ = io.WriteString(w, `{"id":"bad"}`)
	}))
	defer ch1.Close()
	ch2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"from-2"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
	}))
	defer ch2.Close()

	h := newHarness(t, ch1.URL, 0, true, nil)
	h.gw.affinityTTL = 5 * time.Minute
	addChannel(t, h, "c2", ch2.URL, 1)

	// 预置亲和绑定指向渠道1，随后让渠道1 进运行期冷却
	var ch1row store.Channel
	h.db.Where("name = ?", "c1").First(&ch1row)
	h.gw.aff.put(affinityKey(1, "m1", "会话的第一条消息"), ch1row.ID, 5*time.Minute)
	for i := 0; i < 3; i++ {
		h.gw.ch.ReportFailure(&ch1row, channel.FailRate, 0)
	}

	w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"会话的第一条消息"}]}`)
	if !strings.Contains(w.Body.String(), "from-2") {
		t.Fatalf("冷却后亲和应失效并切换渠道2：%d %s", w.Code, w.Body.String())
	}
}