package premium_dm

import (
	"context"

	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/db"
	"gorm.io/gorm"
)

type PremiumDMTaskDAO struct {
	dbProvider db.Provider
}

func NewPremiumDMTaskDAO(dbProvider db.Provider) *PremiumDMTaskDAO {
	return &PremiumDMTaskDAO{dbProvider: dbProvider}
}

func buildDMTaskFilter(
	filter *entity.PremiumDMTaskFilter,
	db *gorm.DB,
	applyPagination bool,
) (*gorm.DB, bizerr.BizError) {
	if filter == nil {
		return nil, bizerr.New(bizerr.ErrCodeDB, "filter is required")
	}
	if filter.ClusterID != "" {
		db = db.Where("tidb_id = ?", filter.ClusterID)
	}
	if filter.TenantID != "" {
		db = db.Where("tenant_id = ?", filter.TenantID)
	}
	if len(filter.Stages) > 0 {
		db = db.Where("stage in ?", filter.Stages)
	}
	if filter.OrderBy != "" {
		db = db.Order(filter.OrderBy)
	}
	if applyPagination {
		if filter.Offset != nil {
			db = db.Offset(*filter.Offset)
		}
		if filter.Limit != nil {
			db = db.Limit(*filter.Limit)
		}
	}
	return db, nil
}

func (r *PremiumDMTaskDAO) ListTasks(
	ctx context.Context,
	filter *entity.PremiumDMTaskFilter,
) (tasks []*entity.PremiumDMTask, err bizerr.BizError) {
	rdb := r.dbProvider.ReadDB(ctx)
	rdb, err = buildDMTaskFilter(filter, rdb, true)
	if err != nil {
		return nil, err
	}
	if err := rdb.Model(&entity.PremiumDMTask{}).Preload("DMSources").Find(&tasks).Error; err != nil {
		return nil, bizerr.WrapGormError(err, "failed to list DMTasks")
	}
	return tasks, nil
}

func (r *PremiumDMTaskDAO) GetTask(
	ctx context.Context,
	clusterID,
	taskID string,
	stages []entity.PremiumDMTaskStage,
) (task *entity.PremiumDMTask, err bizerr.BizError) {
	rdb := r.dbProvider.ReadDB(ctx)
	task = &entity.PremiumDMTask{}
	query := rdb.Model(&entity.PremiumDMTask{}).Preload("DMSources").Where("id = ? AND tidb_id = ?", taskID, clusterID)
	if len(stages) > 0 {
		query = query.Where("stage IN ?", stages)
	}
	if err := query.First(task).Error; err != nil {
		return nil, bizerr.WrapGormError(err, "failed to get DMTask, id: %s", taskID)
	}
	return task, nil
}

func (r *PremiumDMTaskDAO) CreateTask(ctx context.Context, task *entity.PremiumDMTask) (err bizerr.BizError) {
	wdb := r.dbProvider.WriteDB(ctx)
	if err := wdb.Create(task).Error; err != nil {
		return bizerr.WrapGormError(err, "failed to create DMTask")
	}
	return nil
}

func (r *PremiumDMTaskDAO) CountTasks(
	ctx context.Context,
	filter *entity.PremiumDMTaskFilter,
) (_ int64, bizErr bizerr.BizError) {
	rdb := r.dbProvider.ReadDB(ctx)
	// build filter WITHOUT applying pagination (offset/limit) so Count returns total matching rows
	rdb, bizErr = buildDMTaskFilter(filter, rdb, false)
	if bizErr != nil {
		return 0, bizErr
	}
	var count int64
	if err := rdb.Model(&entity.PremiumDMTask{}).Count(&count).Error; err != nil {
		return 0, bizerr.WrapGormError(err, "failed to count DMTasks")
	}
	return count, nil
}

// UpdateTask updates specified columns of a DMTask by id.
func (r *PremiumDMTaskDAO) UpdateTask(
	ctx context.Context,
	id string,
	stages []entity.PremiumDMTaskStage,
	updates map[string]any,
) (err bizerr.BizError) {
	_, err = r.UpdateTaskWithRowsAffected(ctx, id, stages, updates)
	return err
}

// UpdateTaskWithRowsAffected updates specified columns and returns the number of rows
// that matched the task id and expected stages.
func (r *PremiumDMTaskDAO) UpdateTaskWithRowsAffected(
	ctx context.Context,
	id string,
	stages []entity.PremiumDMTaskStage,
	updates map[string]any,
) (rowsAffected int64, err bizerr.BizError) {
	wdb := r.dbProvider.WriteDB(ctx)
	query := wdb.Model(&entity.PremiumDMTask{}).
		Where("id = ?", id)
	if len(stages) > 0 {
		query = query.Where("stage IN ?", stages)
	}
	result := query.Updates(updates)
	if result.Error != nil {
		return 0, bizerr.WrapGormError(result.Error, "failed to update DMTask, id: %s", id)
	}
	return result.RowsAffected, nil
}

func (r *PremiumDMTaskDAO) UpdateSource(ctx context.Context, id string, updates map[string]any) (err bizerr.BizError) {
	wdb := r.dbProvider.WriteDB(ctx)
	result := wdb.Model(&entity.PremiumDMSource{}).
		Where("id = ? AND stage <> ?", id, entity.PremiumDMSourceStageDeleted).
		Updates(updates)
	if result.Error != nil {
		return bizerr.WrapGormError(result.Error, "failed to update DMSource, id: %s", id)
	}
	return nil
}

