package server

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	braction "github.com/tidbcloud/cluster-service-ng/domain/workflow/action/br"
	bractionmocks "github.com/tidbcloud/cluster-service-ng/domain/workflow/action/br/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"gorm.io/gorm"
)

func TestCreateTiKVCapacityReservation(t *testing.T) {
	ctx := context.Background()

	t.Run("returns existing reservation", func(t *testing.T) {
		scaleTikvDecisionRepo := repomocks.NewScaleTikvDecisionRepo(t)
		scaleManager := bractionmocks.NewScaleManager(t)

		existing := &schema.ScaleTikvDecision{
			SharedPoolID: "10001",
			BizType:      schema.BizTypeImport,
			BizID:        "job-1",
			Status:       schema.ScaleStatusReachExpected,
			NeedScale:    false,
			From:         schema.TiKVSpecUsage{NodeCount: 3},
			To:           schema.TiKVSpecUsage{NodeCount: 3},
		}
		scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{
			SharedPoolID:   "10001",
			BizType:        schema.BizTypeImport,
			BizID:          "job-1",
			IncludeDeleted: true,
		}).Return([]*schema.ScaleTikvDecision{existing}, nil).Once()

		s := &serverImpl{
			scaleTikvDecisionRepo: scaleTikvDecisionRepo,
			scaleManager:          scaleManager,
			sharedPoolRepo:        availableSharedPoolRepoForMutationTest(),
		}

		resp, err := s.CreateTiKVCapacityReservation(ctx, &CreateTiKVCapacityReservationInput{
			CreateTiKVCapacityReservationRequest: &apiv1.CreateTiKVCapacityReservationRequest{
				SharedPoolId:      "10001",
				BizType:           apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT,
				ReservationId:     "job-1",
				RequiredDataBytes: 1024,
			},
		})

		require.NoError(t, err)
		require.NotNil(t, resp)
		assert.Equal(t, apiv1.TiKVCapacityReservationStatus_TIKV_CAPACITY_RESERVATION_STATUS_READY, resp.Status)
		assert.Equal(t, uint32(3), resp.GetFrom().GetNodeCount())
	})

	t.Run("rejects released reservation id", func(t *testing.T) {
		scaleTikvDecisionRepo := repomocks.NewScaleTikvDecisionRepo(t)
		scaleManager := bractionmocks.NewScaleManager(t)

		released := &schema.ScaleTikvDecision{
			SharedPoolID: "10001",
			BizType:      schema.BizTypeImport,
			BizID:        "job-1",
		}
		released.DeletedAt = gorm.DeletedAt{Time: time.Now(), Valid: true}

		scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{
			SharedPoolID:   "10001",
			BizType:        schema.BizTypeImport,
			BizID:          "job-1",
			IncludeDeleted: true,
		}).Return([]*schema.ScaleTikvDecision{released}, nil).Once()

		s := &serverImpl{
			scaleTikvDecisionRepo: scaleTikvDecisionRepo,
			scaleManager:          scaleManager,
			sharedPoolRepo:        availableSharedPoolRepoForMutationTest(),
		}

		resp, err := s.CreateTiKVCapacityReservation(ctx, &CreateTiKVCapacityReservationInput{
			CreateTiKVCapacityReservationRequest: &apiv1.CreateTiKVCapacityReservationRequest{
				SharedPoolId:      "10001",
				BizType:           apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT,
				ReservationId:     "job-1",
				RequiredDataBytes: 1024,
			},
		})

		require.Error(t, err)
		assert.Nil(t, resp)
		assert.Equal(t, codes.AlreadyExists, status.Code(err))
	})

	t.Run("creates new reservation", func(t *testing.T) {
		scaleTikvDecisionRepo := repomocks.NewScaleTikvDecisionRepo(t)
		scaleManager := bractionmocks.NewScaleManager(t)

		created := &schema.ScaleTikvDecision{
			SharedPoolID: "10001",
			BizType:      schema.BizTypeImport,
			BizID:        "job-2",
			Status:       schema.ScaleStatusReachExpected,
			NeedScale:    false,
			From:         schema.TiKVSpecUsage{NodeCount: 3},
			To:           schema.TiKVSpecUsage{NodeCount: 3},
		}

		scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{
			SharedPoolID:   "10001",
			BizType:        schema.BizTypeImport,
			BizID:          "job-2",
			IncludeDeleted: true,
		}).Return([]*schema.ScaleTikvDecision{}, nil).Once()
		scaleManager.On("Scale", mock.Anything, mock.MatchedBy(func(input *braction.ScaleInput) bool {
			return input.SharedPoolID == "10001" &&
				input.BizType == schema.BizTypeImport &&
				input.BizID == "job-2" &&
				input.FullfilDataBytes == 2048 &&
				input.PostScaleWriteDB != nil
		})).Return(nil).Once()
		scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{
			SharedPoolID:   "10001",
			BizType:        schema.BizTypeImport,
			BizID:          "job-2",
			IncludeDeleted: false,
		}).Return([]*schema.ScaleTikvDecision{created}, nil).Once()

		s := &serverImpl{
			scaleTikvDecisionRepo: scaleTikvDecisionRepo,
			scaleManager:          scaleManager,
			sharedPoolRepo:        availableSharedPoolRepoForMutationTest(),
		}

		resp, err := s.CreateTiKVCapacityReservation(ctx, &CreateTiKVCapacityReservationInput{
			CreateTiKVCapacityReservationRequest: &apiv1.CreateTiKVCapacityReservationRequest{
				SharedPoolId:      "10001",
				BizType:           apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT,
				ReservationId:     "job-2",
				RequiredDataBytes: 2048,
			},
		})

		require.NoError(t, err)
		require.NotNil(t, resp)
		assert.Equal(t, "job-2", resp.ReservationId)
		assert.Equal(t, apiv1.TiKVCapacityReservationStatus_TIKV_CAPACITY_RESERVATION_STATUS_READY, resp.Status)
	})

	t.Run("auto scaler create piggybacks on active reservation", func(t *testing.T) {
		scaleTikvDecisionRepo := repomocks.NewScaleTikvDecisionRepo(t)
		scaleManager := bractionmocks.NewScaleManager(t)

		active := &schema.ScaleTikvDecision{
			SharedPoolID: "10001",
			BizType:      schema.BizTypeAutoScaler,
			BizID:        "cpu-existing",
			Status:       schema.ScaleStatusExecuting,
			NeedScale:    true,
			From:         schema.TiKVSpecUsage{NodeCount: 6},
			To:           schema.TiKVSpecUsage{NodeCount: 9},
		}
		active.ID = 22

		scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{
			SharedPoolID:   "10001",
			BizType:        schema.BizTypeAutoScaler,
			BizID:          "cpu-new",
			IncludeDeleted: true,
		}).Return([]*schema.ScaleTikvDecision{}, nil).Once()
		scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{
			SharedPoolID:   "10001",
			BizType:        schema.BizTypeAutoScaler,
			IncludeDeleted: false,
		}).Return([]*schema.ScaleTikvDecision{active}, nil).Once()
		scaleManager.On("WaitAndSetScaleFinish", mock.Anything, "10001", uint(22)).Return(false, nil).Once()
		scaleTikvDecisionRepo.On("Load", mock.Anything, uint(22), true).Return(active, nil).Once()

		s := &serverImpl{
			scaleTikvDecisionRepo: scaleTikvDecisionRepo,
			scaleManager:          scaleManager,
			sharedPoolRepo:        availableSharedPoolRepoForMutationTest(),
		}

		resp, err := s.CreateTiKVCapacityReservation(ctx, &CreateTiKVCapacityReservationInput{
			CreateTiKVCapacityReservationRequest: &apiv1.CreateTiKVCapacityReservationRequest{
				SharedPoolId:      "10001",
				BizType:           apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_AUTO_SCALER,
				ReservationId:     "cpu-new",
				RequiredDataBytes: 4096,
			},
		})

		require.NoError(t, err)
		require.NotNil(t, resp)
		assert.Equal(t, "cpu-existing", resp.GetReservationId())
		assert.Equal(t, apiv1.TiKVCapacityReservationStatus_TIKV_CAPACITY_RESERVATION_STATUS_PROVISIONING, resp.GetStatus())
	})

	t.Run("treats duplicate create as idempotent success", func(t *testing.T) {
		scaleTikvDecisionRepo := repomocks.NewScaleTikvDecisionRepo(t)
		scaleManager := bractionmocks.NewScaleManager(t)

		created := &schema.ScaleTikvDecision{
			SharedPoolID: "10001",
			BizType:      schema.BizTypeImport,
			BizID:        "job-4",
			Status:       schema.ScaleStatusReachExpected,
			NeedScale:    false,
			From:         schema.TiKVSpecUsage{NodeCount: 3},
			To:           schema.TiKVSpecUsage{NodeCount: 3},
		}

		scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{
			SharedPoolID:   "10001",
			BizType:        schema.BizTypeImport,
			BizID:          "job-4",
			IncludeDeleted: true,
		}).Return([]*schema.ScaleTikvDecision{}, nil).Once()
		scaleManager.On("Scale", mock.Anything, mock.Anything).Return(status.Error(codes.AlreadyExists, "duplicate")).Once()
		scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{
			SharedPoolID:   "10001",
			BizType:        schema.BizTypeImport,
			BizID:          "job-4",
			IncludeDeleted: false,
		}).Return([]*schema.ScaleTikvDecision{created}, nil).Once()

		s := &serverImpl{
			scaleTikvDecisionRepo: scaleTikvDecisionRepo,
			scaleManager:          scaleManager,
			sharedPoolRepo:        availableSharedPoolRepoForMutationTest(),
		}

		resp, err := s.CreateTiKVCapacityReservation(ctx, &CreateTiKVCapacityReservationInput{
			CreateTiKVCapacityReservationRequest: &apiv1.CreateTiKVCapacityReservationRequest{
				SharedPoolId:      "10001",
				BizType:           apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT,
				ReservationId:     "job-4",
				RequiredDataBytes: 2048,
			},
		})

		require.NoError(t, err)
		require.NotNil(t, resp)
		assert.Equal(t, "job-4", resp.ReservationId)
	})
}

