poimen-workflows/internal/cost/cost_tracker.go

package cost

import (
	"fmt"
	"sync"
	"time"
)

// CostEntry represents a tracked cost
type CostEntry struct {
	ID        string
	Type      string // llm, git, compute
	WorkflowID string
	TaskID    string
	Amount    float64
	Timestamp time.Time
	Metadata  map[string]interface{}
}

// CostTracker tracks and analyzes workflow costs
type CostTracker struct {
	mu       sync.RWMutex
	entries  []*CostEntry
	rates    map[string]float64
}

// NewCostTracker creates a new cost tracker
func NewCostTracker() *CostTracker {
	return &CostTracker{
		entries: make([]*CostEntry, 0),
		rates: map[string]float64{
			"llm_token":     0.0001,  // $0.0001 per token
			"git_push":      0.0,     // Free
			"compute_hour":  0.5,     // $0.5 per hour
		},
	}
}

// TrackLLMCost tracks LLM API costs
func (ct *CostTracker) TrackLLMCost(workflowID, taskID string, tokens int) {
	cost := float64(tokens) * ct.rates["llm_token"]
	ct.mu.Lock()
	defer ct.mu.Unlock()

	entry := &CostEntry{
		ID:         fmt.Sprintf("llm-%d", len(ct.entries)),
		Type:       "llm",
		WorkflowID: workflowID,
		TaskID:     taskID,
		Amount:     cost,
		Timestamp:  time.Now(),
		Metadata: map[string]interface{}{
			"tokens": tokens,
		},
	}

	ct.entries = append(ct.entries, entry)
}

// TrackGitCost tracks git operation costs
func (ct *CostTracker) TrackGitCost(workflowID string, operations int) {
	ct.mu.Lock()
	defer ct.mu.Unlock()

	entry := &CostEntry{
		ID:         fmt.Sprintf("git-%d", len(ct.entries)),
		Type:       "git",
		WorkflowID: workflowID,
		Amount:     0,
		Timestamp:  time.Now(),
		Metadata: map[string]interface{}{
			"operations": operations,
		},
	}

	ct.entries = append(ct.entries, entry)
}

// TrackComputeCost tracks compute resource costs (duration in milliseconds)
func (ct *CostTracker) TrackComputeCost(workflowID, taskID string, durationMs float64) {
	// Convert milliseconds to hours
	durationHours := durationMs / (1000.0 * 3600.0)
	cost := durationHours * ct.rates["compute_hour"]
	ct.mu.Lock()
	defer ct.mu.Unlock()

	entry := &CostEntry{
		ID:         fmt.Sprintf("compute-%d", len(ct.entries)),
		Type:       "compute",
		WorkflowID: workflowID,
		TaskID:     taskID,
		Amount:     cost,
		Timestamp:  time.Now(),
		Metadata: map[string]interface{}{
			"duration_ms": durationMs,
		},
	}

	ct.entries = append(ct.entries, entry)
}

// GetTotalCost returns total cost for all workflows
func (ct *CostTracker) GetTotalCost() float64 {
	ct.mu.RLock()
	defer ct.mu.RUnlock()

	total := 0.0
	for _, entry := range ct.entries {
		total += entry.Amount
	}

	return total
}

// GetWorkflowCost returns total cost for a specific workflow
func (ct *CostTracker) GetWorkflowCost(workflowID string) float64 {
	ct.mu.RLock()
	defer ct.mu.RUnlock()

	total := 0.0
	for _, entry := range ct.entries {
		if entry.WorkflowID == workflowID {
			total += entry.Amount
		}
	}

	return total
}

// GetCostByType returns total cost by type
func (ct *CostTracker) GetCostByType(costType string) float64 {
	ct.mu.RLock()
	defer ct.mu.RUnlock()

	total := 0.0
	for _, entry := range ct.entries {
		if entry.Type == costType {
			total += entry.Amount
		}
	}

	return total
}

// GetAverageCostPerTask returns average cost per task
func (ct *CostTracker) GetAverageCostPerTask(workflowID string) float64 {
	ct.mu.RLock()
	defer ct.mu.RUnlock()

	total := 0.0
	count := 0

	for _, entry := range ct.entries {
		if entry.WorkflowID == workflowID {
			total += entry.Amount
			count++
		}
	}

	if count == 0 {
		return 0
	}

	return total / float64(count)
}

// GetOptimizationSuggestions returns cost optimization recommendations
func (ct *CostTracker) GetOptimizationSuggestions(workflowID string) []string {
	suggestions := make([]string, 0)

	ct.mu.RLock()
	defer ct.mu.RUnlock()

	llmCost := 0.0
	computeCost := 0.0

	for _, entry := range ct.entries {
		if entry.WorkflowID == workflowID {
			if entry.Type == "llm" {
				llmCost += entry.Amount
			} else if entry.Type == "compute" {
				computeCost += entry.Amount
			}
		}
	}

	if llmCost > computeCost*2 {
		suggestions = append(suggestions, "Consider caching LLM results to reduce API calls")
	}

	if computeCost > llmCost*2 {
		suggestions = append(suggestions, "Consider parallelizing compute tasks")
	}

	return suggestions
}

// GetEntries returns all cost entries
func (ct *CostTracker) GetEntries() []*CostEntry {
	ct.mu.RLock()
	defer ct.mu.RUnlock()

	result := make([]*CostEntry, len(ct.entries))
	copy(result, ct.entries)

	return result
}

// GetEntriesForWorkflow returns cost entries for a workflow
func (ct *CostTracker) GetEntriesForWorkflow(workflowID string) []*CostEntry {
	ct.mu.RLock()
	defer ct.mu.RUnlock()

	result := make([]*CostEntry, 0)
	for _, entry := range ct.entries {
		if entry.WorkflowID == workflowID {
			result = append(result, entry)
		}
	}

	return result
}

// SetRate sets the cost rate for a type
func (ct *CostTracker) SetRate(costType string, rate float64) {
	ct.mu.Lock()
	defer ct.mu.Unlock()

	ct.rates[costType] = rate
}

// GetRate gets the cost rate for a type
func (ct *CostTracker) GetRate(costType string) float64 {
	ct.mu.RLock()
	defer ct.mu.RUnlock()

	return ct.rates[costType]
}

// Clear clears all cost entries
func (ct *CostTracker) Clear() {
	ct.mu.Lock()
	defer ct.mu.Unlock()

	ct.entries = make([]*CostEntry, 0)
}