package repository

import (
	"context"
	"errors"
	"fmt"
	"maps"
	"time"

	"github.com/samber/lo"
	"github.com/tidbcloud/pkgv2/db"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"

	openapicommonv1beta1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_cloud_open_api/common/v1beta1"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/domain/repository"
	"github.com/tidbcloud/tidb-management-service/tidb-mgmt-service/infra/repository/schema"
	"github.com/tidbcloud/tidb-management-service/utilities/enum"
	"github.com/tidbcloud/tidb-management-service/utilities/grpc/metadata"
)

type clusterRepositoryImpl struct {
	dbProvider db.Provider
}

func NewClusterRepository(dbProvider db.Provider) repository.ClusterRepository {
	return clusterRepositoryImpl{dbProvider: dbProvider}
}

func (c clusterRepositoryImpl) CountDisplayNameUnderProject(
	ctx context.Context,
	projectId uint64,
	displayName string,
) (int64, error) {
	ctx, span := tracer.Start(ctx, "ClusterRepo.CountDisplayNameUnderProject")
	defer span.End()

	rdb := c.dbProvider.ReadDB(ctx)

	var count int64
	err := rdb.Model(&schema.Cluster{}).
		Where("project_id = ? and display_name = ?", projectId, displayName).
		Count(&count).
		Error
	if err != nil {
		return 0, err
	}
	return count, nil
}

func (c clusterRepositoryImpl) CountDisplayNameUnderOrg(
	ctx context.Context,
	orgId uint64,
	displayName string,
) (int64, error) {
	ctx, span := tracer.Start(ctx, "ClusterRepo.CountDisplayNameUnderOrg")
	defer span.End()

	rdb := c.dbProvider.ReadDB(ctx)

	var count int64
	err := rdb.Model(&schema.Cluster{}).
		Where("organization_id = ? and display_name = ?", orgId, displayName).
		Count(&count).
		Error
	if err != nil {
		return 0, err
	}
	return count, nil
}

func (c clusterRepositoryImpl) InsertNewCluster(ctx context.Context, cluster *entity.Cluster) (uint64, error) {
	ctx, span := tracer.Start(ctx, "ClusterRepo.InsertNewCluster")
	defer span.End()

	wdb := c.dbProvider.WriteDB(ctx)

	clusterRecord := cvtClusterRecord(cluster)

	err := wdb.Create(clusterRecord).Error
	if err != nil {
		return 0, err
	}
	return clusterRecord.ClusterID, nil
}

func (c clusterRepositoryImpl) GetCluster(
	ctx context.Context,
	clusterId uint64,
	includeDeleted bool,
) (*entity.Cluster, error) {
	ctx, span := tracer.Start(ctx, "ClusterRepo.GetCluster")
	defer span.End()

	rdb := c.dbProvider.ReadDB(ctx)

	var cluster schema.Cluster

	tx := rdb.Model(&schema.Cluster{})
	if includeDeleted {
		tx = tx.Unscoped()
	}
	err := tx.Model(&schema.Cluster{}).Where("cluster_id = ?", clusterId).Take(&cluster).Error
	if err != nil {
		return nil, err
	}
	ce, err := cvtClusterEntity(&cluster)
	if err != nil {
		return nil, err
	}
	return ce, nil
}

func (c clusterRepositoryImpl) DeleteCluster(ctx context.Context, clusterId uint64) error {
	ctx, span := tracer.Start(ctx, "ClusterRepo.DeleteCluster")
	defer span.End()

	if err := metadata.Validate(ctx, &metadata.ExpectRequest[uint64]{
		ClusterId: clusterId,
	}); err != nil {
		return err
	}

	wdb := c.dbProvider.WriteDB(ctx)

	err := wdb.Model(&schema.Cluster{}).Where("cluster_id = ?", clusterId).Delete(&schema.Cluster{}).Error
	if err != nil {
		return err
	}
	return nil
}

