package core

import (
	"context"
	"fmt"
	"math"
	"sync"
	"time"

	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/pingcap/schrodinger/buildbot/config"
	ts "github.com/pingcap/schrodinger/buildbot/task"
	"github.com/pingcap/schrodinger/pkg/storage"
	"github.com/pingcap/schrodinger/pkg/types"
	"github.com/pingcap/schrodinger/util"
	"golang.org/x/sync/errgroup"
)

const (
	fetchInterval = 5 * time.Second
)

// TaskRunner is an interface that provides the ability to run BuildTask
type TaskRunner interface {
	// Run runs task.
	Run(ctx context.Context) error
	// Clean cleans up residual data after the task is finished.
	Clean(ctx context.Context) error
}

// BuildManager manages all build tasks
type BuildManager struct {
	store        *storage.Storage
	cfg          config.BaseConfig
	pendingTasks chan *types.BuildTask
	runningTasks sync.Map

	sync.Mutex
}

// NewBuildManager returns an instance of BuildManger
func NewBuildManager(store *storage.Storage, cfg config.BaseConfig) *BuildManager {
	mgr := &BuildManager{
		store:        store,
		cfg:          cfg,
		pendingTasks: make(chan *types.BuildTask, 10000),
	}

	return mgr
}

// Recover recovers the pending and running tasks from store
func (b *BuildManager) Recover() error {
	if err := b.getTasksFromStore(math.MaxInt64, types.TaskRunning, types.TaskPending); err != nil {
		log.Errorf("get build tasks from storage failed, error: %v", err)
	}
	return nil
}

// Run starts run the buildManager serve.
func (b *BuildManager) Run(ctx context.Context) error {
	g := errgroup.Group{}
	// fetch tasks from storage
	g.Go(func() error {
		b.produceTasks(ctx)
		return nil
	})

	// run build tasks
	g.Go(func() error {
		return b.consumeTasks(ctx)
	})

	return g.Wait()
}

func (b *BuildManager) produceTasks(ctx context.Context) {
	ticker := time.NewTicker(fetchInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if err := b.getTasksFromStore(20, types.TaskNew); err != nil {
				log.Errorf("get build tasks from storage failed, error: %v", err)
			}
		}
	}
}

func (b *BuildManager) getTasksFromStore(limit int64, status ...types.TaskStatus) error {
	newTasks, err := b.store.ListBuildTask(storage.BuildTaskConditions{
		Status: status,
		Limit:  limit,
	})

	if err != nil {
		return errors.Trace(err)
	}

	for _, task := range newTasks {
		task.Status = types.TaskPending
		if err := b.store.UpdateBuildTask(task); err != nil {
			log.Errorf("update task %+v failed, error: %v", task, err)
			continue
		}

		b.pendingTasks <- task
	}

	return nil
}

func (b *BuildManager) consumeTasks(ctx context.Context) error {
	g := errgroup.Group{}
	for i := 0; i < b.cfg.Concurrency; i++ {
		g.Go(func() error {
			b.worker(ctx)
			return nil
		})
	}

	return g.Wait()
}

func (b *BuildManager) worker(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case task := <-b.pendingTasks:
			if err := b.runBuildTask(ctx, task); err != nil {
				log.Errorf("[task %d] run task failed, error: %v", task.ID,
					errors.ErrorStack(err))
				continue
			}
			log.Infof("[task %d] run task successfully", task.ID)
		}
	}
}

func (b *BuildManager) runBuildTask(ctx context.Context, task *types.BuildTask) error {
	log.Infof("[task %d] start to prepare build task %+v", task.ID, task)
	if err := b.prepareTask(task); err != nil {
		return errors.Trace(err)
	}

	var (
		runner TaskRunner
		key    string
	)

	switch task.Information.Repo {
	case "tidb", "tikv", "pd":
		key = task.Information.Image
		runner = ts.NewTiDBTask(task, b.store, b.cfg)
	default:
		return errors.NotFoundf("%s", task.Information.Repo)
	}

	_, exist := b.checkDuplicateTask(key, runner)
	if exist {
		log.Infof("[task %d] the build task is duplicate, put it to pending queue", task.ID)
		go b.handleDuplicateTask(ctx, task)
		return nil
	}

	log.Infof("[task %d] start to run build task %+v", task.ID, task)
	if err := runner.Run(ctx); err != nil {
		b.runningTasks.Delete(key)
		return errors.Trace(err)
	}

	if err := b.cleanTask(ctx, runner, key); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (b *BuildManager) checkDuplicateTask(key string, runner TaskRunner) (TaskRunner, bool) {
	b.Lock()
	defer b.Unlock()

	runningTask, exist := b.runningTasks.Load(key)
	if exist {
		return runningTask.(TaskRunner), true
	}

	b.runningTasks.Store(key, runner)
	return nil, false
}

func (b *BuildManager) handleDuplicateTask(ctx context.Context, task *types.BuildTask) {
	// sleep 10 seconds to prevent the case from being dispatched immediately
	util.Sleep(ctx, 10*time.Second)

	task.Status = types.TaskPending
	if err := b.store.UpdateBuildTask(task); err != nil {
		log.Errorf("[task %d] update task %+v failed, error: %v", task.ID, task, err)
	}

	b.pendingTasks <- task
}

func (b *BuildManager) cleanTask(ctx context.Context, runner TaskRunner, key string) error {
	b.Lock()
	defer b.Unlock()

	b.runningTasks.Delete(key)

	if err := runner.Clean(ctx); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (b *BuildManager) prepareTask(task *types.BuildTask) error {
	var (
		info string
		err  error
	)

	defer func() {
		if err != nil {
			task.Status = types.TaskFailed
		}

		task.Information.BuildLog = info
		if err := b.store.UpdateBuildTask(task); err != nil {
			log.Errorf("[task %d] update task %+v failed, error: %v", task.ID, task, err)
		}
	}()

	if err = task.Valid(); err != nil {
		info = fmt.Sprintf("check task failed, %v", err)
		return errors.Trace(err)
	}

	if err = task.Init(b.cfg.DownloadAddr, b.cfg.DockerRegistryAddr); err != nil {
		info = fmt.Sprintf("init task compiler info failed, error: %v", err)
		return errors.Trace(err)
	}

	task.Status = types.TaskRunning

	return nil
}