func (r *PremiumDMTaskDAO) CreatePrecheck(
	ctx context.Context,
	precheck *entity.PremiumDMPrecheck,
) (err bizerr.BizError) {
	if precheck == nil {
		return bizerr.New(bizerr.ErrCodeParams, "precheck is nil")
	}
	wdb := r.dbProvider.WriteDB(ctx)
	if err := wdb.Create(precheck).Error; err != nil {
		return bizerr.WrapGormError(err, "failed to create DMPrecheck")
	}
	return nil
}

func buildDMPrecheckFilter(filter *entity.PremiumDMPrecheckFilter, db *gorm.DB) (*gorm.DB, bizerr.BizError) {
	if filter == nil {
		return nil, bizerr.New(bizerr.ErrCodeDB, "filter is required")
	}
	db = db.Where("canceled_time IS NULL")
	if filter.ClusterID != "" {
		db = db.Where("tidb_id = ?", filter.ClusterID)
	}
	if filter.TenantID != "" {
		db = db.Where("tenant_id = ?", filter.TenantID)
	}
	if len(filter.Stages) > 0 {
		db = db.Where("stage in ?", filter.Stages)
	}
	if filter.OrderBy != "" {
		db = db.Order(filter.OrderBy)
	}
	if filter.Offset != nil {
		db = db.Offset(*filter.Offset)
	}
	if filter.Limit != nil {
		db = db.Limit(*filter.Limit)
	}
	return db, nil
}

func (r *PremiumDMTaskDAO) ListPrechecks(
	ctx context.Context,
	filter *entity.PremiumDMPrecheckFilter,
) (prechecks []*entity.PremiumDMPrecheck, err bizerr.BizError) {
	rdb := r.dbProvider.ReadDB(ctx)
	rdb, err = buildDMPrecheckFilter(filter, rdb)
	if err != nil {
		return nil, err
	}
	if err := rdb.Model(&entity.PremiumDMPrecheck{}).Preload("DMPrecheckItems").Find(&prechecks).Error; err != nil {
		return nil, bizerr.WrapGormError(err, "failed to list DMPrechecks")
	}
	return prechecks, nil
}

// CountPrechecks counts precheck records matching the given filter, ignoring
// any Limit/Offset on the filter. It is used to populate the total_size of
// paged ListPrechecks responses.
func (r *PremiumDMTaskDAO) CountPrechecks(
	ctx context.Context,
	filter *entity.PremiumDMPrecheckFilter,
) (total uint32, err bizerr.BizError) {
	countFilter := *filter
	countFilter.Limit = nil
	countFilter.Offset = nil
	countFilter.OrderBy = ""
	rdb := r.dbProvider.ReadDB(ctx)
	rdb, err = buildDMPrecheckFilter(&countFilter, rdb)
	if err != nil {
		return 0, err
	}
	var count int64
	if err := rdb.Model(&entity.PremiumDMPrecheck{}).Count(&count).Error; err != nil {
		return 0, bizerr.WrapGormError(err, "failed to count DMPrechecks")
	}
	return uint32(count), nil
}

func (r *PremiumDMTaskDAO) GetPrecheck(
	ctx context.Context,
	clusterID, id string,
) (precheck *entity.PremiumDMPrecheck, err bizerr.BizError) {
	rdb := r.dbProvider.ReadDB(ctx)
	precheck = &entity.PremiumDMPrecheck{}
	if err := rdb.Model(&entity.PremiumDMPrecheck{}).
		Preload("DMPrecheckItems").
		Where("id = ? AND tidb_id = ?", id, clusterID).
		First(precheck).Error; err != nil {
		return nil, bizerr.WrapGormError(err, "failed to get DMPrecheck, id: %s", id)
	}
	return precheck, nil
}

func (r *PremiumDMTaskDAO) UpdatePrecheck(
	ctx context.Context,
	id string,
	updates map[string]any,
) (err bizerr.BizError) {
	wdb := r.dbProvider.WriteDB(ctx)
	excluded := []entity.PremiumDMPrecheckStage{entity.PremiumDMPrecheckStageFailed, entity.PremiumDMPrecheckStageCanceled}
	result := wdb.Model(&entity.PremiumDMPrecheck{}).Where("id = ? AND stage NOT IN ?", id, excluded).Updates(updates)
	if result.Error != nil {
		return bizerr.WrapGormError(result.Error, "failed to update DMPrecheck, id: %s", id)
	}
	return nil
}

func (r *PremiumDMTaskDAO) CreatePrecheckItem(
	ctx context.Context,
	item *entity.PremiumDMPrecheckItem,
) (bizErr bizerr.BizError) {
	wdb := r.dbProvider.WriteDB(ctx)
	if err := wdb.Create(item).Error; err != nil {
		return bizerr.WrapGormError(err, "failed to add item to DMPrecheck")
	}
	return nil
}

// UpdatePrecheckItem updates specified columns of a DMPrecheckItem by id.
func (r *PremiumDMTaskDAO) UpdatePrecheckItem(
	ctx context.Context,
	id string,
	updates map[string]any,
) (err bizerr.BizError) {
	wdb := r.dbProvider.WriteDB(ctx)
	excluded := []entity.PremiumDMPrecheckStage{entity.PremiumDMPrecheckStageFailed, entity.PremiumDMPrecheckStageCanceled}
	result := wdb.Model(&entity.PremiumDMPrecheckItem{}).Where("id = ? AND stage NOT IN ?", id, excluded).Updates(updates)
	if result.Error != nil {
		return bizerr.WrapGormError(result.Error, "failed to update DMPrecheckItem, id: %s", id)
	}
	return nil
}