dev-pod-api-build/internal/store/usage.go

package store

import (
	"context"
	"fmt"
	"time"

	"github.com/iliaivanov/spec-kit-remote/cmd/dev-pod-api/internal/model"
)

// RecordPodStart records a pod start event for usage tracking.
func (s *Store) RecordPodStart(ctx context.Context, userID, podName string) error {
	_, err := s.db.Exec(ctx,
		`INSERT INTO usage_records (user_id, pod_name, event_type, value, recorded_at)
		 VALUES ($1, $2, $3, 0, NOW())`,
		userID, podName, model.EventPodStart)
	if err != nil {
		return fmt.Errorf("record pod start: %w", err)
	}
	return nil
}

// RecordPodStop records a pod stop event for usage tracking.
func (s *Store) RecordPodStop(ctx context.Context, userID, podName string) error {
	_, err := s.db.Exec(ctx,
		`INSERT INTO usage_records (user_id, pod_name, event_type, value, recorded_at)
		 VALUES ($1, $2, $3, 0, NOW())`,
		userID, podName, model.EventPodStop)
	if err != nil {
		return fmt.Errorf("record pod stop: %w", err)
	}
	return nil
}

// RecordResourceSample records periodic CPU and memory usage samples.
// cpuMillicores is CPU usage in millicores, memBytes is memory in bytes.
func (s *Store) RecordResourceSample(ctx context.Context, userID, podName string, cpuMillicores, memBytes float64) error {
	now := time.Now().UTC()
	_, err := s.db.Exec(ctx,
		`INSERT INTO usage_records (user_id, pod_name, event_type, value, recorded_at)
		 VALUES ($1, $2, $3, $4, $5), ($1, $2, $6, $7, $5)`,
		userID, podName, model.EventCPUSample, cpuMillicores, now,
		model.EventMemSample, memBytes)
	if err != nil {
		return fmt.Errorf("record resource sample: %w", err)
	}
	return nil
}

// RecordAIRequest records an AI proxy request event.
func (s *Store) RecordAIRequest(ctx context.Context, userID string) error {
	_, err := s.db.Exec(ctx,
		`INSERT INTO usage_records (user_id, pod_name, event_type, value, recorded_at)
		 VALUES ($1, '', $2, 1, NOW())`,
		userID, model.EventAIRequest)
	if err != nil {
		return fmt.Errorf("record ai request: %w", err)
	}
	return nil
}

// GetUsage returns aggregated usage for a user for a given month.
func (s *Store) GetUsage(ctx context.Context, userID string, year int, month time.Month, monthlyBudget int) (*model.UsageSummary, error) {
	return s.getUsage(ctx, userID, year, month, monthlyBudget, time.Now().UTC())
}

