package main

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

	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/pkg/autoscale/app"
	"github.com/tidbcloud/cluster-service-ng/pkg/autoscale/modules/common"
	"google.golang.org/protobuf/encoding/protojson"
)

// --- Pure function tests ---

func TestResolveCloudProvider(t *testing.T) {
	tests := []struct {
		name     string
		pools    []*apiv1.SharedPool
		expected string
	}{
		{
			name:     "nil pools",
			pools:    nil,
			expected: "",
		},
		{
			name:     "empty pools",
			pools:    []*apiv1.SharedPool{},
			expected: "",
		},
		{
			name: "AWS",
			pools: []*apiv1.SharedPool{
				{Provider: &apiv1.Provider{Type: apiv1.ProviderType_AWS}},
			},
			expected: "aws",
		},
		{
			name: "GCP",
			pools: []*apiv1.SharedPool{
				{Provider: &apiv1.Provider{Type: apiv1.ProviderType_GCP}},
			},
			expected: "gcp",
		},
		{
			name: "Azure",
			pools: []*apiv1.SharedPool{
				{Provider: &apiv1.Provider{Type: apiv1.ProviderType_AZURE}},
			},
			expected: "azure",
		},
		{
			name: "Alicloud",
			pools: []*apiv1.SharedPool{
				{Provider: &apiv1.Provider{Type: apiv1.ProviderType_ALICLOUD}},
			},
			expected: "alicloud",
		},
		{
			name: "nil provider skipped",
			pools: []*apiv1.SharedPool{
				{},
				{Provider: &apiv1.Provider{Type: apiv1.ProviderType_GCP}},
			},
			expected: "gcp",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := resolveCloudProvider(tt.pools)
			if got != tt.expected {
				t.Errorf("resolveCloudProvider() = %q, want %q", got, tt.expected)
			}
		})
	}
}

func TestFindControlPlaneBaseURL(t *testing.T) {
	tests := []struct {
		name     string
		ecfg     *engineConfig
		expected string
	}{
		{
			name:     "no inputs",
			ecfg:     &engineConfig{},
			expected: "",
		},
		{
			name: "found in first input",
			ecfg: &engineConfig{
				Autoscale: app.AutoscaleConfig{
					Inputs: []app.InputConfig{
						{Config: map[string]any{"controlPlaneBaseURL": "http://cp:8080"}},
					},
				},
			},
			expected: "http://cp:8080",
		},
		{
			name: "nil config skipped",
			ecfg: &engineConfig{
				Autoscale: app.AutoscaleConfig{
					Inputs: []app.InputConfig{
						{Config: nil},
						{Config: map[string]any{"controlPlaneBaseURL": "http://cp2:8080"}},
					},
				},
			},
			expected: "http://cp2:8080",
		},
		{
			name: "empty string skipped",
			ecfg: &engineConfig{
				Autoscale: app.AutoscaleConfig{
					Inputs: []app.InputConfig{
						{Config: map[string]any{"controlPlaneBaseURL": ""}},
					},
				},
			},
			expected: "",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := findControlPlaneBaseURL(tt.ecfg)
			if got != tt.expected {
				t.Errorf("findControlPlaneBaseURL() = %q, want %q", got, tt.expected)
			}
		})
	}
}

func TestFirstScalerUniqName(t *testing.T) {
	tests := []struct {
		name     string
		cfg      *app.AutoscaleConfig
		expected string
	}{
		{
			name:     "no scalers",
			cfg:      &app.AutoscaleConfig{},
			expected: "",
		},
		{
			name: "single scaler",
			cfg: &app.AutoscaleConfig{
				Scalers: []app.ScalerConfig{{UniqName: "my-scaler"}},
			},
			expected: "my-scaler",
		},
		{
			name: "multiple scalers returns first",
			cfg: &app.AutoscaleConfig{
				Scalers: []app.ScalerConfig{
					{UniqName: "first"},
					{UniqName: "second"},
				},
			},
			expected: "first",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := firstScalerUniqName(tt.cfg)
			if got != tt.expected {
				t.Errorf("firstScalerUniqName() = %q, want %q", got, tt.expected)
			}
		})
	}
}

// --- discover() integration tests using httptest ---

// newTestPoolManager creates a poolManager with a fake cpClient pointing at the
// given httptest server. The app and scalerRefs are nil — discover() will fail
// when trying to start a scaler, but we can still verify the reconcile logic
// by inspecting pm.running and checking cancel calls.
func newTestPoolManager(serverURL string) *poolManager {
	return &poolManager{
		cpClient:       common.NewControlPlaneClient(serverURL, 5*time.Second),
		scalerUniqName: "test-scaler",
		ecfg:           &engineConfig{},
		running:        make(map[string]*scalerHandle),
	}
}

