package sectorstorage

import (
	"context"
	"math/rand"
	"sync"
	"time"

	xerrors "golang.org/x/xerrors"

	"github.com/filecoin-project/go-state-types/abi"

	sealtasks "github.com/filecoin-project/lotus/extern/sector-storage/sealtasks"
	"github.com/filecoin-project/lotus/extern/sector-storage/stores"
	"github.com/filecoin-project/lotus/extern/sector-storage/storiface"
)

type poStScheduler struct {
	lk      sync.RWMutex
	workers map[storiface.WorkerID]*workerHandle
	cond    *sync.Cond

	postType sealtasks.TaskType
}

func newPoStScheduler(t sealtasks.TaskType) *poStScheduler {
	ps := &poStScheduler{
		workers:  map[storiface.WorkerID]*workerHandle{},
		postType: t,
	}
	ps.cond = sync.NewCond(&ps.lk)
	return ps
}

func (ps *poStScheduler) MaybeAddWorker(wid storiface.WorkerID, tasks map[sealtasks.TaskType]struct{}, w *workerHandle) bool {
	if _, ok := tasks[ps.postType]; !ok {
		return false
	}

	ps.lk.Lock()
	defer ps.lk.Unlock()

	ps.workers[wid] = w

	go ps.watch(wid, w)

	ps.cond.Broadcast()

	return true
}

func (ps *poStScheduler) delWorker(wid storiface.WorkerID) *workerHandle {
	ps.lk.Lock()
	defer ps.lk.Unlock()
	var w *workerHandle = nil
	if wh, ok := ps.workers[wid]; ok {
		w = wh
		delete(ps.workers, wid)
	}
	return w
}

func (ps *poStScheduler) CanSched(ctx context.Context) bool {
	ps.lk.RLock()
	defer ps.lk.RUnlock()
	if len(ps.workers) == 0 {
		return false
	}

	for _, w := range ps.workers {
		if w.enabled {
			return true
		}
	}

	return false
}

func (ps *poStScheduler) Schedule(ctx context.Context, primary bool, spt abi.RegisteredSealProof, work WorkerAction) error {
	ps.lk.Lock()
	defer ps.lk.Unlock()

	if len(ps.workers) == 0 {
		return xerrors.Errorf("can't find %s post worker", ps.postType)
	}

	// Get workers by resource
	canDo, candidates := ps.readyWorkers(spt)
	for !canDo {
		//if primary is true, it must be dispatched to a worker
		if primary {
			ps.cond.Wait()
			canDo, candidates = ps.readyWorkers(spt)
		} else {
			return xerrors.Errorf("can't find %s post worker", ps.postType)
		}
	}

	defer func() {
		if ps.cond != nil {
			ps.cond.Broadcast()
		}
	}()

	selected := candidates[0]
	worker := ps.workers[selected.id]

	return worker.active.withResources(selected.id, worker.info, selected.res, &ps.lk, func() error {
		ps.lk.Unlock()
		defer ps.lk.Lock()

		return work(ctx, worker.workerRpc)
	})
}

type candidateWorker struct {
	id  storiface.WorkerID
	res storiface.Resources
}

func (ps *poStScheduler) readyWorkers(spt abi.RegisteredSealProof) (bool, []candidateWorker) {
	var accepts []candidateWorker
	//if the gpus of the worker are insufficient or it's disabled, it cannot be scheduled
	for wid, wr := range ps.workers {
		needRes := wr.info.Resources.ResourceSpec(spt, ps.postType)

		if !wr.active.canHandleRequest(needRes, wid, "post-readyWorkers", wr.info) {
			continue
		}

		accepts = append(accepts, candidateWorker{
			id:  wid,
			res: needRes,
		})
	}

	// todo: round robin or something
	rand.Shuffle(len(accepts), func(i, j int) {
		accepts[i], accepts[j] = accepts[j], accepts[i]
	})

	return len(accepts) != 0, accepts
}

func (ps *poStScheduler) disable(wid storiface.WorkerID) {
	ps.lk.Lock()
	defer ps.lk.Unlock()
	ps.workers[wid].enabled = false
}

func (ps *poStScheduler) enable(wid storiface.WorkerID) {
	ps.lk.Lock()
	defer ps.lk.Unlock()
	ps.workers[wid].enabled = true
}

func (ps *poStScheduler) watch(wid storiface.WorkerID, worker *workerHandle) {
	heartbeatTimer := time.NewTicker(stores.HeartbeatInterval)
	defer heartbeatTimer.Stop()

	ctx, cancel := context.WithCancel(context.TODO())
	defer cancel()

	defer close(worker.closedMgr)

	defer func() {
		log.Warnw("Worker closing", "WorkerID", wid)
		ps.delWorker(wid)
	}()

	for {
		sctx, scancel := context.WithTimeout(ctx, stores.HeartbeatInterval/2)
		curSes, err := worker.workerRpc.Session(sctx)
		scancel()
		if err != nil {
			// Likely temporary error
			log.Warnw("failed to check worker session", "error", err)
			ps.disable(wid)

			select {
			case <-heartbeatTimer.C:
				continue
			case <-worker.closingMgr:
				return
			}
		}

		if storiface.WorkerID(curSes) != wid {
			if curSes != ClosedWorkerID {
				// worker restarted
				log.Warnw("worker session changed (worker restarted?)", "initial", wid, "current", curSes)
			}
			return
		}

		ps.enable(wid)
	}
}

func (ps *poStScheduler) workerCleanup(wid storiface.WorkerID, w *workerHandle) {
	select {
	case <-w.closingMgr:
	default:
		close(w.closingMgr)
	}

	ps.lk.Unlock()
	select {
	case <-w.closedMgr:
	case <-time.After(time.Second):
		log.Errorf("timeout closing worker manager goroutine %s", wid)
	}
	ps.lk.Lock()
}

func (ps *poStScheduler) schedClose() {
	ps.lk.Lock()
	defer ps.lk.Unlock()
	log.Debugf("closing scheduler")

	for i, w := range ps.workers {
		ps.workerCleanup(i, w)
	}
}

func (ps *poStScheduler) WorkerStats(ctx context.Context, cb func(ctx context.Context, wid storiface.WorkerID, worker *workerHandle)) {
	ps.lk.RLock()
	defer ps.lk.RUnlock()
	for id, w := range ps.workers {
		cb(ctx, id, w)
	}
}