//go:generate ${GOBIN}/mockgen -destination repo_mock.go -package=repo ${GO_MODULE}/pkg/service/k8supgrade/task/repo Repo
package repo

import (
	"context"
	"fmt"

	"github.com/tidbcloud/pkgv2/db"
	"gorm.io/gorm"

	entity "github.com/tidbcloud/devops/entity/k8supgrade"
	"github.com/tidbcloud/devops/entity/validation"
	models "github.com/tidbcloud/devops/models/k8supgrade"
	"github.com/tidbcloud/devops/pkg/infra/utils/container"
	"github.com/tidbcloud/devops/pkg/infra/utils/repoopt"
	"github.com/tidbcloud/devops/pkg/service/k8supgrade/task/repo/cvt"
)

type (
	PatchTask func(current *entity.Task) error
)

// GetTaskOpts controls which related data is loaded alongside the task.
type GetTaskOpts struct {
	WithK8sClusters         bool
	WithExcludedK8sClusters bool
}

type Repo interface {
	CreateTask(context.Context, *entity.Task) (*entity.Task, error)
	ListTasks(context.Context, ...repoopt.Option) ([]*entity.Task, error)
	ListTasksWithK8sClusters(context.Context, ...repoopt.Option) ([]*entity.Task, error)
	CountTasks(context.Context, ...repoopt.Option) (int64, error)
	GetTask(context.Context, uint64, bool) (*entity.Task, error)
	GetTaskWithOpts(context.Context, uint64, GetTaskOpts) (*entity.Task, error)
	PatchTask(context.Context, *entity.Task, ...PatchTask) error
	UpdateTask(context.Context, uint64, *entity.Task) error
	DeleteTask(context.Context, uint64) error

	CreateK8sCluster(context.Context, *entity.K8sCluster) (*entity.K8sCluster, error)
	CreateK8sClusterInBatch(context.Context, []*entity.K8sCluster) ([]*entity.K8sCluster, error)
	ListK8sClusters(context.Context, uint64, bool) ([]*entity.K8sCluster, error)
	GetK8sCluster(context.Context, uint64, string) (*entity.K8sCluster, error)
	UpdateK8sCluster(context.Context, uint64, string, *entity.K8sCluster) error
	// UpsertK8sClustersWithVersionGuard inserts new clusters (ID==0) and updates
	// existing ones only when their version matches the value read during sync.
	// It returns the names of clusters whose update was skipped due to a version
	// conflict (i.e. modified concurrently), so the caller can trigger a resync.
	UpsertK8sClustersWithVersionGuard(ctx context.Context, clusters []*entity.K8sCluster) (conflicted []string, err error)
	DeleteK8sClusters(context.Context, uint64) error
}

var _ Repo = &repo{}

type repo struct {
	dbProvider db.Provider
	cvter      cvt.Converter
}

func NewRepo(dbp db.Provider) Repo {
	return &repo{
		dbProvider: dbp,
		cvter:      *cvt.NewConverter(),
	}
}

func (r *repo) CreateTask(ctx context.Context, task *entity.Task) (*entity.Task, error) {
	err := validation.K8sUpgrade.TaskCreation(task)
	if err != nil {
		return nil, fmt.Errorf("invalid arguments to create: %s", err)
	}
	newTask := &entity.Task{
		ID:                  0, // TODO
		Name:                task.Name,
		CloudProvider:       task.CloudProvider,
		VersionToBeUpgraded: task.VersionToBeUpgraded,
		MaxRestartTimes:     task.MaxRestartTimes,
		Deadline:            task.Deadline,
		CreatedAt:           task.CreatedAt,
		Status:              task.Status,
		Applicant:           task.Applicant,
		ClosedAt:            nil,
		ApprovalID:          task.ApprovalID,
		SelectorType:        task.SelectorType,
		SelectorValue:       task.SelectorValue,
	}
	taskMod := r.cvter.EntityToModel().Task(newTask)

	if err := r.dbProvider.WriteDB(ctx).Session(&gorm.Session{FullSaveAssociations: true}).Save(taskMod).Error; err != nil {
		return nil, fmt.Errorf("save k8s upgrade task %s failed: %w", task.Name, err)
	}
	return r.cvter.ModelToEntity().Task(taskMod), nil
}

