package tests

import (
	"context"
	"testing"

	"github.com/pingcap/dbaas/billing2/constant"
	"github.com/pingcap/dbaas/billing2/db/entity"
	"github.com/pingcap/dbaas/billing2/forms"
	"github.com/pingcap/dbaas/billing2/handler/eventbushandler"
	"github.com/pingcap/dbaas/billing2/handler/tidbmgmtsvchandler"
	"github.com/pingcap/dbaas/billing2/model/creditriskcontrol"
	"github.com/pingcap/dbaas/billing2/utils/tests"
	"github.com/stretchr/testify/suite"
)

type TenantCloseEssentialV1ActionSuite struct {
	tests.BillingTestSuite
}

func TestTenantCloseEssentialV1ActionSuite(t *testing.T) {
	suite.Run(t, new(TenantCloseEssentialV1ActionSuite))
}

// essentialV1MgmtMock fakes tidb-mgmt-service cluster listing; embedding the
// interface keeps the mock compatible when other methods are added.
type essentialV1MgmtMock struct {
	tidbmgmtsvchandler.MgmtService
	clusters []*forms.Cluster
}

func (m *essentialV1MgmtMock) GetEssentialV1ClustersByTenantID(ctx context.Context, tenantID uint64) (ret []*forms.Cluster, err error) {
	return m.clusters, nil
}

func (suite *TenantCloseEssentialV1ActionSuite) setupEssentialV1Clusters() []*forms.Cluster {
	return []*forms.Cluster{
		{ID: 101, Name: "essential-a", ServicePlan: "Essential"},
		{ID: 202, Name: "essential-b", ServicePlan: "Essential"},
	}
}

func (suite *TenantCloseEssentialV1ActionSuite) spendLimitEvents() []*eventbushandler.Event {
	return eventbushandler.Get().(*tests.LocalEventBusHandlerImpl).GetAllTxEvents(constant.SpendLimitEventType)
}

func (suite *TenantCloseEssentialV1ActionSuite) TestThrottleEssentialV1Clusters() {
	assert := suite.Assert()
	tenantID := uint(1)
	tidbmgmtsvchandler.Set(&essentialV1MgmtMock{clusters: suite.setupEssentialV1Clusters()})

	err := creditriskcontrol.ThrottleEssentialV1Clusters(tenantID, &creditriskcontrol.ActionParams{})
	assert.NoError(err)

	events := suite.spendLimitEvents()
	assert.Len(events, 2)
	for _, clusterID := range []uint{101, 202} {
		info := suite.spendLimitInfoOfEvent(events, clusterID)
		assert.Equal(entity.SpendLimitStatus_Closing, info.Status)
	}
}

func (suite *TenantCloseEssentialV1ActionSuite) TestRecoverEssentialV1Clusters() {
	assert := suite.Assert()
	tenantID := uint(1)
	tidbmgmtsvchandler.Set(&essentialV1MgmtMock{clusters: suite.setupEssentialV1Clusters()})

	err := creditriskcontrol.RecoverEssentialV1Clusters(tenantID, &creditriskcontrol.ActionParams{})
	assert.NoError(err)

	events := suite.spendLimitEvents()
	assert.Len(events, 2)
	for _, clusterID := range []uint{101, 202} {
		info := suite.spendLimitInfoOfEvent(events, clusterID)
		assert.Equal(entity.SpendLimitStatus_Recovered, info.Status)
	}
}

// TestThrottleThenCancelClosingTenant covers the recover path wired into the
// existing cancel action: throttling publishes closing events, then
// cancel_closing_org must publish recovered events for the essential v1
// clusters (starter rows absent -> the starter leg is a no-op).
func (suite *TenantCloseEssentialV1ActionSuite) TestThrottleThenCancelClosingTenant() {
	assert := suite.Assert()
	tenantID := uint(1)
	tidbmgmtsvchandler.Set(&essentialV1MgmtMock{clusters: suite.setupEssentialV1Clusters()})

	err := creditriskcontrol.ThrottleEssentialV1Clusters(tenantID, &creditriskcontrol.ActionParams{})
	assert.NoError(err)
	err = creditriskcontrol.CancelClosingTenant(tenantID, &creditriskcontrol.ActionParams{})
	assert.NoError(err)

	events := suite.spendLimitEvents()
	assert.Len(events, 4)
	closingCount := 0
	recoveredCount := 0
	for _, e := range events {
		info, ok := e.EventData.(*forms.SpendLimitInfo)
		assert.True(ok)
		if info.Status == entity.SpendLimitStatus_Closing {
			closingCount++
		}
		if info.Status == entity.SpendLimitStatus_Recovered {
			recoveredCount++
		}
	}
	assert.Equal(2, closingCount)
	assert.Equal(2, recoveredCount)
}

func (suite *TenantCloseEssentialV1ActionSuite) TestEssentialV1ActionsNoClusters() {
	assert := suite.Assert()
	tenantID := uint(1)
	tidbmgmtsvchandler.Set(&essentialV1MgmtMock{clusters: nil})

	err := creditriskcontrol.ThrottleEssentialV1Clusters(tenantID, &creditriskcontrol.ActionParams{})
	assert.NoError(err)
	err = creditriskcontrol.RecoverEssentialV1Clusters(tenantID, &creditriskcontrol.ActionParams{})
	assert.NoError(err)
	assert.Empty(suite.spendLimitEvents())
}

// spendLimitInfoOfEvent returns the SpendLimitInfo published for the given
// cluster and asserts the synthetic event carries the cluster identity, the
// same fields a starter spend-limit event carries.
func (suite *TenantCloseEssentialV1ActionSuite) spendLimitInfoOfEvent(events []*eventbushandler.Event, clusterID uint) *forms.SpendLimitInfo {
	assert := suite.Assert()
	for _, e := range events {
		info, ok := e.EventData.(*forms.SpendLimitInfo)
		assert.True(ok)
		if info.RangeID != clusterID {
			continue
		}
		assert.Equal(uint(1), info.TenantID)
		assert.Equal(clusterID, info.RangeIDStr)
		assert.Equal(entity.SpendLimitRangeType_Cluster, info.RangeType)
		assert.Equal(entity.SpendLimitBizType_ServerlessCluster, info.BizType)
		assert.Empty(info.Name)
		assert.Empty(info.BizData)
		return info
	}
	assert.Failf("no spend-limit event for cluster", "clusterID: %d", clusterID)
	return nil
}
