package upgrade

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/version"
)

// resetServerRecommendationState clears all global cache state between tests.
func resetServerRecommendationState() {
	serverRecMu.Lock()
	serverRecs = make(map[ServerKey]*cachedRecommendation)
	activeServerKey = ""
	serverRefresher = nil
	serverInflight = make(map[ServerKey]chan struct{})
	serverRecMu.Unlock()
}

func TestNormalizeServerKey(t *testing.T) {
	require.Equal(t, ServerKey("https://a.com|wsp1"), NormalizeServerKey("https://a.com/", "wsp1"))
	require.Equal(t, ServerKey("https://a.com|"), NormalizeServerKey(" https://a.com ", ""))
	require.Equal(t, ServerKey("https://a.com|wsp1"), NormalizeServerKey("https://a.com//", "wsp1"))
}

func TestSetAndClearServerRecommendation(t *testing.T) {
	resetServerRecommendationState()
	key := NormalizeServerKey("https://server1", "wsp1")

	SetServerRecommendation(key, &ServerRecommendation{
		ServerVersion:             "v1.0.0",
		RecommendedDesktopVersion: "v0.2.5",
		MinDesktopVersion:         "v0.2.0",
	})

	SetActiveServer(key)
	rec, err := EnsureServerRecommendation(context.Background())
	require.NoError(t, err)
	require.NotNil(t, rec)
	require.Equal(t, "v0.2.5", rec.RecommendedDesktopVersion)

	ClearServerRecommendation(key)
	SetActiveServer("")
	rec, err = EnsureServerRecommendation(context.Background())
	require.NoError(t, err)
	require.Nil(t, rec)
}

func TestEnsureServerRecommendationReturnsNilWhenNoActiveServer(t *testing.T) {
	resetServerRecommendationState()
	rec, err := EnsureServerRecommendation(context.Background())
	require.NoError(t, err)
	require.Nil(t, rec)
}

func TestEnsureServerRecommendationReturnsNilWhenNoRefresher(t *testing.T) {
	resetServerRecommendationState()
	SetActiveServer(NormalizeServerKey("https://server1", "wsp1"))
	defer SetActiveServer("")

	rec, err := EnsureServerRecommendation(context.Background())
	require.NoError(t, err)
	require.Nil(t, rec)
}

func TestEnsureServerRecommendationTriggersRefresherWhenCacheEmpty(t *testing.T) {
	resetServerRecommendationState()
	key := NormalizeServerKey("https://server1", "wsp1")
	SetActiveServer(key)
	defer SetActiveServer("")

	var calls int32
	RegisterRecommendationRefresher(func(ctx context.Context, sk string) (*ServerRecommendation, error) {
		atomic.AddInt32(&calls, 1)
		require.Equal(t, string(key), sk)
		return &ServerRecommendation{
			ServerVersion:             "v1.0.0",
			RecommendedDesktopVersion: "v0.2.5",
		}, nil
	})
	defer RegisterRecommendationRefresher(nil)

	rec, err := EnsureServerRecommendation(context.Background())
	require.NoError(t, err)
	require.NotNil(t, rec)
	require.Equal(t, "v0.2.5", rec.RecommendedDesktopVersion)
	require.Equal(t, int32(1), atomic.LoadInt32(&calls))

	// Second call should use cache, not refresher.
	rec, err = EnsureServerRecommendation(context.Background())
	require.NoError(t, err)
	require.NotNil(t, rec)
	require.Equal(t, int32(1), atomic.LoadInt32(&calls))
}

func TestEnsureServerRecommendationCoalescesConcurrentFetches(t *testing.T) {
	resetServerRecommendationState()
	key := NormalizeServerKey("https://server1", "wsp1")
	SetActiveServer(key)
	defer SetActiveServer("")

	var calls int32
	started := make(chan struct{})
	release := make(chan struct{})
	RegisterRecommendationRefresher(func(ctx context.Context, sk string) (*ServerRecommendation, error) {
		atomic.AddInt32(&calls, 1)
		close(started)
		<-release
		return &ServerRecommendation{RecommendedDesktopVersion: "v0.2.5"}, nil
	})
	defer RegisterRecommendationRefresher(nil)

	var wg sync.WaitGroup
	results := make([]*ServerRecommendation, 3)
	errs := make([]error, 3)
	for i := 0; i < 3; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			results[idx], errs[idx] = EnsureServerRecommendation(context.Background())
		}(i)
	}

	<-started
	close(release)
	wg.Wait()

	require.Equal(t, int32(1), atomic.LoadInt32(&calls))
	for i := 0; i < 3; i++ {
		require.NoError(t, errs[i])
		require.NotNil(t, results[i])
		require.Equal(t, "v0.2.5", results[i].RecommendedDesktopVersion)
	}
}

