package cluster

import (
	"context"
	"testing"

	. "github.com/onsi/gomega"
	"github.com/tidbcloud/account/idl/pbgen/proto/account"
	"go.uber.org/mock/gomock"

	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	openapicommonv1beta1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_cloud_open_api/common/v1beta1"
	openapidedicatedv1beta1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_cloud_open_api/dedicated/v1beta1"
	tidbmgmtservicev1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/app/internalapi/cluster/mock_dedicatedclustersvc"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/domain/adapter/account/mock_account"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/domain/entity"
)

// fakeRegionServer implements openapidedicatedv1beta1.RegionServiceServer for testing.
type fakeRegionServer struct {
	openapidedicatedv1beta1.UnimplementedRegionServiceServer
}

func (f *fakeRegionServer) GetNodeSpec(
	_ context.Context,
	req *openapidedicatedv1beta1.GetNodeSpecRequest,
) (*openapidedicatedv1beta1.NodeSpec, error) {
	return &openapidedicatedv1beta1.NodeSpec{
		NodeSpecKey: req.GetNodeSpecKey(),
		StorageTypes: []openapidedicatedv1beta1.Cluster_StorageNodeSetting_StorageType{
			openapidedicatedv1beta1.Cluster_StorageNodeSetting_Basic,
			openapidedicatedv1beta1.Cluster_StorageNodeSetting_Standard,
		},
	}, nil
}

// fakeCustomStatusSvc implements customstatus.Service for testing.
type fakeCustomStatusSvc struct{}

func (f *fakeCustomStatusSvc) SetClusterCustomStatus(
	context.Context,
	*entity.ClusterCustomStatus,
) (a, b *entity.ClusterCustomStatus, _ openapicommonv1beta1.ClusterState, _ error) {
	return &entity.ClusterCustomStatus{}, &entity.ClusterCustomStatus{}, 0, nil
}

func (f *fakeCustomStatusSvc) GetClusterCustomStatusByID(context.Context, uint64, uint64) (*entity.ClusterCustomStatus, error) {
	return &entity.ClusterCustomStatus{}, nil
}

func (f *fakeCustomStatusSvc) GetClusterCustomStatusByName(
	context.Context,
	uint64,
	string,
	tidbmgmtservicev1.CustomStatusType,
) (*entity.ClusterCustomStatus, error) {
	return &entity.ClusterCustomStatus{}, nil
}

func (f *fakeCustomStatusSvc) ListClusterCustomStatus(
	_ context.Context,
	clusterIds []uint64,
	_ []tidbmgmtservicev1.CustomStatusType,
	_ bool,
) (map[uint64][]*entity.ClusterCustomStatus, error) {
	result := make(map[uint64][]*entity.ClusterCustomStatus)
	for _, id := range clusterIds {
		result[id] = []*entity.ClusterCustomStatus{}
	}
	return result, nil
}

func (f *fakeCustomStatusSvc) DeleteClusterCustomStatus(context.Context, uint64, uint64) error {
	return nil
}

func (f *fakeCustomStatusSvc) ClearClusterCustomStatusRecordByClusterId(context.Context, uint64) error {
	return nil
}

// baseRestoreRequest returns a valid RestoreDedicatedClusterRequest for testing.
func baseRestoreRequest() *tidbmgmtservicev1.RestoreDedicatedClusterRequest {
	return &tidbmgmtservicev1.RestoreDedicatedClusterRequest{
		CreatedBy:      "1",
		OrganizationId: 1,
		ProjectId:      1,
		CloudProvider:  "aws",
		Region:         "us-east-1",
		ClusterName:    "test-cluster",
		Port:           4000,
		Version:        "v8.5.0",
		TidbComp: &tidbmgmtservicev1.RestoreDedicatedClusterRequest_ComputeComp{
			SpuName: "8C16G",
			Count:   1,
		},
		TikvComp: &tidbmgmtservicev1.RestoreDedicatedClusterRequest_StorageComp{
			SpuName: "8C32G",
			Count:   3,
			SizeGi:  500,
		},
	}
}

