package entity

import (
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	schemav1 "github.com/tidbcloud/cluster-service-ng/pkg/storage/etcd/schema/v1"
	mysqlschema "github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema"
)

func TestComputeTiDBConfig_PremiumReservedSettings(t *testing.T) {
	t.Run("tikv worker url is still configured", func(t *testing.T) {
		sharedPool := buildSharedPool(apiv1.CapacityMode_TOPOLOGY_DRIVEN, []*Component{
			{ComponentType: apiv1.ComponentType_PD, Id: "sp-1"},
			{ComponentType: apiv1.ComponentType_TIKV, Id: "sp-1"},
			{ComponentType: apiv1.ComponentType_TIKV_WORKER, Id: "sp-1", NodeQuota: &apiv1.NodeQuota{NodeCount: 1}},
		})
		tidb := buildTiDBComponent(apiv1.ComponentType_TIDB)

		err := ComputeTiDBConfig(sharedPool, tidb)
		require.NoError(t, err)

		conf := parseTiDBToml(t, tidb.GetConfiguration().GetConfigurationFileContent())
		assert.Equal(t, "https://db-sp-1-default-tikv-worker-peer.tidbsp-1.svc:19000", conf.Get("tikv-worker-url"))
	})
}

func TestComputeTiProxyConfig_RoutingRule(t *testing.T) {
	t.Run("shared tiproxy uses proxy cidr routing", func(t *testing.T) {
		sharedPool := buildSharedPool(apiv1.CapacityMode_RCU_DRIVEN, nil)
		sharedPool.DeployType = mysqlschema.SharedPoolDeployTypeSharedTiProxy
		tiproxy := buildTiProxyComponent()
		tiproxy.GetConfiguration().ConfigurationFileContent = "[balance]\nlabel-name = \"keyspace\"\n[log]\nlevel = \"info\"\n"
		tiproxy.Labels = map[string]string{"keyspace": "stale", "custom": "keep"}

		err := ComputeTiProxyConfig(sharedPool, tiproxy)
		require.NoError(t, err)

		conf := parseToml(t, tiproxy.GetConfiguration().GetConfigurationFileContent())
		assert.Equal(t, "proxy_cidr", conf.Get("balance.routing-rule"))
		assert.Nil(t, conf.Get("balance.label-name"))
		assert.Equal(t, "info", conf.Get("log.level"))
		assert.Equal(t, "stale", tiproxy.GetLabels()["keyspace"])
		assert.Equal(t, "keep", tiproxy.GetLabels()["custom"])
	})

	t.Run("dedicated tiproxy keeps routing rule unset and uses keyspace label", func(t *testing.T) {
		sharedPool := buildSharedPool(apiv1.CapacityMode_RCU_DRIVEN, nil)
		tiproxy := buildTiProxyComponent()

		err := ComputeTiProxyConfig(sharedPool, tiproxy)
		require.NoError(t, err)

		conf := parseToml(t, tiproxy.GetConfiguration().GetConfigurationFileContent())
		assert.Nil(t, conf.Get("balance.routing-rule"))
		assert.Equal(t, "keyspace", conf.Get("balance.label-name"))
		assert.Equal(t, "info", conf.Get("log.level"))
	})
}

func buildSharedPool(capacityMode apiv1.CapacityMode, components []*Component) *SharedPool {
	return &SharedPool{SharedPool: apiv1.SharedPool{
		Id:            "sp-1",
		CapacityMode:  capacityMode,
		DataplaneName: "dp-1",
		ObjectStorage: &apiv1.SharedPool_ObjectStorage{
			Region: "us-west-2",
			Bucket: "bucket-1",
			Prefix: "prefix-1",
		},
		Provider: &apiv1.Provider{
			Type:      apiv1.ProviderType_AWS,
			Region:    "us-west-2",
			NodeZones: []string{"us-west-2a"},
		},
		ComponentGroups: components,
	}}
}

func buildTiDBComponent(componentType apiv1.ComponentType) *Component {
	return &Component{
		Id:            "sp-1",
		ComponentType: componentType,
		Configuration: &ComponentConfiguration{
			ConfigurationFileContent: "log-level = \"info\"\n",
		},
	}
}

func buildTiProxyComponent() *Component {
	return &Component{
		Id:            "sp-1",
		ComponentType: apiv1.ComponentType_TIPROXY,
		Configuration: &ComponentConfiguration{
			ConfigurationFileContent: "[log]\nlevel = \"info\"\n",
		},
	}
}

