package services

import (
	"sync"
	"time"
)

type TaskStore struct {
	mu    sync.RWMutex
	tasks map[string]*Task
}

func NewTaskStore() *TaskStore {
	return &TaskStore{tasks: map[string]*Task{}}
}

func (s *TaskStore) Put(t *Task) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.tasks[t.ID] = t
}

func (s *TaskStore) Get(id string) (*Task, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	t, ok := s.tasks[id]
	return t, ok
}

func (s *TaskStore) Delete(id string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.tasks, id)
}

func (s *TaskStore) Cleanup(retention time.Duration) {
	now := time.Now()
	var expired []*Task

	s.mu.RLock()
	for _, t := range s.tasks {
		if t.Status == TaskCompleted || t.Status == TaskCanceled || t.Status == TaskFailed {
			if !t.EndedAt.IsZero() && now.Sub(t.EndedAt) > retention {
				expired = append(expired, t)
			}
		}
	}
	s.mu.RUnlock()

	for _, t := range expired {
		t.Cleanup()
		s.mu.Lock()
		if cur, ok := s.tasks[t.ID]; ok && cur == t {
			delete(s.tasks, t.ID)
		}
		s.mu.Unlock()
	}
}