// setupRestoreTest creates a serverImpl with mocked dependencies for RestoreDedicatedCluster tests.
func setupRestoreTest(t *testing.T, ctrl *gomock.Controller) (*serverImpl, *mock_dedicatedclustersvc.MockService, *mock_account.MockClient) {
	t.Helper()

	accountCli := mock_account.NewMockClient(ctrl)
	clusterSvc := mock_dedicatedclustersvc.NewMockService(ctrl)

	server := &serverImpl{
		accountCli:                accountCli,
		dedicatedClusterSvc:       clusterSvc,
		openApiDedicatedRegionApp: &fakeRegionServer{},
		customStatusSvc:           &fakeCustomStatusSvc{},
	}
	return server, clusterSvc, accountCli
}

// TestClusterConfigurationPassthrough verifies that ClusterConfiguration
// from the request is passed through directly to RestoreCluster.
func TestClusterConfigurationPassthrough(t *testing.T) {
	g := NewWithT(t)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	server, clusterSvc, accountCli := setupRestoreTest(t, ctrl)

	accountCli.EXPECT().
		GetUserById(gomock.Any(), uint64(1)).
		Return(&account.User{Id: 1, Email: "test@example.com"}, nil)

	clusterSvc.EXPECT().
		RestoreCluster(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(
			_ context.Context,
			_ *entity.Cluster,
			_, _ uint64,
			_ bool,
			clusterCfg *globaldedicatedv1.ClusterConfiguration,
		) (*entity.Cluster, error) {
			g.Expect(clusterCfg).ShouldNot(BeNil())
			g.Expect(clusterCfg.ComponentConfigurations).Should(HaveLen(2))

			g.Expect(clusterCfg.ComponentConfigurations[0].ComponentType).Should(Equal(globaldedicatedv1.ComponentType_TIDB))
			g.Expect(clusterCfg.ComponentConfigurations[0].OverwriteImageTag).Should(Equal("v1.0.0"))

			g.Expect(clusterCfg.ComponentConfigurations[1].ComponentType).Should(Equal(globaldedicatedv1.ComponentType_TIKV))
			g.Expect(clusterCfg.ComponentConfigurations[1].OverwriteImageTag).Should(Equal("v2.0.0"))

			return &entity.Cluster{ClusterId: 100}, nil
		})

	req := baseRestoreRequest()
	req.ClusterConfiguration = &globaldedicatedv1.ClusterConfiguration{
		ComponentConfigurations: []*globaldedicatedv1.ClusterConfiguration_ComponentConfiguration{
			{ComponentType: globaldedicatedv1.ComponentType_TIDB, OverwriteImageTag: "v1.0.0"},
			{ComponentType: globaldedicatedv1.ComponentType_TIKV, OverwriteImageTag: "v2.0.0"},
		},
	}

	_, err := server.RestoreDedicatedCluster(context.Background(), req)
	g.Expect(err).ShouldNot(HaveOccurred())
}