func TestEnsureServerRecommendationDoesNotLeakAcrossServerKeys(t *testing.T) {
	resetServerRecommendationState()
	key1 := NormalizeServerKey("https://server1", "wsp1")
	key2 := NormalizeServerKey("https://server2", "wsp2")

	SetServerRecommendation(key1, &ServerRecommendation{RecommendedDesktopVersion: "v0.2.5"})
	SetActiveServer(key2)
	defer SetActiveServer("")

	RegisterRecommendationRefresher(func(ctx context.Context, sk string) (*ServerRecommendation, error) {
		// key2 should trigger a fetch, not reuse key1's cache.
		require.Equal(t, string(key2), sk)
		return &ServerRecommendation{RecommendedDesktopVersion: "v1.1.3"}, nil
	})
	defer RegisterRecommendationRefresher(nil)

	rec, err := EnsureServerRecommendation(context.Background())
	require.NoError(t, err)
	require.NotNil(t, rec)
	require.Equal(t, "v1.1.3", rec.RecommendedDesktopVersion)
}

func TestCheckVersionUsesServerRecommendationAsLatest(t *testing.T) {
	resetServerRecommendationState()
	setUpgradeTestPlatform(t, "linux-amd64", true, true)

	oldVersion := version.Version
	version.Version = "v0.2.0"
	defer func() { version.Version = oldVersion }()

	mux := http.NewServeMux()
	server := httptest.NewServer(mux)
	defer server.Close()

	mux.HandleFunc("/stable.json", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, `{"schema_version":"1","product":"agentteam-daemon","channel":"stable","current_version":"v0.9.0","published_at":"2026-04-16T04:47:16Z","manifest_url":"%s/v0.9.0.json"}`, server.URL)
	})
	mux.HandleFunc("/v0.9.0.json", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, `{"schema_version":"1","product":"agentteam-daemon","version":"v0.9.0","published_at":"2026-04-16T04:47:16Z","manifest_url":"%s/v0.9.0.json","client":{"install_script":"https://static.pingcap.com/agentteam/client/install.sh","checksums":"https://static.pingcap.com/agentteam/client/v0.9.0/agentteam_0.9.0_checksums.txt"}}`, server.URL)
	})
	// The pinned version manifest must also be served.
	mux.HandleFunc("/v0.2.5.json", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, `{"schema_version":"1","product":"agentteam-daemon","version":"v0.2.5","published_at":"2026-04-16T04:47:16Z","manifest_url":"%s/v0.2.5.json","client":{"install_script":"https://static.pingcap.com/agentteam/client/install.sh","checksums":"https://static.pingcap.com/agentteam/client/v0.2.5/agentteam_0.2.5_checksums.txt"}}`, server.URL)
	})

	key := NormalizeServerKey("https://private-server", "wsp1")
	SetServerRecommendation(key, &ServerRecommendation{
		ServerVersion:             "v1.0.0",
		RecommendedDesktopVersion: "v0.2.5",
		MinDesktopVersion:         "v0.2.0",
	})
	SetActiveServer(key)
	defer func() {
		SetActiveServer("")
		ClearServerRecommendation(key)
	}()

	result, err := CheckVersion(context.Background(), server.URL+"/stable.json")
	require.NoError(t, err)
	// Latest should be the server-recommended v0.2.5, not the CDN v0.9.0.
	require.Equal(t, "v0.2.5", result.Latest)
	require.True(t, result.HasUpdate)
	require.Equal(t, "v1.0.0", result.ServerVersion)
	require.Equal(t, "v0.2.0", result.MinDesktopVersion)
}

