package loadbalance

import (
	"sync"
	"testing"
	"time"

	"github.com/valyala/fasthttp"
)

// TestStickySession_BasicRoute 测试基本的会话粘性路由。
// 第一次请求应设置 cookie，第二次携带相同 cookie 应路由到同一目标。
func TestStickySession_BasicRoute(t *testing.T) {
	t.Parallel()
	t.Run("首次请求设置cookie并路由", func(_ *testing.T) {
		config := DefaultStickyConfig()
		config.Enabled = true
		fallback := NewRoundRobin()
		sticky := NewStickySession(config, fallback)
		sticky.Start()
		defer sticky.Stop()

		targets := []*Target{
			createHealthyTarget("http://backend1:8080", true),
			createHealthyTarget("http://backend2:8080", true),
		}

		ctx := &fasthttp.RequestCtx{}
		got := sticky.Select(ctx, targets)
		if got == nil {
			t.Fatal("Select() = nil, want non-nil")
		}

		// 验证设置了 cookie
		cookieValue := ctx.Response.Header.PeekCookie(config.Name)
		if len(cookieValue) == 0 {
			t.Error("首次请求未设置 cookie")
		}
	})

	t.Run("相同cookie路由到同一目标", func(_ *testing.T) {
		config := DefaultStickyConfig()
		config.Enabled = true
		fallback := NewRoundRobin()
		sticky := NewStickySession(config, fallback)
		sticky.Start()
		defer sticky.Stop()

		targets := []*Target{
			createHealthyTarget("http://backend1:8080", true),
			createHealthyTarget("http://backend2:8080", true),
		}

		// 第一次请求
		ctx1 := &fasthttp.RequestCtx{}
		got1 := sticky.Select(ctx1, targets)
		if got1 == nil {
			t.Fatal("第一次 Select() = nil")
		}

		// 提取 cookie
		cookie := &fasthttp.Cookie{}
		cookie.SetKey(config.Name)
		if err := cookie.ParseBytes(ctx1.Response.Header.PeekCookie(config.Name)); err != nil {
			t.Fatalf("解析 cookie 失败: %v", err)
		}

		// 第二次请求携带相同 cookie
		ctx2 := &fasthttp.RequestCtx{}
		ctx2.Request.Header.SetCookie(config.Name, string(cookie.Value()))
		got2 := sticky.Select(ctx2, targets)
		if got2 == nil {
			t.Fatal("第二次 Select() = nil")
		}

		if got2.URL != got1.URL {
			t.Errorf("相同 cookie 路由到不同目标: 第一次=%q, 第二次=%q", got1.URL, got2.URL)
		}
	})

	t.Run("禁用时不设置cookie", func(_ *testing.T) {
		config := DefaultStickyConfig()
		config.Enabled = false
		fallback := NewRoundRobin()
		sticky := NewStickySession(config, fallback)
		sticky.Start()
		defer sticky.Stop()

		targets := []*Target{
			createHealthyTarget("http://backend1:8080", true),
		}

		ctx := &fasthttp.RequestCtx{}
		got := sticky.Select(ctx, targets)
		if got == nil {
			t.Fatal("Select() = nil")
		}

		cookieValue := ctx.Response.Header.PeekCookie(config.Name)
		if len(cookieValue) > 0 {
			t.Error("禁用时不应设置 cookie")
		}
	})
}

// TestStickySession_TargetUnavailable 测试目标不可用时回退到 fallback。
func TestStickySession_TargetUnavailable(t *testing.T) {
	t.Parallel()
	t.Run("目标不健康时回退", func(_ *testing.T) {
		config := DefaultStickyConfig()
		config.Enabled = true
		fallback := NewRoundRobin()
		sticky := NewStickySession(config, fallback)
		sticky.Start()
		defer sticky.Stop()

		targets := []*Target{
			createHealthyTarget("http://backend1:8080", true),
			createHealthyTarget("http://backend2:8080", true),
		}

		// 第一次请求，记录会话
		ctx1 := &fasthttp.RequestCtx{}
		got1 := sticky.Select(ctx1, targets)
		if got1 == nil {
			t.Fatal("第一次 Select() = nil")
		}

		// 提取 cookie
		cookie := &fasthttp.Cookie{}
		cookie.SetKey(config.Name)
		if err := cookie.ParseBytes(ctx1.Response.Header.PeekCookie(config.Name)); err != nil {
			t.Fatalf("解析 cookie 失败: %v", err)
		}

		// 使之前选中的目标不健康
		for _, target := range targets {
			if target.URL == got1.URL {
				target.Healthy.Store(false)
				break
			}
		}

		// 第二次请求，应回退到其他目标
		ctx2 := &fasthttp.RequestCtx{}
		ctx2.Request.Header.SetCookie(config.Name, string(cookie.Value()))
		got2 := sticky.Select(ctx2, targets)
		if got2 == nil {
			t.Fatal("第二次 Select() = nil")
		}

		if got2.URL == got1.URL {
			t.Errorf("不健康目标未回退: %q", got2.URL)
		}
	})
}

