package vercel

import (
	"context"
	"time"

	"github.com/pingcap/errors"
	"github.com/tidbcloud/pkgv2/tracing"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"gorm.io/gorm"
)

type BranchTaskDAO interface {
	AddTask(ctx context.Context, identifier BranchTask) error
	GetTasksByStatus(ctx context.Context, status TaskStatus) ([]BranchTask, error)
	GetTasksToClean(ctx context.Context, timeBefore time.Time) ([]BranchTask, error)
	DeleteTasks(ctx context.Context, tasks []BranchTask) error
	CountTasksByStatus(ctx context.Context) ([]TaskStatusCount, error)
	UpdateTasksStatus(ctx context.Context, status TaskStatus, taskIDs []uint64) error
}

type branchTaskDAOImpl struct {
	db *gorm.DB
}

func NewBranchTaskDAO(db *gorm.DB) BranchTaskDAO {
	return &branchTaskDAOImpl{
		db: db,
	}
}

func (d *branchTaskDAOImpl) AddTask(ctx context.Context, task BranchTask) (err error) {
	ctx, span := Tracer.Start(ctx, "branchTaskDAOImpl.AddTask")
	defer func() {
		tracing.EndSpan(span, err)
	}()

	err = d.db.Create(&task).Error
	if err != nil {
		return errors.Trace(err)
	}

	tilog.Info(ctx, "Succeed to add branch task")
	return nil
}

func (d *branchTaskDAOImpl) GetTasksByStatus(ctx context.Context, status TaskStatus) (_ []BranchTask, err error) {
	_, span := Tracer.Start(ctx, "branchTaskDAOImpl.GetTasksByStatus")
	defer func() {
		tracing.EndSpan(span, err)
	}()

	var tasks []BranchTask
	if result := d.db.Where("status = ?", string(status)).Find(&tasks); result.Error != nil {
		return nil, errors.Trace(result.Error)
	}
	return tasks, nil
}

func (d *branchTaskDAOImpl) GetTasksToClean(ctx context.Context, timeBefore time.Time) (_ []BranchTask, err error) {
	_, span := Tracer.Start(ctx, "branchTaskDAOImpl.GetTasksToClean")
	defer func() {
		tracing.EndSpan(span, err)
	}()

	var tasks []BranchTask
	if result := d.db.Where("created_at < ?", timeBefore).Find(&tasks); result.Error != nil {
		return nil, errors.Trace(result.Error)
	}
	return tasks, nil
}

func (d *branchTaskDAOImpl) DeleteTasks(ctx context.Context, tasks []BranchTask) (err error) {
	_, span := Tracer.Start(ctx, "branchTaskDAOImpl.DeleteTasks")
	defer func() {
		tracing.EndSpan(span, err)
	}()

	if len(tasks) == 0 {
		return nil
	}

	if result := d.db.Delete(&tasks); result.Error != nil {
		return errors.Trace(result.Error)
	}
	return nil
}

func (d *branchTaskDAOImpl) CountTasksByStatus(ctx context.Context) ([]TaskStatusCount, error) {
	_, span := Tracer.Start(ctx, "branchTaskDAOImpl.CountTasksByStatus")
	defer func() {
		tracing.EndSpan(span, nil)
	}()

	var counts []TaskStatusCount
	if result := d.db.Model(&BranchTask{}).Select("status, count(*) as count").Group("status").Scan(&counts); result.Error != nil {
		return nil, errors.Trace(result.Error)
	}
	return counts, nil
}

func (d *branchTaskDAOImpl) UpdateTasksStatus(ctx context.Context, status TaskStatus, taskIDs []uint64) (err error) {
	_, span := Tracer.Start(ctx, "branchTaskDAOImpl.UpdateTasksStatus")
	defer func() {
		tracing.EndSpan(span, err)
	}()

	if len(taskIDs) == 0 {
		return nil
	}
	if result := d.db.Model(&BranchTask{}).Where("id in ?", taskIDs).Update("status", string(status)); result.Error != nil {
		return errors.Trace(result.Error)
	}
	return nil
}