func TestGetTiKVCapacityReservation(t *testing.T) {
	ctx := context.Background()
	scaleTikvDecisionRepo := repomocks.NewScaleTikvDecisionRepo(t)
	scaleManager := bractionmocks.NewScaleManager(t)

	decision := &schema.ScaleTikvDecision{
		SharedPoolID: "10001",
		BizType:      schema.BizTypeAutoScaler,
		BizID:        "as-1",
		Status:       schema.ScaleStatusExecuting,
		NeedScale:    true,
		From:         schema.TiKVSpecUsage{NodeCount: 3},
		To:           schema.TiKVSpecUsage{NodeCount: 5},
	}
	decision.ID = 11

	reloaded := *decision
	reloaded.Status = schema.ScaleStatusReachExpected

	scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{
		SharedPoolID:   "10001",
		BizType:        schema.BizTypeAutoScaler,
		BizID:          "as-1",
		IncludeDeleted: true,
	}).Return([]*schema.ScaleTikvDecision{decision}, nil).Once()
	scaleManager.On("WaitAndSetScaleFinish", mock.Anything, "10001", uint(11)).Return(true, nil).Once()
	scaleTikvDecisionRepo.On("Load", mock.Anything, uint(11), true).Return(&reloaded, nil).Once()

	s := &serverImpl{
		scaleTikvDecisionRepo: scaleTikvDecisionRepo,
		scaleManager:          scaleManager,
		sharedPoolRepo:        availableSharedPoolRepoForMutationTest(),
	}

	resp, err := s.GetTiKVCapacityReservation(ctx, &GetTiKVCapacityReservationInput{
		GetTiKVCapacityReservationRequest: &apiv1.GetTiKVCapacityReservationRequest{
			SharedPoolId:  "10001",
			BizType:       apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_AUTO_SCALER,
			ReservationId: "as-1",
		},
	})

	require.NoError(t, err)
	require.NotNil(t, resp)
	assert.Equal(t, apiv1.TiKVCapacityReservationStatus_TIKV_CAPACITY_RESERVATION_STATUS_READY, resp.Status)
	assert.Equal(t, uint32(5), resp.GetTo().GetNodeCount())
}