// TestStickySession_CookieEncodeDecode 测试 cookie 编解码。
func TestStickySession_CookieEncodeDecode(t *testing.T) {
	t.Parallel()
	t.Run("编码解码round-trip", func(_ *testing.T) {
		url := "http://backend1:8080"
		expires := time.Now().Add(time.Hour)
		encoded := encodeStickyCookie(url, expires)
		if encoded == "" {
			t.Fatal("encodeStickyCookie() 返回空字符串")
		}

		decodedURL, decodedExpires, ok := decodeStickyCookie(encoded)
		if !ok {
			t.Fatal("decodeStickyCookie() returned ok=false")
		}

		if decodedURL != url {
			t.Errorf("解码后 URL = %q, want %q", decodedURL, url)
		}
		if decodedExpires.Unix() != expires.Unix() {
			t.Errorf("解码后 expires = %v, want %v", decodedExpires, expires)
		}
	})

	t.Run("空URL编码解码", func(_ *testing.T) {
		expires := time.Now().Add(time.Hour)
		encoded := encodeStickyCookie("", expires)
		decodedURL, decodedExpires, ok := decodeStickyCookie(encoded)
		if !ok {
			t.Fatal("decodeStickyCookie() returned ok=false")
		}
		if decodedURL != "" {
			t.Errorf("解码后 URL = %q, want 空字符串", decodedURL)
		}
		if decodedExpires.Unix() != expires.Unix() {
			t.Errorf("解码后 expires = %v, want %v", decodedExpires, expires)
		}
	})

	t.Run("无效编码", func(_ *testing.T) {
		decodedURL, decodedExpires, ok := decodeStickyCookie("invalid-base64!!!")
		if ok {
			t.Errorf("decodeStickyCookie() = (%q, %v, %v), want ok=false", decodedURL, decodedExpires, ok)
		}
	})
}

// TestStickySession_Concurrent 测试并发安全。
// 100 个 goroutine 同时访问会话存储。
func TestStickySession_Concurrent(t *testing.T) {
	t.Parallel()
	config := DefaultStickyConfig()
	config.Enabled = true
	fallback := NewRoundRobin()
	sticky := NewStickySession(config, fallback)
	sticky.Start()
	defer sticky.Stop()

	targets := []*Target{
		createHealthyTarget("http://backend1:8080", true),
		createHealthyTarget("http://backend2:8080", true),
		createHealthyTarget("http://backend3:8080", true),
	}

	var wg sync.WaitGroup
	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			ctx := &fasthttp.RequestCtx{}
			// 交替使用有 cookie 和没有 cookie 的请求
			if idx%2 == 0 {
				ctx.Request.Header.SetCookie(config.Name, encodeStickyCookie("http://backend1:8080", time.Now().Add(time.Hour)))
			}
			got := sticky.Select(ctx, targets)
			if got == nil {
				t.Error("并发 Select() = nil")
			}
		}(i)
	}
	wg.Wait()
}

// TestStickySession_ExpiredCookie 测试过期 cookie 会导致回退到 fallback。
func TestStickySession_ExpiredCookie(t *testing.T) {
	t.Parallel()
	fallback := NewRoundRobin()
	config := DefaultStickyConfig()
	config.Enabled = true
	config.Expires = -time.Hour // Negative = already expired

	sticky := NewStickySession(config, fallback)
	sticky.Start()
	defer sticky.Stop()

	// Make backend1 unavailable so we know for sure fallback picks backend2
	targets := []*Target{
		createHealthyTarget("http://backend1:8080", false),
		createHealthyTarget("http://backend2:8080", true),
	}

	// Create an expired cookie manually
	expiredCookie := encodeStickyCookie("http://backend1:8080", time.Now().Add(-time.Hour))

	ctx := &fasthttp.RequestCtx{}
	ctx.Request.Header.SetCookie(config.Name, expiredCookie)

	// Should fallback because cookie is expired (and even if not, backend1 is unavailable)
	selected := sticky.Select(ctx, targets)
	if selected == nil {
		t.Fatal("expected a target")
	}
	// Should not route to backend1 because cookie expired / unavailable
	if selected.URL == "http://backend1:8080" {
		t.Error("should not route using expired cookie")
	}
	// Should set a new cookie
	newCookie := ctx.Response.Header.PeekCookie(config.Name)
	if len(newCookie) == 0 {
		t.Error("expected new cookie to be set")
	}
}

// TestStickySession_DoubleStop 测试重复调用 Stop() 不会 panic。
func TestStickySession_DoubleStop(t *testing.T) {
	t.Parallel()
	fallback := NewRoundRobin()
	config := DefaultStickyConfig()

	sticky := NewStickySession(config, fallback)
	sticky.Start()
	sticky.Stop()
	sticky.Stop() // Should not panic
}

// TestStickySession_NilFallback 测试 nil fallback 会使用默认 RoundRobin。
func TestStickySession_NilFallback(t *testing.T) {
	t.Parallel()
	config := DefaultStickyConfig()
	config.Enabled = true

	sticky := NewStickySession(config, nil)
	sticky.Start()
	defer sticky.Stop()

	targets := []*Target{
		createHealthyTarget("http://backend1:8080", true),
	}

	ctx := &fasthttp.RequestCtx{}
	selected := sticky.Select(ctx, targets)
	if selected == nil {
		t.Fatal("expected a target with default fallback")
	}
}

// TestStickySession_SelectExcluding 测试排除选择委托给 fallback。
func TestStickySession_SelectExcluding(t *testing.T) {
	t.Parallel()
	t.Run("SelectExcluding委托给fallback", func(_ *testing.T) {
		config := DefaultStickyConfig()
		config.Enabled = true
		fallback := NewRoundRobin()
		sticky := NewStickySession(config, fallback)
		sticky.Start()
		defer sticky.Stop()

		targets := []*Target{
			createHealthyTarget("http://backend1:8080", true),
			createHealthyTarget("http://backend2:8080", true),
		}

		excluded := []*Target{targets[0]}
		got := sticky.SelectExcluding(targets, excluded)
		if got == nil {
			t.Fatal("SelectExcluding() = nil")
		}
		if got.URL != "http://backend2:8080" {
			t.Errorf("SelectExcluding() = %q, want %q", got.URL, "http://backend2:8080")
		}
	})
}