// getUsage is the internal implementation that accepts a "now" parameter for testing.
func (s *Store) getUsage(ctx context.Context, userID string, year int, month time.Month, monthlyBudget int, now time.Time) (*model.UsageSummary, error) {
	monthStart := time.Date(year, month, 1, 0, 0, 0, 0, time.UTC)
	monthEnd := monthStart.AddDate(0, 1, 0)

	// Calculate pod-hours from pod_start/pod_stop pairs.
	// For each pod_start, find the nearest subsequent pod_stop for the same user/pod.
	// If no stop found, use "now" (pod is still running).
	// Sessions that cross month boundaries are clamped to [monthStart, monthEnd].
	var podHours float64
	err := s.db.QueryRow(ctx, `
		WITH sessions AS (
			SELECT
				GREATEST(recorded_at, $2) AS start_time,
				(
					SELECT MIN(r2.recorded_at)
					FROM usage_records r2
					WHERE r2.user_id = r.user_id
					  AND r2.pod_name = r.pod_name
					  AND r2.event_type = 'pod_stop'
					  AND r2.recorded_at > r.recorded_at
				) AS stop_time
			FROM usage_records r
			WHERE r.user_id = $1
			  AND r.event_type = 'pod_start'
			  AND r.recorded_at < $3
		)
		SELECT COALESCE(SUM(
			EXTRACT(EPOCH FROM LEAST(COALESCE(stop_time, $4), $3) - start_time) / 3600.0
		), 0)
		FROM sessions
		WHERE COALESCE(stop_time, $4) > $2`,
		userID, monthStart, monthEnd, now).Scan(&podHours)
	if err != nil {
		return nil, fmt.Errorf("calculate pod hours: %w", err)
	}

	// CPU-hours from resource samples.
	// Each sample represents 60s of CPU at the sampled rate (millicores).
	// CPU-hours = sum(millicores) * 60s / 3600s / 1000m = sum(millicores) / 60000
	var cpuHours float64
	err = s.db.QueryRow(ctx, `
		SELECT COALESCE(SUM(value) / 60000.0, 0)
		FROM usage_records
		WHERE user_id = $1 AND event_type = $2
		  AND recorded_at >= $3 AND recorded_at < $4`,
		userID, model.EventCPUSample, monthStart, monthEnd).Scan(&cpuHours)
	if err != nil {
		return nil, fmt.Errorf("calculate cpu hours: %w", err)
	}

	// Count AI requests.
	var aiRequests int64
	err = s.db.QueryRow(ctx, `
		SELECT COUNT(*)
		FROM usage_records
		WHERE user_id = $1 AND event_type = $2
		  AND recorded_at >= $3 AND recorded_at < $4`,
		userID, model.EventAIRequest, monthStart, monthEnd).Scan(&aiRequests)
	if err != nil {
		return nil, fmt.Errorf("count ai requests: %w", err)
	}

	var budgetUsedPct float64
	if monthlyBudget > 0 {
		budgetUsedPct = (podHours / float64(monthlyBudget)) * 100
	}

	return &model.UsageSummary{
		PodHours:      podHours,
		CPUHours:      cpuHours,
		AIRequests:    aiRequests,
		BudgetUsedPct: budgetUsedPct,
	}, nil
}

// GetDailyUsage returns daily usage breakdown for a user for a given month.
// Pod-hours per day are estimated from CPU sample count (each sample = 1/60 hour of pod time).
func (s *Store) GetDailyUsage(ctx context.Context, userID string, year int, month time.Month) ([]model.DailyUsage, error) {
	monthStart := time.Date(year, month, 1, 0, 0, 0, 0, time.UTC)
	monthEnd := monthStart.AddDate(0, 1, 0)

	rows, err := s.db.Query(ctx, `
		SELECT
			DATE(recorded_at) AS day,
			COALESCE(SUM(CASE WHEN event_type = 'cpu_sample' THEN 1.0/60.0 ELSE 0 END), 0) AS pod_hours,
			COALESCE(SUM(CASE WHEN event_type = 'cpu_sample' THEN value / 60000.0 ELSE 0 END), 0) AS cpu_hours,
			COUNT(*) FILTER (WHERE event_type = 'ai_request') AS ai_requests
		FROM usage_records
		WHERE user_id = $1
		  AND recorded_at >= $2 AND recorded_at < $3
		  AND event_type IN ('cpu_sample', 'ai_request')
		GROUP BY DATE(recorded_at)
		ORDER BY DATE(recorded_at)`,
		userID, monthStart, monthEnd)
	if err != nil {
		return nil, fmt.Errorf("query daily usage: %w", err)
	}
	defer rows.Close()

	var result []model.DailyUsage
	for rows.Next() {
		var d model.DailyUsage
		var day time.Time
		if err := rows.Scan(&day, &d.PodHours, &d.CPUHours, &d.AIRequests); err != nil {
			return nil, fmt.Errorf("scan daily usage: %w", err)
		}
		d.Date = day.Format("2006-01-02")
		result = append(result, d)
	}
	return result, rows.Err()
}