// Package utils provides shared utilities for the fetch_ml project.
package metrics

import (
	"sync/atomic"
	"time"
)

func max(a, b int64) int64 {
	if a > b {
		return a
	}
	return b
}

type Metrics struct {
	TasksProcessed  atomic.Int64
	TasksFailed     atomic.Int64
	DataFetchTime   atomic.Int64 // Total nanoseconds
	ExecutionTime   atomic.Int64
	DataTransferred atomic.Int64 // Total bytes
	ActiveTasks     atomic.Int64
	QueuedTasks     atomic.Int64
}

func (m *Metrics) RecordTaskSuccess(duration time.Duration) {
	m.TasksProcessed.Add(1)
	m.ExecutionTime.Add(duration.Nanoseconds())
}

func (m *Metrics) RecordTaskFailure() {
	m.TasksFailed.Add(1)
}

func (m *Metrics) RecordTaskStart() {
	m.ActiveTasks.Add(1)
}

// RecordTaskCompletion decrements the number of active tasks. It is safe to call
// even if no tasks are currently recorded; the caller should ensure calls are
// balanced with RecordTaskStart.
func (m *Metrics) RecordTaskCompletion() {
	m.ActiveTasks.Add(-1)
}

func (m *Metrics) RecordDataTransfer(bytes int64, duration time.Duration) {
	m.DataTransferred.Add(bytes)
	m.DataFetchTime.Add(duration.Nanoseconds())
}

func (m *Metrics) SetQueuedTasks(count int64) {
	m.QueuedTasks.Store(count)
}

func (m *Metrics) GetStats() map[string]any {
	processed := m.TasksProcessed.Load()
	failed := m.TasksFailed.Load()
	dataTransferred := m.DataTransferred.Load()
	dataFetchTime := m.DataFetchTime.Load()

	return map[string]any{
		"tasks_processed":     processed,
		"tasks_failed":        failed,
		"active_tasks":        m.ActiveTasks.Load(),
		"queued_tasks":        m.QueuedTasks.Load(),
		"success_rate":        float64(processed-failed) / float64(max(processed, 1)),
		"avg_exec_time":       time.Duration(m.ExecutionTime.Load() / max(processed, 1)),
		"data_transferred_gb": float64(dataTransferred) / (1024 * 1024 * 1024),
		"avg_fetch_time":      time.Duration(dataFetchTime / max(processed, 1)),
	}
}