func (c clusterRepositoryImpl) UpdateDisplayName(ctx context.Context, clusterId uint64, displayName string) error {
	ctx, span := tracer.Start(ctx, "ClusterRepo.UpdateDisplayName")
	defer span.End()

	wdb := c.dbProvider.WriteDB(ctx)
	err := wdb.Model(&schema.Cluster{}).Where("cluster_id = ?", clusterId).Update("display_name", displayName).Error
	if err != nil {
		return err
	}
	return nil
}

func (c clusterRepositoryImpl) UpdateServicePlan(ctx context.Context, clusterId uint64, servicePlan string) error {
	ctx, span := tracer.Start(ctx, "ClusterRepo.UpdateServicePlan")
	defer span.End()

	wdb := c.dbProvider.WriteDB(ctx)
	err := wdb.Model(&schema.Cluster{}).Where("cluster_id = ?", clusterId).Update("service_plan", servicePlan).Error
	if err != nil {
		return err
	}
	return nil
}

func (c clusterRepositoryImpl) UpdateRuVersion(ctx context.Context, clusterId uint64, ruVersion int32) error {
	ctx, span := tracer.Start(ctx, "ClusterRepo.UpdateRuVersion")
	defer span.End()

	wdb := c.dbProvider.WriteDB(ctx)
	err := wdb.Model(&schema.Cluster{}).Where("cluster_id = ?", clusterId).Update("ru_version", ruVersion).Error
	if err != nil {
		return err
	}
	return nil
}

func (c clusterRepositoryImpl) UpdateProjectId(ctx context.Context, clusterId, projectId uint64) error {
	ctx, span := tracer.Start(ctx, "ClusterRepo.UpdateProjectId")
	defer span.End()

	wdb := c.dbProvider.WriteDB(ctx)
	err := wdb.Model(&schema.Cluster{}).Where("cluster_id = ?", clusterId).Update("project_id", projectId).Error
	if err != nil {
		return err
	}
	return nil
}

func (c clusterRepositoryImpl) UpdateShadowState(ctx context.Context,
	clusterId uint64,
	shadowState entity.ShadowState,
	originalState entity.ShadowState) error {
	ctx, span := tracer.Start(ctx, "ClusterRepo.UpdateShadowState")
	defer span.End()

	wdb := c.dbProvider.WriteDB(ctx)
	result := wdb.Model(&schema.Cluster{}).
		Where("cluster_id = ?", clusterId).
		Where("shadow_state = ?", originalState).
		Update("shadow_state", shadowState)

	if result.Error != nil {
		return result.Error
	}

	if result.RowsAffected == 0 {
		return errors.New("optimistic lock failed: state has changed or record not found")
	}
	return nil
}

func (c clusterRepositoryImpl) UpdateLabels(ctx context.Context, clusterId uint64, labels map[string]string) error {
	ctx, span := tracer.Start(ctx, "ClusterRepo.UpdateLabels")
	defer span.End()

	wdb := c.dbProvider.WriteDB(ctx)

	var cluster schema.Cluster
	err := wdb.Model(&schema.Cluster{}).Where("cluster_id = ?", clusterId).Take(&cluster).Error
	if err != nil {
		return err
	}
	maps.Copy(cluster.Labels, labels)
	err = wdb.Model(&schema.Cluster{}).Where("cluster_id = ?", clusterId).Update("labels", cluster.Labels).Error
	if err != nil {
		return err
	}
	return nil
}

func (c clusterRepositoryImpl) TouchCluster(ctx context.Context, clusterId uint64) error {
	ctx, span := tracer.Start(ctx, "ClusterRepo.TouchCluster")
	defer span.End()

	wdb := c.dbProvider.WriteDB(ctx)

	err := wdb.Model(&schema.Cluster{}).
		Where("cluster_id = ?", clusterId).
		Update("updated_at", gorm.Expr("NOW()")).
		Error
	if err != nil {
		return err
	}
	return nil
}