func TestDeleteTiKVCapacityReservation(t *testing.T) {
	ctx := context.Background()
	scaleTikvDecisionRepo := repomocks.NewScaleTikvDecisionRepo(t)
	scaleManager := bractionmocks.NewScaleManager(t)

	decision := &schema.ScaleTikvDecision{
		SharedPoolID: "10001",
		BizType:      schema.BizTypeImport,
		BizID:        "job-3",
	}
	decision.ID = 12

	scaleTikvDecisionRepo.On("List", mock.Anything, &repo.ScaleTikvDecisionFilter{
		SharedPoolID:   "10001",
		BizType:        schema.BizTypeImport,
		BizID:          "job-3",
		IncludeDeleted: true,
	}).Return([]*schema.ScaleTikvDecision{decision}, nil).Once()
	scaleTikvDecisionRepo.On("Delete", mock.Anything, uint(12)).Return(nil).Once()

	s := &serverImpl{
		scaleTikvDecisionRepo: scaleTikvDecisionRepo,
		scaleManager:          scaleManager,
		sharedPoolRepo:        availableSharedPoolRepoForMutationTest(),
	}

	err := s.DeleteTiKVCapacityReservation(ctx, &DeleteTiKVCapacityReservationInput{
		DeleteTiKVCapacityReservationRequest: &apiv1.DeleteTiKVCapacityReservationRequest{
			SharedPoolId:  "10001",
			BizType:       apiv1.TiKVCapacityReservationBizType_TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT,
			ReservationId: "job-3",
		},
	})

	require.NoError(t, err)
}