package resourcepool

import (
	"context"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"
	"google.golang.org/grpc/codes"
	"google.golang.org/protobuf/proto"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/resourcepool/mock_resourcepool"
	regionservice "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/service/region"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

type capacityRegionServiceStub struct {
	regionservice.Service

	constraints func(context.Context, *globalngv1.GetByocSharedComponentScaleConstraintsRequest) (*globalngv1.ByocSharedComponentScaleConstraints, bizerr.BizError)
	scale       func(context.Context, *globalngv1.ScaleByocSharedComponentsRequest) (*globalngv1.ScaleByocSharedComponentsResponse, bizerr.BizError)
	listSpecs   func(context.Context, *globalngv1.ListNodeSpecsRequest) (*globalngv1.ListNodeSpecsResponse, bizerr.BizError)
}

func (s capacityRegionServiceStub) GetByocSharedComponentScaleConstraints(
	ctx context.Context,
	req *globalngv1.GetByocSharedComponentScaleConstraintsRequest,
) (*globalngv1.ByocSharedComponentScaleConstraints, bizerr.BizError) {
	return s.constraints(ctx, req)
}

func (s capacityRegionServiceStub) ScaleByocSharedComponents(
	ctx context.Context,
	req *globalngv1.ScaleByocSharedComponentsRequest,
) (*globalngv1.ScaleByocSharedComponentsResponse, bizerr.BizError) {
	return s.scale(ctx, req)
}

func (s capacityRegionServiceStub) ListNodeSpecs(
	ctx context.Context,
	req *globalngv1.ListNodeSpecsRequest,
) (*globalngv1.ListNodeSpecsResponse, bizerr.BizError) {
	return s.listSpecs(ctx, req)
}

func TestResizeResourcePoolCapacityPreviewFirstAndNoWriteTarget(t *testing.T) {
	t.Parallel()

	t.Run("blocked validate-only stops after preview", func(t *testing.T) {
		t.Parallel()

		calls := 0
		svc := &serviceImpl{regionService: capacityRegionServiceStub{
			scale: func(_ context.Context, req *globalngv1.ScaleByocSharedComponentsRequest) (*globalngv1.ScaleByocSharedComponentsResponse, bizerr.BizError) {
				calls++
				require.True(t, req.GetValidateOnly())
				require.Nil(t, req.GetTiflashWrite())
				return &globalngv1.ScaleByocSharedComponentsResponse{
					Preview: &globalngv1.ByocSharedComponentScalePreviewSummary{
						Allowed:         false,
						BlockingReasons: []string{"pool is modifying; regional-secret=do-not-leak"},
					},
				}, nil
			},
		}}

		_, err := svc.ResizeResourcePoolCapacity(t.Context(), resizeRequest(true, nil))
		require.Error(t, err)
		require.Equal(t, bizerr.ErrCode(codes.FailedPrecondition), err.ErrCode())
		require.NotContains(t, err.Error(), "regional-secret")
		require.Equal(t, 1, calls)
	})

	t.Run("real resize validates then submits and records lifecycle", func(t *testing.T) {
		t.Parallel()

		ctrl := gomock.NewController(t)
		repo := mock_resourcepool.NewMockRepo(ctrl)
		order := make([]bool, 0, 2)
		stub := capacityRegionServiceStub{
			scale: func(_ context.Context, req *globalngv1.ScaleByocSharedComponentsRequest) (*globalngv1.ScaleByocSharedComponentsResponse, bizerr.BizError) {
				order = append(order, req.GetValidateOnly())
				require.Nil(t, req.GetTiflashWrite())
				require.Nil(t, req.GetTiflashCompute(), "omitted Compute must remain absent")
				if !req.GetValidateOnly() {
					require.EqualValues(t, 6, req.GetTikv().GetNodeCount())
					require.EqualValues(t, 4, req.GetTikv().GetCpu())
					require.EqualValues(t, 16, req.GetTikv().GetMemoryGi())
					require.Len(t, req.GetTikv().GetStorages(), 2)
					require.Equal(t, "DATA", req.GetTikv().GetStorages()[0].GetStorageName())
					require.EqualValues(t, 500, req.GetTikv().GetStorages()[0].GetSizeGi())
					require.Equal(t, "RAFTLOG", req.GetTikv().GetStorages()[1].GetStorageName())
					require.EqualValues(t, 100, req.GetTikv().GetStorages()[1].GetSizeGi())
					return &globalngv1.ScaleByocSharedComponentsResponse{}, nil
				}
				require.Empty(t, req.GetTikv().GetStorages(), "public preview must not set raftlog")
				return allowedScalePreview(), nil
			},
		}
		svc := &serviceImpl{regionService: stub, resourcePoolRepo: repo}

		response, err := svc.ResizeResourcePoolCapacity(t.Context(), resizeRequest(false, nil))
		require.NoError(t, err)
		require.Equal(t, []bool{true, false}, order)
		require.EqualValues(t, 3, response.GetCurrentCapacity().GetTikv().GetNodeCount())
		require.EqualValues(t, 6, response.GetTargetCapacity().GetTikv().GetNodeCount())
		require.Len(t, response.GetCurrentCapacity().GetTikv().GetStorages(), 1)
		require.Len(t, response.GetTargetCapacity().GetTikv().GetStorages(), 1)
	})

	t.Run("terminal workflow failure is not overwritten after submission", func(t *testing.T) {
		t.Parallel()

		ctrl := gomock.NewController(t)
		repo := mock_resourcepool.NewMockRepo(ctrl)
		repo.EXPECT().UpdateLifecycle(
			gomock.Any(), uint64(41), globalngv1.ResourcePool_FAILED, "",
		).Return(nil)
		calls := 0
		svc := &serviceImpl{
			resourcePoolRepo: repo,
			regionService: capacityRegionServiceStub{
				scale: func(ctx context.Context, req *globalngv1.ScaleByocSharedComponentsRequest) (*globalngv1.ScaleByocSharedComponentsResponse, bizerr.BizError) {
					calls++
					if req.GetValidateOnly() {
						return allowedScalePreview(), nil
					}
					// Simulate the workflow reaching its terminal-failure hook before
					// the regional submission call returns to the public adapter.
					require.NoError(t, repo.UpdateLifecycle(ctx, 41, globalngv1.ResourcePool_FAILED, ""))
					return &globalngv1.ScaleByocSharedComponentsResponse{}, nil
				},
			},
		}

		_, err := svc.ResizeResourcePoolCapacity(t.Context(), resizeRequest(false, nil))
		require.NoError(t, err)
		require.Equal(t, 2, calls)
	})

	t.Run("explicit Compute zero remains present in preview request", func(t *testing.T) {
		t.Parallel()

		calls := 0
		svc := &serviceImpl{regionService: capacityRegionServiceStub{
			scale: func(_ context.Context, req *globalngv1.ScaleByocSharedComponentsRequest) (*globalngv1.ScaleByocSharedComponentsResponse, bizerr.BizError) {
				calls++
				require.True(t, req.GetValidateOnly())
				require.Nil(t, req.GetTiflashWrite())
				require.NotNil(t, req.GetTiflashCompute())
				require.NotNil(t, req.GetTiflashCompute().NodeCount)
				require.Zero(t, req.GetTiflashCompute().GetNodeCount())
				preview := allowedScalePreview()
				preview.Components = append(preview.Components, &globalngv1.ByocSharedComponentScalePlanComponent{
					ComponentType: globalngv1.ByocSharedComponentType_BYOC_SHARED_COMPONENT_TYPE_TIFLASH_COMPUTE,
					Current:       scaleTarget(8, 32, 2, "DATA", 500),
					Target:        scaleTarget(8, 32, 0, "DATA", 500),
					Allowed:       true,
				})
				return preview, nil
			},
		}}
		zero := int64(0)

		response, err := svc.ResizeResourcePoolCapacity(t.Context(), resizeRequest(true, &globalngv1.ComponentCapacity{
			NodeCount: &zero,
		}))
		require.NoError(t, err)
		require.Equal(t, 1, calls)
		require.NotNil(t, response.GetTargetCapacity().GetTiflash())
		require.Zero(t, response.GetTargetCapacity().GetTiflash().GetNodeCount())
	})
}

func TestBuildPublicScaleRequestRejectsRaftlog(t *testing.T) {
	t.Parallel()

	request := resizeRequest(true, nil)
	request.TargetCapacity.Tikv.Storages = []*globalngv1.Storage{
		{Name: "data", SizeGi: 500},
		{Name: "raftlog", SizeGi: 100},
	}

	scaleRequest, err := buildPublicScaleRequest(request, true)

	require.Nil(t, scaleRequest)
	require.ErrorContains(t, err, "unsupported storage name")
}

func TestGetResourcePoolCapacityConstraintsProjection(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	repo := mock_resourcepool.NewMockRepo(ctrl)
	repo.EXPECT().Load(gomock.Any(), uint64(41), false).Return(&entity.ResourcePool{
		ResourcePool: &globalngv1.ResourcePool{
			Id:          41,
			OwnerType:   globalngv1.ResourcePool_ORGANIZATION,
			OwnerId:     501,
			ProviderId:  globalcommonv1.CloudProvider_AWS,
			RegionId:    globalcommonv1.Region_AWS_US_WEST_2,
			ServicePlan: globalcommonv1.ServicePlan_BYOC,
		},
	}, nil)
	stub := capacityRegionServiceStub{
		constraints: func(_ context.Context, _ *globalngv1.GetByocSharedComponentScaleConstraintsRequest) (*globalngv1.ByocSharedComponentScaleConstraints, bizerr.BizError) {
			return &globalngv1.ByocSharedComponentScaleConstraints{
				Preview: &globalngv1.ByocSharedComponentScalePreviewSummary{
					Allowed:         false,
					BlockingReasons: []string{"byoc vcpu pool is modifying"},
				},
				Components: []*globalngv1.ByocSharedComponentScaleConstraint{constraintComponent(
					globalngv1.ByocSharedComponentType_BYOC_SHARED_COMPONENT_TYPE_TIKV,
					scaleTarget(4, 16, 3, "DATA", 500, "RAFTLOG", 40),
				),
					constraintComponent(
						globalngv1.ByocSharedComponentType_BYOC_SHARED_COMPONENT_TYPE_TIFLASH_WRITE,
						scaleTarget(16, 64, 2, "DATA", 900),
					),
					constraintComponent(
						globalngv1.ByocSharedComponentType_BYOC_SHARED_COMPONENT_TYPE_TIFLASH_COMPUTE,
						scaleTarget(8, 32, 0, "DATA", 500),
					)},
			}, nil
		},
		listSpecs: func(_ context.Context, req *globalngv1.ListNodeSpecsRequest) (*globalngv1.ListNodeSpecsResponse, bizerr.BizError) {
			cpu, memory := int64(4), int64(16)
			if req.GetComponentType() == globalngv1.NodeSpecComponentType_NODE_SPEC_COMPONENT_TYPE_TIFLASH {
				cpu, memory = 8, 32
			}
			return &globalngv1.ListNodeSpecsResponse{NodeSpecs: []*globalngv1.ResourcePoolNodeSpec{{
				NodeSpecKey:   "discovered",
				DisplayName:   "Discovered spec",
				ComponentType: req.GetComponentType(),
				Cpu:           cpu,
				MemoryGi:      memory,
			}}}, nil
		},
	}
	svc := &serviceImpl{regionService: stub, resourcePoolRepo: repo}

	response, err := svc.GetResourcePoolCapacityConstraints(t.Context(), &globalngv1.GetResourcePoolCapacityConstraintsRequest{
		ResourcePoolId: 41,
		OrganizationId: 501,
	})
	require.NoError(t, err)
	require.EqualValues(t, 3, response.GetCurrentCapacity().GetTikv().GetNodeCount())
	require.Len(t, response.GetCurrentCapacity().GetTikv().GetStorages(), 1)
	require.Equal(t, "data", response.GetCurrentCapacity().GetTikv().GetStorages()[0].GetName())
	require.Nil(t, response.GetCurrentCapacity().GetTiflash(), "zero Compute is disabled in current capacity")
	require.NotNil(t, response.GetTikv())
	require.NotNil(t, response.GetTiflash())
	require.Equal(t, "RESOURCE_POOL_MODIFYING", response.GetTikv().GetBlockingReasons()[0].GetCode())
	require.Equal(t, "RESOURCE_POOL_MODIFYING", response.GetTiflash().GetBlockingReasons()[0].GetCode())
}

func TestGetResourcePoolCapacityConstraintsMissingComputeFailsClosed(t *testing.T) {
	t.Parallel()

	svc := &serviceImpl{regionService: capacityRegionServiceStub{
		constraints: func(_ context.Context, _ *globalngv1.GetByocSharedComponentScaleConstraintsRequest) (*globalngv1.ByocSharedComponentScaleConstraints, bizerr.BizError) {
			return &globalngv1.ByocSharedComponentScaleConstraints{Components: []*globalngv1.ByocSharedComponentScaleConstraint{
				constraintComponent(globalngv1.ByocSharedComponentType_BYOC_SHARED_COMPONENT_TYPE_TIKV, scaleTarget(4, 16, 3, "DATA", 500)),
			}}, nil
		},
	}}

	_, err := svc.GetResourcePoolCapacityConstraints(t.Context(), &globalngv1.GetResourcePoolCapacityConstraintsRequest{
		ResourcePoolId: 41,
		OrganizationId: 501,
	})
	require.Error(t, err)
	require.Equal(t, bizerr.ErrCode(codes.FailedPrecondition), err.ErrCode())
}

func TestPublicComponentConstraintsSanitizesAndEnriches(t *testing.T) {
	t.Parallel()

	component := constraintComponent(
		globalngv1.ByocSharedComponentType_BYOC_SHARED_COMPONENT_TYPE_TIKV,
		scaleTarget(4, 16, 3, "DATA", 500),
	)
	component.BlockingReasons = []*globalngv1.ByocSharedComponentScaleConstraintReason{{
		Code:    "UNRECOGNIZED_REGIONAL_CODE",
		Message: "raw regional secret",
	}}
	component.NodeCountConstraints = &globalngv1.ByocSharedComponentScaleNodeCountConstraints{
		Current: 3,
		Min:     3,
		Max:     12,
		Rule:    &globalngv1.ByocSharedComponentScaleNodeCountRule{Type: "MULTIPLE_OF", Value: 3},
	}
	component.NodeSpecConstraints = &globalngv1.ByocSharedComponentScaleNodeSpecConstraints{
		Current: "4C16G",
		Candidates: []*globalngv1.ByocSharedComponentScaleNodeSpecCandidate{{
			Name:      "4C16G",
			Cpu:       4,
			MemoryGi:  16,
			Available: true,
			IsCurrent: true,
			StorageConstraints: []*globalngv1.ByocSharedComponentScaleStorageConstraint{
				{
					StorageName:   "DATA",
					Resizable:     true,
					CurrentSizeGi: 500,
					MaxSizeGi:     proto.Uint32(2000),
				},
				{
					StorageName:   "RAFTLOG",
					Resizable:     true,
					CurrentSizeGi: 100,
					MaxSizeGi:     proto.Uint32(500),
				},
			},
		}},
	}
	discovered := []*globalngv1.ResourcePoolNodeSpec{{
		NodeSpecKey: "public-4c16g",
		DisplayName: "4 vCPU, 16 GiB",
		Cpu:         4,
		MemoryGi:    16,
		Storages: []*globalngv1.NodeSpecStorage{{
			Name: "data", DefaultSizeGi: 500, MinSizeGi: 100, MaxSizeGi: 4000,
		}},
	}}

	result, err := publicComponentConstraints(
		component,
		globalngv1.NodeSpecComponentType_NODE_SPEC_COMPONENT_TYPE_TIKV,
		discovered,
	)
	require.NoError(t, err)
	require.Equal(t, "CAPACITY_CHANGE_NOT_ALLOWED", result.GetBlockingReasons()[0].GetCode())
	require.NotContains(t, result.GetBlockingReasons()[0].GetMessage(), "secret")
	require.Equal(t, "public-4c16g", result.GetNodeSpec().GetCurrent().GetNodeSpecKey())
	require.Len(t, result.GetNodeSpec().GetCurrent().GetStorages(), 1)
	storage := result.GetNodeSpec().GetCurrent().GetStorages()[0]
	require.EqualValues(t, 500, storage.GetCurrentSizeGi())
	require.True(t, storage.GetResizable())
	require.EqualValues(t, 2000, storage.GetMaxSizeGi())
}

func TestPublicComponentConstraintsLegacyCurrentFallbackAndUnknownRule(t *testing.T) {
	t.Parallel()

	component := constraintComponent(
		globalngv1.ByocSharedComponentType_BYOC_SHARED_COMPONENT_TYPE_TIFLASH_COMPUTE,
		scaleTarget(7, 28, 1, "DATA", 700),
	)
	component.NodeCountConstraints = &globalngv1.ByocSharedComponentScaleNodeCountConstraints{
		Rule: &globalngv1.ByocSharedComponentScaleNodeCountRule{Type: "FUTURE_RULE"},
	}

	_, err := publicComponentConstraints(
		component,
		globalngv1.NodeSpecComponentType_NODE_SPEC_COMPONENT_TYPE_TIFLASH,
		nil,
	)
	require.Error(t, err)

	component.NodeCountConstraints.Rule.Type = "NONE"
	component.NodeSpecConstraints = &globalngv1.ByocSharedComponentScaleNodeSpecConstraints{
		Current: "legacy-7c28g",
		Candidates: []*globalngv1.ByocSharedComponentScaleNodeSpecCandidate{{
			Name:      "legacy-7c28g",
			Cpu:       7,
			MemoryGi:  28,
			Available: true,
			IsCurrent: true,
		}},
	}
	result, err := publicComponentConstraints(
		component,
		globalngv1.NodeSpecComponentType_NODE_SPEC_COMPONENT_TYPE_TIFLASH,
		nil,
	)
	require.NoError(t, err)
	current := result.GetNodeSpec().GetCurrent()
	require.Equal(t, "legacy-7c28g", current.GetNodeSpecKey())
	require.False(t, current.GetIsDefault())
	require.EqualValues(t, 700, current.GetStorages()[0].GetDefaultSizeGi())
	require.EqualValues(t, 700, current.GetStorages()[0].GetMinSizeGi())
	require.EqualValues(t, 700, current.GetStorages()[0].GetMaxSizeGi())
}

func resizeRequest(validateOnly bool, compute *globalngv1.ComponentCapacity) *globalngv1.ResizeResourcePoolCapacityRequest {
	tikvCount := int64(6)
	return &globalngv1.ResizeResourcePoolCapacityRequest{
		ResourcePoolId: 41,
		OrganizationId: 501,
		ValidateOnly:   validateOnly,
		TargetCapacity: &globalngv1.ResourcePoolCapacity{
			Tikv:    &globalngv1.ComponentCapacity{NodeCount: &tikvCount},
			Tiflash: compute,
		},
	}
}

func allowedScalePreview() *globalngv1.ScaleByocSharedComponentsResponse {
	return &globalngv1.ScaleByocSharedComponentsResponse{
		Preview: &globalngv1.ByocSharedComponentScalePreviewSummary{Allowed: true},
		Components: []*globalngv1.ByocSharedComponentScalePlanComponent{{
			ComponentType: globalngv1.ByocSharedComponentType_BYOC_SHARED_COMPONENT_TYPE_TIKV,
			Current:       scaleTarget(4, 16, 3, "DATA", 500, "RAFTLOG", 100),
			Target:        scaleTarget(4, 16, 6, "DATA", 500, "RAFTLOG", 100),
			Allowed:       true,
		}},
	}
}

func constraintComponent(
	componentType globalngv1.ByocSharedComponentType,
	current *globalngv1.ByocSharedComponentScaleTarget,
) *globalngv1.ByocSharedComponentScaleConstraint {
	return &globalngv1.ByocSharedComponentScaleConstraint{
		ComponentType: componentType,
		Exists:        true,
		Current:       current,
	}
}

func scaleTarget(cpu, memoryGi, nodeCount uint32, storagePairs ...any) *globalngv1.ByocSharedComponentScaleTarget {
	if len(storagePairs)%2 != 0 {
		panic("storage pairs must contain name and size")
	}
	target := &globalngv1.ByocSharedComponentScaleTarget{
		Cpu:       proto.Uint32(cpu),
		MemoryGi:  proto.Uint32(memoryGi),
		NodeCount: proto.Uint32(nodeCount),
	}
	for index := 0; index+1 < len(storagePairs); index += 2 {
		name, ok := storagePairs[index].(string)
		if !ok {
			panic("storage name must be a string")
		}
		size, ok := storagePairs[index+1].(int)
		if !ok {
			panic("storage size must be an int")
		}
		target.Storages = append(target.Storages, &globalngv1.ByocSharedComponentStorage{
			StorageName: strings.ToUpper(name),
			SizeGi:      proto.Uint32(uint32(size)),
		})
	}
	return target
}