func (c clusterRepositoryImpl) QueryClusterByFilterAndPage(
	ctx context.Context,
	filter *repository.ClusterFilter,
	offset, pageSize int,
) ([]*entity.Cluster, error) {
	ctx, span := tracer.Start(ctx, "ClusterRepo.QueryClusterByFilterAndPage")
	defer span.End()

	rdb := c.dbProvider.ReadDB(ctx)

	var clusters []*schema.Cluster

	if filter == nil {
		filter = &repository.ClusterFilter{}
	}

	tx := rdb.Model(&schema.Cluster{})
	tx = filterQuery(tx, filter)

	// sort
	if filter.CreatedAtAsc {
		tx = tx.Order("created_at asc")
	} else {
		tx = tx.Order("created_at desc")
	}

	// page
	if offset > 0 {
		tx = tx.Offset(offset)
	}
	if pageSize > 0 {
		tx = tx.Limit(pageSize)
	}

	err := tx.Find(&clusters).Error
	if err != nil {
		return nil, err
	}

	if len(clusters) == 0 {
		return []*entity.Cluster{}, nil
	}

	if len(clusters) > 1000 {
		tilog.Warn(ctx, "QueryClusterByFilterAndPage returns more than 1000 clusters, which may cause performance issues")
	}

	clustersEntities := make([]*entity.Cluster, 0, len(clusters))
	// convert to entity
	for _, cluster := range clusters {
		ce, err := cvtClusterEntity(cluster)
		if err != nil {
			return nil, err
		}
		clustersEntities = append(clustersEntities, ce)
	}

	return clustersEntities, nil
}

func (c clusterRepositoryImpl) QueryPreparedShadowClustersForUpdate(
	ctx context.Context,
	shadowPoolId uint64,
	count int,
) ([]*entity.Cluster, error) {
	ctx, span := tracer.Start(ctx, "ClusterRepo.QueryPreparedShadowClustersForUpdate")
	defer span.End()

	if count <= 0 {
		return nil, fmt.Errorf("count must be greater than 0")
	}

	var clusters []*schema.Cluster
	wdb := c.dbProvider.WriteDB(ctx)
	err := wdb.Clauses(clause.Locking{
		Strength: "UPDATE",
		Options:  "SKIP LOCKED",
	}).
		Where("shadow_pool_id = ?", shadowPoolId).
		Where("shadow_state = ?", int8(entity.ShadowStatePrepared)).
		Limit(count).
		Find(&clusters).Error

	if err != nil {
		return nil, err
	}

	if len(clusters) < count {
		return nil, fmt.Errorf("not enough prepared clusters in shadow pool %d", shadowPoolId)
	}

	clustersEntities := make([]*entity.Cluster, 0, len(clusters))
	// convert to entity
	for _, cluster := range clusters {
		ce, err := cvtClusterEntity(cluster)
		if err != nil {
			return nil, err
		}
		clustersEntities = append(clustersEntities, ce)
	}

	return clustersEntities, nil
}

func (c clusterRepositoryImpl) CountClusterByFilter(
	ctx context.Context,
	filter *repository.ClusterFilter,
) (int64, error) {
	ctx, span := tracer.Start(ctx, "ClusterRepo.CountClusterByFilter")
	defer span.End()

	rdb := c.dbProvider.ReadDB(ctx)

	if filter == nil {
		filter = &repository.ClusterFilter{}
	}

	var count int64

	tx := rdb.Model(&schema.Cluster{})
	if shouldForceOrgDeployPlanCountIndex(filter) {
		tx = tx.Set("gorm:query_option", "FORCE INDEX (`idx_clusters_org_deploy_plan_deleted`)")
	}
	tx = filterQuery(tx, filter)

	err := tx.Count(&count).Error
	if err != nil {
		return 0, err
	}
	return count, nil
}