func TestCheckVersionFallsBackToCDNWhenNoActiveServer(t *testing.T) {
	resetServerRecommendationState()
	setUpgradeTestPlatform(t, "linux-amd64", true, true)

	oldVersion := version.Version
	version.Version = "v1.0.1"
	defer func() { version.Version = oldVersion }()

	mux := http.NewServeMux()
	server := httptest.NewServer(mux)
	defer server.Close()

	mux.HandleFunc("/stable.json", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, `{"schema_version":"1","product":"agentteam-daemon","channel":"stable","current_version":"v1.0.2-rc.1","published_at":"2026-04-16T04:47:16Z","manifest_url":"%s/v1.0.2-rc.1.json"}`, server.URL)
	})
	mux.HandleFunc("/v1.0.2-rc.1.json", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, `{"schema_version":"1","product":"agentteam-daemon","version":"v1.0.2-rc.1","published_at":"2026-04-16T04:47:16Z","manifest_url":"%s/v1.0.2-rc.1.json","client":{"install_script":"https://static.pingcap.com/agentteam/client/install.sh","checksums":"https://static.pingcap.com/agentteam/client/v1.0.2-rc.1/agentteam_1.0.2-rc.1_checksums.txt"}}`, server.URL)
	})

	// No active server, no refresher — should use CDN latest.
	result, err := CheckVersion(context.Background(), server.URL+"/stable.json")
	require.NoError(t, err)
	require.Equal(t, "v1.0.2-rc.1", result.Latest)
	require.Empty(t, result.ServerVersion)
}

func TestCheckVersionTriggersRefresherWhenConnectedButCacheEmpty(t *testing.T) {
	resetServerRecommendationState()
	setUpgradeTestPlatform(t, "linux-amd64", true, true)

	oldVersion := version.Version
	version.Version = "v0.2.0"
	defer func() { version.Version = oldVersion }()

	mux := http.NewServeMux()
	server := httptest.NewServer(mux)
	defer server.Close()

	mux.HandleFunc("/stable.json", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, `{"schema_version":"1","product":"agentteam-daemon","channel":"stable","current_version":"v0.9.0","published_at":"2026-04-16T04:47:16Z","manifest_url":"%s/v0.9.0.json"}`, server.URL)
	})
	mux.HandleFunc("/v0.9.0.json", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, `{"schema_version":"1","product":"agentteam-daemon","version":"v0.9.0","published_at":"2026-04-16T04:47:16Z","manifest_url":"%s/v0.9.0.json","client":{"install_script":"https://static.pingcap.com/agentteam/client/install.sh","checksums":"https://static.pingcap.com/agentteam/client/v0.9.0/agentteam_0.9.0_checksums.txt"}}`, server.URL)
	})
	mux.HandleFunc("/v0.2.5.json", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, `{"schema_version":"1","product":"agentteam-daemon","version":"v0.2.5","published_at":"2026-04-16T04:47:16Z","manifest_url":"%s/v0.2.5.json","client":{"install_script":"https://static.pingcap.com/agentteam/client/install.sh","checksums":"https://static.pingcap.com/agentteam/client/v0.2.5/agentteam_0.2.5_checksums.txt"}}`, server.URL)
	})

	key := NormalizeServerKey("https://private-server", "wsp1")
	SetActiveServer(key)
	defer SetActiveServer("")

	var refresherCalled int32
	RegisterRecommendationRefresher(func(ctx context.Context, sk string) (*ServerRecommendation, error) {
		atomic.AddInt32(&refresherCalled, 1)
		return &ServerRecommendation{
			ServerVersion:             "v1.0.0",
			RecommendedDesktopVersion: "v0.2.5",
		}, nil
	})
	defer RegisterRecommendationRefresher(nil)

	result, err := CheckVersion(context.Background(), server.URL+"/stable.json")
	require.NoError(t, err)
	require.Equal(t, int32(1), atomic.LoadInt32(&refresherCalled))
	require.Equal(t, "v0.2.5", result.Latest)
	require.Equal(t, "v1.0.0", result.ServerVersion)
}

func TestEnsureServerRecommendationRespectsContextCancellation(t *testing.T) {
	resetServerRecommendationState()
	key := NormalizeServerKey("https://server1", "wsp1")
	SetActiveServer(key)
	defer SetActiveServer("")

	RegisterRecommendationRefresher(func(ctx context.Context, sk string) (*ServerRecommendation, error) {
		<-ctx.Done()
		return nil, ctx.Err()
	})
	defer RegisterRecommendationRefresher(nil)

	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	rec, err := EnsureServerRecommendation(ctx)
	require.Error(t, err)
	require.Nil(t, rec)
}