package grpc

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"
	"time"
)

type importAsyncTaskStatus string

const (
	importAsyncTaskStatusRunning   importAsyncTaskStatus = "running"
	importAsyncTaskStatusSucceeded importAsyncTaskStatus = "succeeded"
	importAsyncTaskStatusFailed    importAsyncTaskStatus = "failed"
)

type importAsyncTask struct {
	ID        string
	Status    importAsyncTaskStatus
	Result    any
	Error     error
	CreatedAt time.Time
	UpdatedAt time.Time
}

type importAsyncTaskManager struct {
	nextID        atomic.Uint64
	taskRetention time.Duration
	mu            sync.RWMutex
	tasks         map[string]*importAsyncTask
}

func newImportAsyncTaskManager() *importAsyncTaskManager {
	return newImportAsyncTaskManagerWithRetention(10 * time.Minute)
}

func newImportAsyncTaskManagerWithRetention(taskRetention time.Duration) *importAsyncTaskManager {
	return &importAsyncTaskManager{
		taskRetention: taskRetention,
		tasks:         make(map[string]*importAsyncTask),
	}
}

func (m *importAsyncTaskManager) Create(ctx context.Context, run func(context.Context) (any, error)) string {
	id := fmt.Sprintf("import-task-%d", m.nextID.Add(1))
	m.createWithID(ctx, id, run)
	return id
}

func (m *importAsyncTaskManager) CreateIfAbsent(ctx context.Context, id string, run func(context.Context) (any, error)) (string, bool) {
	m.mu.Lock()
	m.pruneExpiredTasksLocked(time.Now())
	if _, ok := m.tasks[id]; ok {
		m.mu.Unlock()
		return id, false
	}
	task := &importAsyncTask{
		ID:        id,
		Status:    importAsyncTaskStatusRunning,
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
	}
	m.tasks[id] = task
	m.mu.Unlock()

	go m.execute(context.WithoutCancel(ctx), id, run)
	return id, true
}

func (m *importAsyncTaskManager) createWithID(ctx context.Context, id string, run func(context.Context) (any, error)) {
	m.CreateIfAbsent(ctx, id, run)
}

func (m *importAsyncTaskManager) execute(ctx context.Context, id string, run func(context.Context) (any, error)) {
	result, err := run(ctx)

	m.mu.Lock()
	defer m.mu.Unlock()

	task, ok := m.tasks[id]
	if !ok {
		return
	}

	task.UpdatedAt = time.Now()
	if err != nil {
		task.Status = importAsyncTaskStatusFailed
		task.Error = err
		task.Result = nil
		return
	}

	task.Status = importAsyncTaskStatusSucceeded
	task.Result = result
	task.Error = nil
}

func (m *importAsyncTaskManager) Get(id string) (*importAsyncTask, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.pruneExpiredTasksLocked(time.Now())

	task, ok := m.tasks[id]
	if !ok {
		return nil, false
	}
	cp := *task
	return &cp, true
}

func (m *importAsyncTaskManager) pruneExpiredTasksLocked(now time.Time) {
	if m.taskRetention <= 0 {
		return
	}
	for id, task := range m.tasks {
		if task.Status == importAsyncTaskStatusRunning {
			continue
		}
		if now.Sub(task.UpdatedAt) <= m.taskRetention {
			continue
		}
		delete(m.tasks, id)
	}
}