func shouldForceOrgDeployPlanCountIndex(filter *repository.ClusterFilter) bool {
	if filter == nil {
		return false
	}
	if filter.QueryDeleted {
		return false
	}
	if len(filter.OrganizationIds) != 1 {
		return false
	}
	if filter.DeployType == "" {
		return false
	}

	return true
}

// shadowPoolClusterCountRow is an internal scan target for CountShadowPoolClustersGrouped.
type shadowPoolClusterCountRow struct {
	PoolID      uint64 `gorm:"column:shadow_pool_id"`
	ShadowState int8   `gorm:"column:shadow_state"`
	IsDeleted   bool   `gorm:"column:is_deleted"`
	Count       int64  `gorm:"column:count"`
}

func (c clusterRepositoryImpl) CountShadowPoolClustersGrouped(
	ctx context.Context,
	poolIDs []uint64,
) ([]repository.ShadowPoolClusterCount, error) {
	ctx, span := tracer.Start(ctx, "ClusterRepo.CountShadowPoolClustersGrouped")
	defer span.End()

	if len(poolIDs) == 0 {
		return nil, nil
	}

	rdb := c.dbProvider.ReadDB(ctx)

	var rows []shadowPoolClusterCountRow
	err := rdb.Unscoped().
		Model(&schema.Cluster{}).
		Select("shadow_pool_id, shadow_state, (deleted_at IS NOT NULL) AS is_deleted, COUNT(*) AS count").
		Where("shadow_pool_id IN ?", poolIDs).
		Group("shadow_pool_id, shadow_state, is_deleted").
		Scan(&rows).Error
	if err != nil {
		return nil, err
	}

	result := make([]repository.ShadowPoolClusterCount, 0, len(rows))
	for _, row := range rows {
		result = append(result, repository.ShadowPoolClusterCount{
			PoolID:      row.PoolID,
			ShadowState: entity.ShadowState(row.ShadowState),
			IsDeleted:   row.IsDeleted,
			Count:       row.Count,
		})
	}
	return result, nil
}

func filterQuery(tx *gorm.DB, filter *repository.ClusterFilter) *gorm.DB {
	if filter.QueryDeleted {
		tx = tx.Unscoped()
	}
	if len(filter.ClusterIds) > 0 {
		tx = tx.Where("cluster_id in ?", filter.ClusterIds)
	}
	if len(filter.OrganizationIds) > 0 {
		tx = tx.Where("organization_id in ?", filter.OrganizationIds)
	}
	if len(filter.ProjectIds) > 0 {
		tx = tx.Where("project_id in ?", filter.ProjectIds)
	}
	if len(filter.ExcludeProjectIds) > 0 {
		tx = tx.Where("project_id NOT IN ?", filter.ExcludeProjectIds)
	}
	if len(filter.CloudProviders) > 0 {
		tx = tx.Where("cloud_provider in ?", filter.CloudProviders)
	}
	if len(filter.Regions) > 0 {
		tx = tx.Where("region in ?", filter.Regions)
	}
	if len(filter.RegionPairs) > 0 {
		pairedRegionQuery := tx.Session(&gorm.Session{}).Where(
			"cloud_provider = ? AND region = ?",
			filter.RegionPairs[0].CloudProvider,
			filter.RegionPairs[0].Region,
		)
		for _, pair := range filter.RegionPairs[1:] {
			pairedRegionQuery = pairedRegionQuery.Or("cloud_provider = ? AND region = ?", pair.CloudProvider, pair.Region)
		}
		tx = tx.Where(pairedRegionQuery)
	}
	if filter.DeployType != "" {
		tx = tx.Where("deploy_type = ?", filter.DeployType)
	}
	for k, v := range filter.Labels {
		tx = tx.Where(fmt.Sprintf("JSON_EXTRACT(labels,'$.%q')= ?", k), v)
	}
	if len(filter.ServicePlans) > 0 {
		tx = tx.Where("service_plan in ?", filter.ServicePlans)
	}
	if len(filter.EngineTypes) > 0 {
		tx = tx.Where("engine_type in ?", filter.EngineTypes)
	}
	if len(filter.TagCanonicalNames) > 0 {
		tx = addTagFilterQuery(tx, filter.TagCanonicalNames)
	}
	if filter.Name != "" {
		tx = tx.Where("display_name like ?", fmt.Sprintf("%%%s%%", filter.Name))
	}
	if filter.PermissionFilter != nil {
		pf := filter.PermissionFilter
		if !pf.All {
			tx = tx.Where(tx.Session(&gorm.Session{}).
				Where("project_id IN ?", pf.ProjectIds).
				Or("cluster_id IN ?", pf.ClusterIds),
			)
		}
	}
	if len(filter.ShadowPoolIds) > 0 {
		tx = tx.Where("shadow_pool_id in ?", filter.ShadowPoolIds)
	}
	if len(filter.ShadowStates) > 0 {
		tx = tx.Where("shadow_state in ?", lo.Map(filter.ShadowStates, func(item entity.ShadowState, index int) int8 {
			return int8(item)
		}))
	}
	return tx
}