func serveListSharedPools(pools []*apiv1.SharedPool) *httptest.Server {
	resp := &apiv1.ListSharedPoolsResponse{SharedPools: pools}
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		data, _ := protojson.Marshal(resp)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write(data)
	}))
}

func TestDiscover_EmptyPools(t *testing.T) {
	srv := serveListSharedPools(nil)
	defer srv.Close()

	pm := newTestPoolManager(srv.URL)
	pm.discover(context.Background())

	if len(pm.running) != 0 {
		t.Errorf("expected 0 running, got %d", len(pm.running))
	}
}

func TestDiscover_AddPool_ExistingNotRestarted(t *testing.T) {
	pools := []*apiv1.SharedPool{
		{Id: "pool-1", Phase: apiv1.SharedPool_AVAILABLE},
	}
	srv := serveListSharedPools(pools)
	defer srv.Close()

	pm := newTestPoolManager(srv.URL)
	// Pre-populate running with a fake cancel so we test the "already exists" path.
	called := false
	pm.running["pool-1"] = &scalerHandle{cancel: func() { called = true }}

	pm.discover(context.Background())

	// pool-1 already existed — should not be restarted.
	if called {
		t.Error("existing pool should not have been cancelled")
	}
	if len(pm.running) != 1 {
		t.Errorf("expected 1 running, got %d", len(pm.running))
	}
}

func TestDiscover_AddPool_StartsScaler(t *testing.T) {
	pools := []*apiv1.SharedPool{
		{Id: "pool-new", Phase: apiv1.SharedPool_AVAILABLE},
	}
	srv := serveListSharedPools(pools)
	defer srv.Close()

	// Build a real App with mock config so RunOneClusterScaler works.
	a, scalerRefs := newTestAppWithMockConfig(t)

	pm := &poolManager{
		app:            a,
		scalerRefs:     scalerRefs,
		scalerUniqName: "scaler-1",
		executorName:   "ScalerInterfaceV1ExecutorSyncLoop",
		cpClient:       common.NewControlPlaneClient(srv.URL, 5*time.Second),
		ecfg:           &engineConfig{},
		running:        make(map[string]*scalerHandle),
	}

	pm.discover(context.Background())

	if _, ok := pm.running["pool-new"]; !ok {
		t.Fatal("expected pool-new to be started")
	}
	if len(pm.running) != 1 {
		t.Errorf("expected 1 running, got %d", len(pm.running))
	}

	// Cleanup: cancel the scaler goroutine.
	pm.stopAll()
}

// newTestAppWithMockConfig creates an App using mock input/output/scaler types
// and returns it together with the scalerRefs from LoadBaseConfig.
func newTestAppWithMockConfig(t *testing.T) (*app.App, map[string]app.ScalerRefInfo) {
	t.Helper()
	a, err := app.New(nil)
	if err != nil {
		t.Fatalf("app.New: %v", err)
	}
	cfg := &app.AutoscaleConfig{
		Inputs: []app.InputConfig{{
			UniqName:         "input-1",
			Type:             "InputProviderMockV1",
			InterfaceVersion: "InputInterfaceV1",
			Config:           map[string]any{"step": 1},
		}},
		Outputs: []app.OutputConfig{{
			UniqName:         "output-1",
			Type:             "OutputMockEmpty",
			InterfaceVersion: "OutputInterfaceV1",
		}},
		Scalers: []app.ScalerConfig{{
			UniqName:         "scaler-1",
			Type:             "ClusterScalerPhasedComponentScalerV1",
			InterfaceVersion: "ClusterScalerInterfaceV1",
			InputUniqName:    "input-1",
			OutputUniqName:   "output-1",
			DefaultConfig:    map[string]any{"componentName": "tidb"},
		}},
	}
	refs, err := a.LoadBaseConfig(cfg)
	if err != nil {
		t.Fatalf("LoadBaseConfig: %v", err)
	}
	return a, refs
}

func TestDiscover_RemovePool(t *testing.T) {
	// Return empty pool list — previously running pool should be stopped.
	srv := serveListSharedPools(nil)
	defer srv.Close()

	pm := newTestPoolManager(srv.URL)
	cancelled := false
	pm.running["pool-old"] = &scalerHandle{cancel: func() { cancelled = true }}

	pm.discover(context.Background())

	if !cancelled {
		t.Error("expected cancel to be called for removed pool")
	}
	if len(pm.running) != 0 {
		t.Errorf("expected 0 running after removal, got %d", len(pm.running))
	}
}