func buildTiKVComponent() *Component {
	return &Component{
		Id:            "sp-1",
		ComponentType: apiv1.ComponentType_TIKV,
		Storages: []*Storage{
			{
				Name:      apiv1.Component_Storage_DATA,
				SizeGi:    200,
				Type:      "gp3",
				MountPath: "/data",
			},
			{
				Name:      apiv1.Component_Storage_RAFTLOG,
				SizeGi:    100,
				Type:      "gp3",
				MountPath: "/raftlog",
			},
		},
		Configuration: &ComponentConfiguration{
			ConfigurationFileContent: "[log]\nlevel = \"info\"\n",
		},
	}
}

func buildTiKVWorkerComponent(componentType apiv1.ComponentType) *Component {
	return &Component{
		Id:            "sp-1",
		ComponentType: componentType,
		Configuration: &ComponentConfiguration{
			ConfigurationFileContent: "[log]\nlevel = \"info\"\n",
		},
	}
}

func buildTiFlashComputeComponent() *Component {
	return &Component{
		Id:            "sp-1",
		ComponentType: apiv1.ComponentType_TIFLASH_COMPUTE,
		Storages: []*Storage{
			{
				Name:      apiv1.Component_Storage_DATA,
				SizeGi:    200,
				Type:      "gp3",
				MountPath: "/data0",
			},
		},
		Configuration: &ComponentConfiguration{
			ConfigurationFileContent: "[logger]\ncount = 1\n",
		},
	}
}

func parseTiDBToml(t *testing.T, config string) *schemav1.TomlConfig {
	t.Helper()
	return parseToml(t, config)
}

func parseToml(t *testing.T, config string) *schemav1.TomlConfig {
	t.Helper()
	conf := schemav1.NewTomlConfig(map[string]any{})
	require.NoError(t, conf.UnmarshalTOML([]byte(config)))
	return conf
}

func TestComputeTiDBConfig_ColumnarSharedTiProxy(t *testing.T) {
	sharedPool := buildSharedPool(apiv1.CapacityMode_RCU_DRIVEN, []*Component{
		{ComponentType: apiv1.ComponentType_PD, Id: "sp-1"},
		{ComponentType: apiv1.ComponentType_TIKV_WORKER, Id: "sp-1", NodeQuota: &apiv1.NodeQuota{NodeCount: 1}},
		{ComponentType: apiv1.ComponentType_TIFLASH_COMPUTE, Id: "sp-1"},
	})
	sharedPool.DeployType = mysqlschema.SharedPoolDeployTypeSharedTiProxy
	tidb := buildTiDBComponent(apiv1.ComponentType_TIDB)

	err := ComputeTiDBConfig(sharedPool, tidb)
	require.NoError(t, err)

	conf := parseTiDBToml(t, tidb.GetConfiguration().GetConfigurationFileContent())
	assert.Equal(t, true, conf.Get("disaggregated-tiflash"))
	assert.Equal(t, "columnar", conf.Get("cse.columnar-store-type"))
	assert.Equal(t, "https://db-sp-1-default-tikv-worker-peer.tidbsp-1.svc:19000", conf.Get("tikv-worker-url"))
}

func TestComputeTiKVConfig_ColumnarSharedTiProxy(t *testing.T) {
	sharedPool := buildSharedPool(apiv1.CapacityMode_RCU_DRIVEN, []*Component{
		{ComponentType: apiv1.ComponentType_SCHEMA_MANAGER, Id: "sp-1", NodeQuota: &apiv1.NodeQuota{NodeCount: 1}},
	})
	sharedPool.DeployType = mysqlschema.SharedPoolDeployTypeSharedTiProxy
	tikv := buildTiKVComponent()

	err := ComputeTiKVConfig(sharedPool, tikv, "")
	require.NoError(t, err)

	conf := parseToml(t, tikv.GetConfiguration().GetConfigurationFileContent())
	assert.Equal(t, true, conf.Get("kvengine.build-columnar"))
	assert.Equal(t, true, conf.Get("kvengine.build-fts-index"))
}

func TestComputeSchemaManagerConfig_ColumnarSharedTiProxy(t *testing.T) {
	sharedPool := buildSharedPool(apiv1.CapacityMode_RCU_DRIVEN, nil)
	sharedPool.DeployType = mysqlschema.SharedPoolDeployTypeSharedTiProxy

	comp := buildTiKVWorkerComponent(apiv1.ComponentType_SCHEMA_MANAGER)
	err := ComputeSchemaManagerConfig(sharedPool, comp)
	require.NoError(t, err)
	conf := parseToml(t, comp.GetConfiguration().GetConfigurationFileContent())
	assert.Equal(t, true, conf.Get("schema-manager.enabled"))
	assert.Equal(t, "/var/lib/tikv-worker/schemas", conf.Get("schema-manager.dir"))
	assert.Equal(t, "10s", conf.Get("schema-manager.keyspace-refresh-interval"))
}