func (c clusterRepositoryImpl) SetClusterRegion(ctx context.Context, clusterId uint64, provider, region string) error {
	ctx, span := tracer.Start(ctx, "ClusterRepo.SetClusterRegion")
	defer span.End()

	wdb := c.dbProvider.WriteDB(ctx)

	var cluster schema.Cluster
	err := wdb.Model(&schema.Cluster{}).Where("cluster_id = ?", clusterId).Take(&cluster).Error
	if err != nil {
		return err
	}

	data := map[string]any{"cloud_provider": provider, "region": region}
	err = wdb.Model(&schema.Cluster{}).Where("cluster_id = ?", clusterId).Updates(data).Error
	if err != nil {
		return err
	}

	return nil
}

func cvtClusterRecord(cluster *entity.Cluster) *schema.Cluster {
	engineType := cluster.EngineType
	if engineType == openapicommonv1beta1.EngineType_ENGINE_TYPE_UNSPECIFIED {
		engineType = openapicommonv1beta1.EngineType_MYSQL_COMPATIBLE
	}

	// Then persist cluster info to database
	return &schema.Cluster{
		Name:              fmt.Sprintf("clusters/%d", cluster.ClusterId),
		DisplayName:       cluster.DisplayName,
		OrganizationID:    cluster.OrganizationId,
		ProjectID:         cluster.ProjectId,
		OriginalProjectID: cluster.OriginalProjectId,
		ClusterID:         cluster.ClusterId,
		DeployType:        cluster.DeployType,
		CloudProvider:     cluster.CloudProvider.String(),
		Region:            cluster.Region,
		CreatedBy:         cluster.CreatedBy,
		Labels:            cluster.Labels,
		Annotations:       cluster.Annotations,
		ServicePlan:       cluster.ServicePlan.String(),
		EngineType:        engineType.String(),
		ShadowPoolID:      cluster.ShadowPoolId,
		ShadowState:       int8(cluster.ShadowState),
		RuVersion:         cluster.RuVersion,
	}
}