func TestDiscover_DeletingPoolSkipped(t *testing.T) {
	pools := []*apiv1.SharedPool{
		{Id: "pool-active", Phase: apiv1.SharedPool_AVAILABLE},
		{Id: "pool-deleting", Phase: apiv1.SharedPool_DELETING},
	}
	srv := serveListSharedPools(pools)
	defer srv.Close()

	pm := newTestPoolManager(srv.URL)
	// Pre-add the active pool so we don't need a real app.
	pm.running["pool-active"] = &scalerHandle{cancel: func() {}}

	pm.discover(context.Background())

	if _, ok := pm.running["pool-deleting"]; ok {
		t.Error("DELETING pool should not be started")
	}
	if _, ok := pm.running["pool-active"]; !ok {
		t.Error("ACTIVE pool should still be running")
	}
}

func TestDiscover_ListError_NoEffect(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		_, _ = w.Write([]byte("internal error"))
	}))
	defer srv.Close()

	pm := newTestPoolManager(srv.URL)
	pm.running["existing"] = &scalerHandle{cancel: func() {}}

	pm.discover(context.Background())

	// On error, running should not be modified.
	if len(pm.running) != 1 {
		t.Errorf("expected running unchanged (1), got %d", len(pm.running))
	}
}

func TestDiscover_Idempotent(t *testing.T) {
	pools := []*apiv1.SharedPool{
		{Id: "pool-1", Phase: apiv1.SharedPool_AVAILABLE},
	}
	srv := serveListSharedPools(pools)
	defer srv.Close()

	pm := newTestPoolManager(srv.URL)
	callCount := 0
	pm.running["pool-1"] = &scalerHandle{cancel: func() { callCount++ }}

	// Run discover twice — pool-1 already exists both times.
	pm.discover(context.Background())
	pm.discover(context.Background())

	if callCount != 0 {
		t.Errorf("existing pool should never be cancelled, but cancel called %d times", callCount)
	}
	if len(pm.running) != 1 {
		t.Errorf("expected 1 running, got %d", len(pm.running))
	}
}

func TestStopAll_CancelsAll(t *testing.T) {
	pm := newTestPoolManager("http://unused")
	cancelled := map[string]bool{}
	pm.running["a"] = &scalerHandle{cancel: func() { cancelled["a"] = true }}
	pm.running["b"] = &scalerHandle{cancel: func() { cancelled["b"] = true }}

	pm.stopAll()

	if !cancelled["a"] || !cancelled["b"] {
		t.Errorf("expected all cancels called, got %v", cancelled)
	}
	if pm.running != nil {
		t.Error("expected running to be nil after stopAll")
	}
}

func TestWatchScaler_DoesNotDeleteReplacedHandle(t *testing.T) {
	pm := newTestPoolManager("http://unused")

	oldDone := make(chan struct{})
	oldHandle := &scalerHandle{
		cancel: func() {},
		done:   oldDone,
	}
	newDone := make(chan struct{})
	newHandle := &scalerHandle{
		cancel: func() {},
		done:   newDone,
	}

	pm.mu.Lock()
	pm.running["pool-1"] = oldHandle
	pm.mu.Unlock()

	go pm.watchScaler("pool-1", oldHandle)

	pm.mu.Lock()
	pm.running["pool-1"] = newHandle
	autoscaleTiKVActiveScalers.WithLabelValues("pool-1").Set(1)
	pm.mu.Unlock()

	close(oldDone)

	time.Sleep(50 * time.Millisecond)

	pm.mu.Lock()
	current := pm.running["pool-1"]
	pm.mu.Unlock()

	if current != newHandle {
		t.Error("watchScaler should not have deleted the newer handle")
	}

	val := getGaugeValue(autoscaleTiKVActiveScalers, "pool-1")
	if val != 1 {
		t.Errorf("expected gauge=1 for pool-1, got %v", val)
	}
}

