package flash

import (
	"visiona-local/server/internal/driver"
	"sync"
)

type FlashTask struct {
	ID         string
	DeviceID   string
	ModelID    string
	ProgressCh chan driver.FlashProgress
	Done       bool
}

type ProgressTracker struct {
	tasks map[string]*FlashTask
	mu    sync.RWMutex
}

func NewProgressTracker() *ProgressTracker {
	return &ProgressTracker{
		tasks: make(map[string]*FlashTask),
	}
}

func (pt *ProgressTracker) Create(taskID, deviceID, modelID string) *FlashTask {
	pt.mu.Lock()
	defer pt.mu.Unlock()
	task := &FlashTask{
		ID:         taskID,
		DeviceID:   deviceID,
		ModelID:    modelID,
		ProgressCh: make(chan driver.FlashProgress, 20),
	}
	pt.tasks[taskID] = task
	return task
}

func (pt *ProgressTracker) Get(taskID string) (*FlashTask, bool) {
	pt.mu.RLock()
	defer pt.mu.RUnlock()
	t, ok := pt.tasks[taskID]
	return t, ok
}

func (pt *ProgressTracker) Remove(taskID string) {
	pt.mu.Lock()
	defer pt.mu.Unlock()
	delete(pt.tasks, taskID)
}