// TestAdditionalIopsSetOnDataStorage verifies that AdditionalIops is set on DATA
// storage for TiKV and TiFlash when additional_iops > 0.
func TestAdditionalIopsSetOnDataStorage(t *testing.T) {
	g := NewWithT(t)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	server, clusterSvc, accountCli := setupRestoreTest(t, ctrl)

	accountCli.EXPECT().
		GetUserById(gomock.Any(), uint64(1)).
		Return(&account.User{Id: 1, Email: "test@example.com"}, nil)

	clusterSvc.EXPECT().
		RestoreCluster(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(
			_ context.Context,
			restoreCluster *entity.Cluster,
			_, _ uint64,
			_ bool,
			_ *globaldedicatedv1.ClusterConfiguration,
		) (*entity.Cluster, error) {
			gc, ok := restoreCluster.GlobalCluster.(*globaldedicatedv1.Cluster)
			g.Expect(ok).Should(BeTrue())
			comps := gc.ClusterComponents

			// Find TiKV component and check DATA storage has AdditionalIops
			tikvFound := false
			for _, comp := range comps {
				if comp.ComponentType != globaldedicatedv1.ComponentType_TIKV {
					continue
				}
				tikvFound = true
				dataStorage, ok := comp.Storages[globaldedicatedv1.Sku_DATA.String()]
				g.Expect(ok).Should(BeTrue())
				g.Expect(dataStorage.AdditionalIops).ShouldNot(BeNil())
				g.Expect(*dataStorage.AdditionalIops).Should(Equal(uint32(1000)))
			}
			g.Expect(tikvFound).Should(BeTrue())

			// Find TiFlash component and check DATA storage has AdditionalIops
			tiflashFound := false
			for _, comp := range comps {
				if comp.ComponentType != globaldedicatedv1.ComponentType_TIFLASH {
					continue
				}
				tiflashFound = true
				dataStorage, ok := comp.Storages[globaldedicatedv1.Sku_DATA.String()]
				g.Expect(ok).Should(BeTrue())
				g.Expect(dataStorage.AdditionalIops).ShouldNot(BeNil())
				g.Expect(*dataStorage.AdditionalIops).Should(Equal(uint32(500)))
			}
			g.Expect(tiflashFound).Should(BeTrue())

			return &entity.Cluster{ClusterId: 100}, nil
		})

	req := baseRestoreRequest()
	req.TikvComp.AdditionalIops = 1000
	req.TiflashComp = &tidbmgmtservicev1.RestoreDedicatedClusterRequest_StorageComp{
		SpuName:        "8C32G",
		Count:          2,
		SizeGi:         200,
		AdditionalIops: 500,
	}

	_, err := server.RestoreDedicatedCluster(context.Background(), req)
	g.Expect(err).ShouldNot(HaveOccurred())
}

// TestNilClusterConfigurationPassthrough verifies that nil ClusterConfiguration
// in the request results in nil being passed to RestoreCluster.
func TestNilClusterConfigurationPassthrough(t *testing.T) {
	g := NewWithT(t)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	server, clusterSvc, accountCli := setupRestoreTest(t, ctrl)

	accountCli.EXPECT().
		GetUserById(gomock.Any(), uint64(1)).
		Return(&account.User{Id: 1, Email: "test@example.com"}, nil)

	clusterSvc.EXPECT().
		RestoreCluster(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(
			_ context.Context,
			_ *entity.Cluster,
			_, _ uint64,
			_ bool,
			clusterCfg *globaldedicatedv1.ClusterConfiguration,
		) (*entity.Cluster, error) {
			g.Expect(clusterCfg).Should(BeNil())
			return &entity.Cluster{ClusterId: 100}, nil
		})

	req := baseRestoreRequest()
	req.ClusterConfiguration = nil

	_, err := server.RestoreDedicatedCluster(context.Background(), req)
	g.Expect(err).ShouldNot(HaveOccurred())
}

// TestAdditionalIopsZeroNotSet verifies that additional_iops = 0 does not set
// AdditionalIops on the DATA storage.
func TestAdditionalIopsZeroNotSet(t *testing.T) {
	g := NewWithT(t)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	server, clusterSvc, accountCli := setupRestoreTest(t, ctrl)

	accountCli.EXPECT().
		GetUserById(gomock.Any(), uint64(1)).
		Return(&account.User{Id: 1, Email: "test@example.com"}, nil)

	clusterSvc.EXPECT().
		RestoreCluster(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(
			_ context.Context,
			restoreCluster *entity.Cluster,
			_, _ uint64,
			_ bool,
			_ *globaldedicatedv1.ClusterConfiguration,
		) (*entity.Cluster, error) {
			gc, ok := restoreCluster.GlobalCluster.(*globaldedicatedv1.Cluster)
			g.Expect(ok).Should(BeTrue())
			comps := gc.ClusterComponents

			// Verify TiKV DATA storage does NOT have AdditionalIops set
			for _, comp := range comps {
				if comp.ComponentType != globaldedicatedv1.ComponentType_TIKV {
					continue
				}
				dataStorage, ok := comp.Storages[globaldedicatedv1.Sku_DATA.String()]
				g.Expect(ok).Should(BeTrue())
				g.Expect(dataStorage.AdditionalIops).Should(BeNil())
			}

			return &entity.Cluster{ClusterId: 100}, nil
		})

	req := baseRestoreRequest()
	req.TikvComp.AdditionalIops = 0

	_, err := server.RestoreDedicatedCluster(context.Background(), req)
	g.Expect(err).ShouldNot(HaveOccurred())
}