func (r *repo) ListTasks(ctx context.Context, opts ...repoopt.Option) ([]*entity.Task, error) {
	var tasks []*models.Task
	rdb := r.dbProvider.ReadDB(ctx).Model(&models.Task{})
	rdb = repoopt.Aggregate(opts...).Apply(rdb)
	if err := rdb.Find(&tasks).Error; err != nil {
		return nil, fmt.Errorf("list k8s upgrade tasks failed: %w", err)
	}
	return container.TransformSlice(tasks, r.cvter.ModelToEntity().Task), nil
}

func (r *repo) ListTasksWithK8sClusters(ctx context.Context, opts ...repoopt.Option) ([]*entity.Task, error) {
	var tasks []*models.Task
	rdb := r.dbProvider.ReadDB(ctx).Model(&models.Task{})
	rdb = repoopt.Aggregate(opts...).Apply(rdb)
	if err := rdb.Find(&tasks).Error; err != nil {
		return nil, fmt.Errorf("list k8s upgrade tasks with clusters failed: %w", err)
	}

	entityTasks := container.TransformSlice(tasks, r.cvter.ModelToEntity().Task)
	for _, task := range entityTasks {
		clusters, err := r.ListK8sClusters(ctx, task.ID, false)
		if err != nil {
			return nil, fmt.Errorf("list k8s clusters of the task %d failed: %w", task.ID, err)
		}
		task.K8sClusterList = clusters

		excludedClusters, err := r.ListK8sClusters(ctx, task.ID, true)
		if err != nil {
			return nil, fmt.Errorf("list excluded k8s clusters of the task %d failed: %w", task.ID, err)
		}
		task.ExcludedK8sClusters = excludedClusters
	}

	return entityTasks, nil
}

func (r *repo) CountTasks(ctx context.Context, opts ...repoopt.Option) (int64, error) {
	rdb := r.dbProvider.ReadDB(ctx).Model(&models.Task{})
	rdb = repoopt.Aggregate(opts...).Apply(rdb)

	var count int64
	if err := rdb.Count(&count).Error; err != nil {
		return 0, fmt.Errorf("list k8s upgrade tasks failed: %w", err)
	}
	return count, nil
}

func (r *repo) GetTask(ctx context.Context, id uint64, withK8sClusters bool) (*entity.Task, error) {
	ret := &models.Task{}

	if err := r.dbProvider.ReadDB(ctx).Model(&models.Task{}).
		Where("id = ?", id).
		First(ret).
		Error; err != nil {
		return nil, fmt.Errorf("load k8s upgrade task %d failed: %w", id, err)
	}

	task := r.cvter.ModelToEntity().Task(ret)
	if withK8sClusters {
		var err error
		task.K8sClusterList, err = r.ListK8sClusters(ctx, id, false)
		if err != nil {
			return nil, fmt.Errorf("list k8s clusters of the task %d failed: %w", id, err)
		}
		task.ExcludedK8sClusters, err = r.ListK8sClusters(ctx, id, true)
		if err != nil {
			return nil, fmt.Errorf("list excluded k8s clusters of the task %d failed: %w", id, err)
		}
	}
	return task, nil
}

func (r *repo) GetTaskWithOpts(ctx context.Context, id uint64, opts GetTaskOpts) (*entity.Task, error) {
	ret := &models.Task{}

	if err := r.dbProvider.ReadDB(ctx).Model(&models.Task{}).
		Where("id = ?", id).
		First(ret).
		Error; err != nil {
		return nil, fmt.Errorf("load k8s upgrade task %d failed: %w", id, err)
	}

	task := r.cvter.ModelToEntity().Task(ret)
	if opts.WithK8sClusters {
		var err error
		task.K8sClusterList, err = r.ListK8sClusters(ctx, id, false)
		if err != nil {
			return nil, fmt.Errorf("list k8s clusters of the task %d failed: %w", id, err)
		}
	}
	if opts.WithExcludedK8sClusters {
		var err error
		task.ExcludedK8sClusters, err = r.ListK8sClusters(ctx, id, true)
		if err != nil {
			return nil, fmt.Errorf("list excluded k8s clusters of the task %d failed: %w", id, err)
		}
	}
	return task, nil
}