func cvtClusterEntity(clusterRecord *schema.Cluster) (*entity.Cluster, error) {
	var deletedAt *time.Time
	if clusterRecord.DeletedAt.Valid {
		deletedAt = &clusterRecord.DeletedAt.Time
	}

	cloudProvider, err := enum.FromString[openapicommonv1beta1.Region_CloudProvider](clusterRecord.CloudProvider)
	if err != nil {
		return nil, err
	}

	servicePlan, err := enum.FromString[openapicommonv1beta1.ServicePlan](clusterRecord.ServicePlan)
	if err != nil {
		return nil, err
	}

	if clusterRecord.EngineType == "" {
		clusterRecord.EngineType = openapicommonv1beta1.EngineType_MYSQL_COMPATIBLE.String()
	}

	engineType, err := enum.FromString[openapicommonv1beta1.EngineType](clusterRecord.EngineType)
	if err != nil {
		return nil, err
	}

	if clusterRecord.Labels == nil {
		clusterRecord.Labels = make(map[string]string)
	}
	if clusterRecord.Annotations == nil {
		clusterRecord.Annotations = make(map[string]string)
	}

	return &entity.Cluster{
		ClusterId:         clusterRecord.ClusterID,
		ProjectId:         clusterRecord.ProjectID,
		OriginalProjectId: clusterRecord.OriginalProjectID,
		OrganizationId:    clusterRecord.OrganizationID,
		DisplayName:       clusterRecord.DisplayName,
		DeployType:        clusterRecord.DeployType,
		CloudProvider:     openapicommonv1beta1.Region_CloudProvider(cloudProvider),
		Region:            clusterRecord.Region,
		UserId:            clusterRecord.Labels[metadata.UserIDHeaderKey],
		CreatedBy:         clusterRecord.CreatedBy,
		CreatedAt:         clusterRecord.CreatedAt,
		UpdatedAt:         clusterRecord.UpdatedAt,
		DeletedAt:         deletedAt,
		Labels:            clusterRecord.Labels,
		Annotations:       clusterRecord.Annotations,
		ServicePlan:       openapicommonv1beta1.ServicePlan(servicePlan),
		EngineType:        openapicommonv1beta1.EngineType(engineType),
		ShadowPoolId:      clusterRecord.ShadowPoolID,
		ShadowState:       entity.ShadowState(clusterRecord.ShadowState),
		RuVersion:         clusterRecord.RuVersion,
	}, nil
}

func (c clusterRepositoryImpl) SupplementOriginalProjectIdAndServicePlan(ctx context.Context, startId uint64, pageSize int) (uint64, error) {
	if pageSize == 0 {
		return 0, errors.New("pageSize required")
	}
	if pageSize > 1000 {
		return 0, errors.New("pageSize too large")
	}

	rdb := c.dbProvider.ReadDB(ctx)
	var ids []uint64
	err := rdb.Model(&schema.Cluster{}).
		Where("id > ?", startId).
		Order("id ASC").
		Limit(pageSize).
		Pluck("id", &ids).Error

	if err != nil {
		return 0, fmt.Errorf("fetch ids failed: %w", err)
	}

	if len(ids) == 0 {
		tilog.Warn(ctx, "no clusters to handle")
		return 0, nil
	}

	lastId := ids[len(ids)-1]

	wdb := c.dbProvider.WriteDB(ctx)

	err = wdb.Transaction(func(tx *gorm.DB) error {
		// Supplement original_project_id
		err = tx.Exec(`
            UPDATE clusters
            SET original_project_id = project_id
            WHERE id in ? AND original_project_id = 0 AND deleted_at IS NULL`, ids).Error
		if err != nil {
			return err
		}

		// Supplement service plan for serverless clusters
		// All Essential records already have service plans, so only the Starter needs to be added.
		err = tx.Exec(`
            UPDATE clusters
            SET service_plan = 'Starter'
            WHERE id IN ?
            AND deploy_type = 'DEPLOY_TYPE_SERVERLESS'
            AND service_plan NOT IN ('Essential', 'Starter')
            AND deleted_at IS NULL`, ids).Error
		if err != nil {
			return err
		}

		// Supplement service plan for dedicated clusters
		err = tx.Exec(`
            UPDATE clusters
            SET service_plan = 'Dedicated'
            WHERE id IN ?
            AND deploy_type = 'DEPLOY_TYPE_DEDICATED'
            AND service_plan != 'Dedicated'
            AND deleted_at IS NULL`, ids).Error
		if err != nil {
			return err
		}

		// All Premium or BYOC records already have service plans
		return nil
	})

	return lastId, err
}