func TestComputeTiflashComputeConfig_ColumnarSharedTiProxy(t *testing.T) {
	sharedPool := buildSharedPool(apiv1.CapacityMode_RCU_DRIVEN, nil)
	sharedPool.DeployType = mysqlschema.SharedPoolDeployTypeSharedTiProxy
	comp := buildTiFlashComputeComponent()

	err := ComputeTiflashComputeConfig(sharedPool, comp, "")
	require.NoError(t, err)

	conf := parseToml(t, comp.GetConfiguration().GetConfigurationFileContent())
	assert.Equal(t, true, conf.Get("flash.use_columnar"))
	assert.Equal(t, "http://s3.us-west-2.amazonaws.com", conf.Get("storage.s3.endpoint"))
	assert.Equal(t, "bucket-1", conf.Get("storage.s3.bucket"))
	assert.Equal(t, filepath.Join("prefix-1", "tiflash-root"), conf.Get("storage.s3.root"))
}

func TestComputeColumnarConfigs_TiflashMode(t *testing.T) {
	tests := []struct {
		name          string
		deployType    string
		tiflashMode   string
		wantColumnar  bool
		wantProxyCIDR bool
	}{
		{
			name:          "marker pool enables columnar without shared tiproxy routing",
			tiflashMode:   mysqlschema.SharedPoolTiflashModeColumnar,
			wantColumnar:  true,
			wantProxyCIDR: false,
		},
		{
			name:         "legacy pool keeps non columnar config",
			wantColumnar: false,
		},
		{
			name:          "legacy shared tiproxy remains columnar",
			deployType:    mysqlschema.SharedPoolDeployTypeSharedTiProxy,
			wantColumnar:  true,
			wantProxyCIDR: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			sharedPool := buildSharedPool(apiv1.CapacityMode_RCU_DRIVEN, []*Component{
				{ComponentType: apiv1.ComponentType_PD, Id: "sp-1"},
				{ComponentType: apiv1.ComponentType_TIKV_WORKER, Id: "sp-1", NodeQuota: &apiv1.NodeQuota{NodeCount: 1}},
				{ComponentType: apiv1.ComponentType_TIFLASH_COMPUTE, Id: "sp-1"},
				{ComponentType: apiv1.ComponentType_SCHEMA_MANAGER, Id: "sp-1", NodeQuota: &apiv1.NodeQuota{NodeCount: 1}},
			})
			sharedPool.DeployType = tt.deployType
			sharedPool.TiflashMode = tt.tiflashMode

			tidb := buildTiDBComponent(apiv1.ComponentType_TIDB)
			require.NoError(t, ComputeTiDBConfig(sharedPool, tidb))
			tidbConf := parseTiDBToml(t, tidb.GetConfiguration().GetConfigurationFileContent())
			if tt.wantColumnar {
				assert.Equal(t, "columnar", tidbConf.Get("cse.columnar-store-type"))
			} else {
				assert.Nil(t, tidbConf.Get("cse.columnar-store-type"))
			}

			tikv := buildTiKVComponent()
			require.NoError(t, ComputeTiKVConfig(sharedPool, tikv, ""))
			tikvConf := parseToml(t, tikv.GetConfiguration().GetConfigurationFileContent())
			if tt.wantColumnar {
				assert.Equal(t, true, tikvConf.Get("kvengine.build-columnar"))
				assert.Equal(t, true, tikvConf.Get("kvengine.build-fts-index"))
			} else {
				assert.Nil(t, tikvConf.Get("kvengine.build-columnar"))
				assert.Nil(t, tikvConf.Get("kvengine.build-fts-index"))
			}

			tiflash := buildTiFlashComputeComponent()
			require.NoError(t, ComputeTiflashComputeConfig(sharedPool, tiflash, ""))
			tiflashConf := parseToml(t, tiflash.GetConfiguration().GetConfigurationFileContent())
			if tt.wantColumnar {
				assert.Equal(t, true, tiflashConf.Get("flash.use_columnar"))
			} else {
				assert.Nil(t, tiflashConf.Get("flash.use_columnar"))
			}

			tiproxy := buildTiProxyComponent()
			require.NoError(t, ComputeTiProxyConfig(sharedPool, tiproxy))
			tiproxyConf := parseToml(t, tiproxy.GetConfiguration().GetConfigurationFileContent())
			if tt.wantProxyCIDR {
				assert.Equal(t, "proxy_cidr", tiproxyConf.Get("balance.routing-rule"))
			} else {
				assert.Nil(t, tiproxyConf.Get("balance.routing-rule"))
				assert.Equal(t, "keyspace", tiproxyConf.Get("balance.label-name"))
			}
		})
	}
}