func (r *repo) DeleteTask(ctx context.Context, tid uint64) error {
	return r.dbProvider.WriteDB(ctx).Delete(&models.Task{}, tid).Error
}

func (r *repo) DeleteK8sClusters(ctx context.Context, tid uint64) error {
	return r.dbProvider.WriteDB(ctx).Where("task_id = ?", tid).Delete(&models.K8sCluster{}).Error
}

func (r *repo) UpdateTask(ctx context.Context, taskID uint64, task *entity.Task) error {
	err := validation.K8sUpgrade.TaskUpdate(task)
	if err != nil {
		return err
	}
	taskModel := r.cvter.EntityToModel().Task(task)
	err = r.dbProvider.WriteDB(ctx).Model(&models.Task{}).Where("id = ?", taskID).Updates(taskModel).Error
	if err != nil {
		return fmt.Errorf("update k8s upgrade task %s failed: %w", task.Name, err)
	}

	return nil
}

func (r *repo) PatchTask(ctx context.Context, current *entity.Task, patch ...PatchTask) error {
	if len(patch) == 0 {
		return nil
	}

	for _, p := range patch {
		err := p(current)
		if err != nil {
			return err
		}
	}

	taskModel := r.cvter.EntityToModel().Task(current)

	err := r.dbProvider.WriteDB(ctx).Session(&gorm.Session{FullSaveAssociations: true}).Save(taskModel).Error
	if err != nil {
		return fmt.Errorf("save k8s upgrade task failed: %w", err)
	}
	return nil
}

func (r *repo) CreateK8sCluster(ctx context.Context, cluster *entity.K8sCluster) (*entity.K8sCluster, error) {
	err := validation.K8sUpgrade.K8sClusterCreation(cluster)
	if err != nil {
		return nil, fmt.Errorf("invalid arguments to create: %s", err)
	}
	clusterMod := r.cvter.EntityToModel().K8sCluster(cluster)

	if err := r.dbProvider.WriteDB(ctx).Session(&gorm.Session{FullSaveAssociations: true}).Create(clusterMod).Error; err != nil {
		return nil, fmt.Errorf("save k8s cluster %s of upgrade task %d failed: %w", cluster.Name, cluster.TaskID, err)
	}
	return r.cvter.ModelToEntity().K8sCluster(clusterMod), nil
}

func (r *repo) CreateK8sClusterInBatch(ctx context.Context, clusters []*entity.K8sCluster) ([]*entity.K8sCluster, error) {
	for _, cluster := range clusters {
		err := validation.K8sUpgrade.K8sClusterCreation(cluster)
		if err != nil {
			return nil, fmt.Errorf("invalid arguments to create: %s", err)
		}
	}

	clusterMods := r.cvtK8sClustersToModels(clusters)

	if err := r.dbProvider.WriteDB(ctx).Session(&gorm.Session{FullSaveAssociations: true}).Create(clusterMods).Error; err != nil {
		return nil, fmt.Errorf("save k8s clusters in batches failed: %w", err)
	}
	return r.cvtK8sClustersToEntities(clusterMods), nil
}

func (r *repo) ListK8sClusters(ctx context.Context, taskID uint64, isExcluded bool) ([]*entity.K8sCluster, error) {
	var clusters []*models.K8sCluster

	result := r.dbProvider.ReadDB(ctx).Where("task_id = ?", taskID).
		Where("is_excluded = ?", isExcluded).
		Find(&clusters)
	if result.Error != nil {
		return nil, fmt.Errorf("list k8s clusters failed: %w", result.Error)
	}
	// TODO container.TransformSlice(ret, r.cvter.ModelToEntity().K8sCluster)
	entities := r.cvtK8sClustersToEntities(clusters)
	return entities, nil
}