func TestWatchScaler_DeletesOwnHandle(t *testing.T) {
	pm := newTestPoolManager("http://unused")

	done := make(chan struct{})
	handle := &scalerHandle{
		cancel: func() {},
		done:   done,
	}

	pm.mu.Lock()
	pm.running["pool-1"] = handle
	autoscaleTiKVActiveScalers.WithLabelValues("pool-1").Set(1)
	pm.mu.Unlock()

	go pm.watchScaler("pool-1", handle)

	close(done)

	time.Sleep(50 * time.Millisecond)

	pm.mu.Lock()
	_, exists := pm.running["pool-1"]
	pm.mu.Unlock()

	if exists {
		t.Error("expected pool-1 to be removed after its own handle exited")
	}
}

func TestDiscoveryLoop_StopsOnContext(t *testing.T) {
	srv := serveListSharedPools(nil)
	defer srv.Close()

	pm := newTestPoolManager(srv.URL)
	ctx, cancel := context.WithCancel(context.Background())
	cancel() // cancel immediately

	done := make(chan struct{})
	go func() {
		pm.discoveryLoop(ctx)
		close(done)
	}()

	select {
	case <-done:
		// ok
	case <-time.After(2 * time.Second):
		t.Fatal("discoveryLoop did not exit after context cancel")
	}
}

func TestDiscover_ReResolvesCloudProvider(t *testing.T) {
	pools := []*apiv1.SharedPool{
		{
			Id:       "pool-1",
			Phase:    apiv1.SharedPool_AVAILABLE,
			Provider: &apiv1.Provider{Type: apiv1.ProviderType_AWS},
		},
	}
	srv := serveListSharedPools(pools)
	defer srv.Close()

	pm := newTestPoolManager(srv.URL)
	pm.ecfg.Autoscale.Inputs = []app.InputConfig{
		{Config: map[string]any{}},
	}
	// Pre-add pool-1 so discover doesn't try to start a scaler.
	pm.running["pool-1"] = &scalerHandle{cancel: func() {}}

	// cloudProvider is empty — discover should resolve it.
	pm.discover(context.Background())

	if pm.ecfg.Engine.CloudProvider != "aws" {
		t.Errorf("expected cloudProvider=aws, got %q", pm.ecfg.Engine.CloudProvider)
	}
	if cp, ok := pm.ecfg.Autoscale.Inputs[0].Config["cloudProvider"]; !ok || cp != "aws" {
		t.Errorf("expected input cloudProvider=aws, got %v", cp)
	}

	// Second call: switch server to return GCP pools.
	srv.Close()
	gcpPools := []*apiv1.SharedPool{
		{
			Id:       "pool-1",
			Phase:    apiv1.SharedPool_AVAILABLE,
			Provider: &apiv1.Provider{Type: apiv1.ProviderType_GCP},
		},
	}
	srv2 := serveListSharedPools(gcpPools)
	defer srv2.Close()
	pm.cpClient = common.NewControlPlaneClient(srv2.URL, 5*time.Second)

	// cloudProvider already set to "aws" — should NOT be overwritten even though server returns GCP.
	pm.discover(context.Background())
	if pm.ecfg.Engine.CloudProvider != "aws" {
		t.Errorf("expected cloudProvider to stay aws, got %q", pm.ecfg.Engine.CloudProvider)
	}
	if cp := pm.ecfg.Autoscale.Inputs[0].Config["cloudProvider"]; cp != "aws" {
		t.Errorf("expected input cloudProvider to stay aws, got %v", cp)
	}
}

// --- resolveCloudProviderFromPools tests ---

func TestResolveCloudProviderFromPools_Success(t *testing.T) {
	pools := []*apiv1.SharedPool{
		{Provider: &apiv1.Provider{Type: apiv1.ProviderType_GCP}},
	}
	srv := serveListSharedPools(pools)
	defer srv.Close()

	ecfg := &engineConfig{
		Autoscale: app.AutoscaleConfig{
			Inputs: []app.InputConfig{
				{Config: map[string]any{"controlPlaneBaseURL": srv.URL}},
			},
		},
	}

	err := resolveCloudProviderFromPools(ecfg)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if ecfg.Engine.CloudProvider != "gcp" {
		t.Errorf("expected cloudProvider=gcp, got %q", ecfg.Engine.CloudProvider)
	}
	if cp := ecfg.Autoscale.Inputs[0].Config["cloudProvider"]; cp != "gcp" {
		t.Errorf("expected input cloudProvider=gcp, got %v", cp)
	}
}

func TestResolveCloudProviderFromPools_NoPools(t *testing.T) {
	srv := serveListSharedPools(nil)
	defer srv.Close()

	ecfg := &engineConfig{
		Autoscale: app.AutoscaleConfig{
			Inputs: []app.InputConfig{
				{Config: map[string]any{"controlPlaneBaseURL": srv.URL}},
			},
		},
	}

	err := resolveCloudProviderFromPools(ecfg)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	// No pools — cloudProvider stays empty.
	if ecfg.Engine.CloudProvider != "" {
		t.Errorf("expected empty cloudProvider, got %q", ecfg.Engine.CloudProvider)
	}
}

