package dbutil

import (
	"context"
	"time"

	"github.com/pingcap/dbaas/central/db/entity"
	"github.com/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
	"gorm.io/gorm"
)

func GetPOCTrialEndedAt(ctx context.Context, tx *gorm.DB, tenantID entity.ID) (*time.Time, error) {
	trial := entity.Trial{}
	if err := entity.NewTrialQuerySet(tx).TenantIDEq(tenantID).TrialTypeEq(entity.TrialTypePOC).One(&trial); err != nil {
		return nil, err
	}

	return trial.EndedAt, nil
}

func ResetPOCTrialEndedAt(ctx context.Context, tx *gorm.DB, tenant entity.Tenant) error {
	var (
		err   error
		trial entity.Trial
	)

	err = entity.NewTrialQuerySet(tx).TenantIDEq(tenant.ID).TrialTypeEq(entity.TrialTypePOC).One(&trial)

	if err != nil {
		if err != gorm.ErrRecordNotFound {
			return err
		}

		log.Warn(ctx, "tenant does not have trial record, create one", zap.Uint64("tenantID", tenant.ID.Uint64()))
		_, err = createTrial(tx, tenant, entity.TrialTypePOC, nil, nil)

		return err
	}

	updater := entity.NewTrialQuerySet(tx).TenantIDEq(tenant.ID).TrialTypeEq(entity.TrialTypePOC).GetUpdater()
	return updater.SetEndedAt(nil).Update()
}

func createTrial(tx *gorm.DB, tenant entity.Tenant, trialType entity.TrialType, estimatedEndedAt, endedAt *time.Time) (*entity.Trial, error) {
	var (
		err   error
		trial entity.Trial
	)

	systemUser, err := SystemUser(tx, tenant.ID)
	if err != nil {
		return nil, err
	}

	trial = entity.Trial{
		CreatedByID:      systemUser.ID,
		TenantID:         tenant.ID,
		StartedAt:        tenant.CreatedAt,
		EstimatedEndedAt: estimatedEndedAt,
		EndedAt:          endedAt,
		TrialType:        trialType,
	}
	err = tx.Create(&trial).Error

	return &trial, err
}

func AddPOCTrialRecord(ctx context.Context, tx *gorm.DB, tenant entity.Tenant) error {
	var (
		err   error
		trial entity.Trial
	)

	err = entity.NewTrialQuerySet(tx).TenantIDEq(tenant.ID).TrialTypeEq(entity.TrialTypePOC).One(&trial)

	if err != nil {
		if err != gorm.ErrRecordNotFound {
			return err
		}

		_, err = createTrial(tx, tenant, entity.TrialTypePOC, nil, nil)
		return err
	}

	log.Warn(ctx, "tenant already had trial record", zap.Uint64("tenantID", tenant.ID.Uint64()))
	return nil
}

func MarkPOCTrialAsEnd(ctx context.Context, tx *gorm.DB, tenant entity.Tenant) error {
	now := time.Now()
	updater := entity.NewTrialQuerySet(tx).TenantIDEq(tenant.ID).TrialTypeEq(entity.TrialTypePOC).EndedAtIsNull().GetUpdater()
	return updater.SetEndedAt(&now).Update()
}