func (r *repo) GetK8sCluster(ctx context.Context, taskID uint64, k8sName string) (*entity.K8sCluster, error) {
	ret := &models.K8sCluster{}
	if err := r.dbProvider.ReadDB(ctx).Model(&models.K8sCluster{}).
		Where("task_id = ?", taskID).
		Where("name = ?", k8sName).
		First(ret).
		Error; err != nil {
		return nil, fmt.Errorf("load k8s cluster %s of task %d failed: %w", k8sName, taskID, err)
	}
	return r.cvter.ModelToEntity().K8sCluster(ret), nil
}

func (r *repo) UpdateK8sCluster(ctx context.Context, taskID uint64, name string, cluster *entity.K8sCluster) error {
	err := validation.K8sUpgrade.K8sClusterUpdate(cluster)
	if err != nil {
		return err
	}
	clusterModel := r.cvter.EntityToModel().K8sCluster(cluster)
	clusterModel.Version = cluster.Version + 1
	result := r.dbProvider.WriteDB(ctx).Model(&models.K8sCluster{}).
		Where("id = ?", cluster.ID).
		Where("task_id = ?", taskID).
		Where("name = ?", name).
		Where("version = ?", cluster.Version).
		Updates(clusterModel)
	if result.Error != nil {
		return fmt.Errorf("update k8s cluster %d of task %d failed: %w", cluster.ID, cluster.TaskID, result.Error)
	} else if result.RowsAffected == 0 {
		return fmt.Errorf("k8s cluster %s of task %d is modified concurrently, please refresh the page and try again", name, taskID)
	}
	return nil
}

func (r *repo) UpsertK8sClustersWithVersionGuard(ctx context.Context, clusters []*entity.K8sCluster) (conflicted []string, err error) {
	var toInsert []*entity.K8sCluster
	var toUpdate []*entity.K8sCluster
	for _, c := range clusters {
		if c.ID == 0 {
			toInsert = append(toInsert, c)
		} else {
			toUpdate = append(toUpdate, c)
		}
	}

	if len(toInsert) > 0 {
		if _, err = r.CreateK8sClusterInBatch(ctx, toInsert); err != nil {
			return nil, fmt.Errorf("insert new k8s clusters failed: %w", err)
		}
	}

	for _, cluster := range toUpdate {
		if err = validation.K8sUpgrade.K8sClusterUpdate(cluster); err != nil {
			return nil, fmt.Errorf("invalid arguments to update: %s", err)
		}
		clusterModel := r.cvter.EntityToModel().K8sCluster(cluster)
		// Write version+1 explicitly so the SET and WHERE are in the same statement:
		// UPDATE ... SET ..., version=<v+1> WHERE id=? AND version=<v>
		// If RowsAffected==0 the row was modified concurrently.
		clusterModel.Version = cluster.Version + 1
		result := r.dbProvider.WriteDB(ctx).Model(&models.K8sCluster{}).
			Select("*").
			Omit("id", "task_id", "name", "created_at").
			Where("id = ?", cluster.ID).
			Where("version = ?", cluster.Version).
			Updates(clusterModel)
		if result.Error != nil {
			return nil, fmt.Errorf("update k8s cluster %s with version guard failed: %w", cluster.Name, result.Error)
		}
		if result.RowsAffected == 0 {
			conflicted = append(conflicted, cluster.Name)
		}
	}

	return conflicted, nil
}

func (r *repo) cvtK8sClustersToEntities(clusters []*models.K8sCluster) []*entity.K8sCluster {
	clusterEntities := make([]*entity.K8sCluster, 0)
	for _, c := range clusters {
		clusterEntities = append(clusterEntities, r.cvter.ModelToEntity().K8sCluster(c))
	}
	return clusterEntities
}

func (r *repo) cvtK8sClustersToModels(clusters []*entity.K8sCluster) []*models.K8sCluster {
	clusterModels := make([]*models.K8sCluster, 0)
	for _, c := range clusters {
		clusterModels = append(clusterModels, r.cvter.EntityToModel().K8sCluster(c))
	}
	return clusterModels
}