func TestResolveCloudProviderFromPools_NoBaseURL(t *testing.T) {
	ecfg := &engineConfig{}
	err := resolveCloudProviderFromPools(ecfg)
	if err == nil {
		t.Fatal("expected error when no controlPlaneBaseURL")
	}
}

func TestResolveCloudProviderFromPools_APIError(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	}))
	defer srv.Close()

	ecfg := &engineConfig{
		Autoscale: app.AutoscaleConfig{
			Inputs: []app.InputConfig{
				{Config: map[string]any{"controlPlaneBaseURL": srv.URL}},
			},
		},
	}

	err := resolveCloudProviderFromPools(ecfg)
	if err == nil {
		t.Fatal("expected error on API failure")
	}
}

func TestResolveCloudProviderFromPools_DoesNotOverwrite(t *testing.T) {
	pools := []*apiv1.SharedPool{
		{Provider: &apiv1.Provider{Type: apiv1.ProviderType_GCP}},
	}
	srv := serveListSharedPools(pools)
	defer srv.Close()

	ecfg := &engineConfig{
		Engine: engineSpecific{CloudProvider: "aws"},
		Autoscale: app.AutoscaleConfig{
			Inputs: []app.InputConfig{
				{Config: map[string]any{
					"controlPlaneBaseURL": srv.URL,
					"cloudProvider":       "aws",
				}},
			},
		},
	}

	err := resolveCloudProviderFromPools(ecfg)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	// Pre-existing values should not be overwritten.
	if ecfg.Engine.CloudProvider != "aws" {
		t.Errorf("expected cloudProvider=aws, got %q", ecfg.Engine.CloudProvider)
	}
	if cp := ecfg.Autoscale.Inputs[0].Config["cloudProvider"]; cp != "aws" {
		t.Errorf("expected input cloudProvider=aws, got %v", cp)
	}
}

// --- injectCloudProvider tests ---

func TestInjectCloudProvider(t *testing.T) {
	ecfg := &engineConfig{
		Autoscale: app.AutoscaleConfig{
			Inputs: []app.InputConfig{
				{Config: nil},
				{Config: map[string]any{"cloudProvider": "existing"}},
				{Config: map[string]any{}},
			},
		},
	}

	injectCloudProvider(ecfg, "aws")

	if ecfg.Engine.CloudProvider != "aws" {
		t.Errorf("expected Engine.CloudProvider=aws, got %q", ecfg.Engine.CloudProvider)
	}
	// Input 0: had nil config, should be created with cloudProvider.
	if ecfg.Autoscale.Inputs[0].Config["cloudProvider"] != "aws" {
		t.Errorf("input[0]: expected cloudProvider=aws, got %v", ecfg.Autoscale.Inputs[0].Config["cloudProvider"])
	}
	// Input 1: had pre-existing cloudProvider, should NOT be overwritten.
	if ecfg.Autoscale.Inputs[1].Config["cloudProvider"] != "existing" {
		t.Errorf("input[1]: expected cloudProvider=existing, got %v", ecfg.Autoscale.Inputs[1].Config["cloudProvider"])
	}
	// Input 2: had empty config, should get cloudProvider.
	if ecfg.Autoscale.Inputs[2].Config["cloudProvider"] != "aws" {
		t.Errorf("input[2]: expected cloudProvider=aws, got %v", ecfg.Autoscale.Inputs[2].Config["cloudProvider"])
	}
}

func TestInjectCloudProvider_DoesNotOverwriteEngine(t *testing.T) {
	ecfg := &engineConfig{
		Engine: engineSpecific{CloudProvider: "gcp"},
	}
	injectCloudProvider(ecfg, "aws")
	if ecfg.Engine.CloudProvider != "gcp" {
		t.Errorf("expected Engine.CloudProvider=gcp (not overwritten), got %q", ecfg.Engine.CloudProvider)
	}
}

func getGaugeValue(gv *prometheus.GaugeVec, labels ...string) float64 {
	g, err := gv.GetMetricWithLabelValues(labels...)
	if err != nil {
		return 0
	}
	var m dto.Metric
	if err := g.Write(&m); err != nil {
		return 0
	}
	return m.GetGauge().GetValue()
}