core/server/store/postgres/usage.go

package postgres

import (
	"context"
	"fmt"
	"strings"
	"time"

	"switchboard-core/models"
	"switchboard-core/store"
)

type UsageStore struct{}

func NewUsageStore() *UsageStore { return &UsageStore{} }

// Log inserts a usage entry.
func (s *UsageStore) Log(ctx context.Context, entry *models.UsageEntry) error {
	_, err := DB.ExecContext(ctx, `
		INSERT INTO usage_log (
			channel_id, user_id, provider_config_id, provider_scope,
			model_id, role,
			prompt_tokens, completion_tokens,
			cache_creation_tokens, cache_read_tokens,
			cost_input, cost_output
		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
	`,
		entry.ChannelID, entry.UserID, entry.ProviderConfigID, entry.ProviderScope,
		entry.ModelID, entry.Role,
		entry.PromptTokens, entry.CompletionTokens,
		entry.CacheCreationTokens, entry.CacheReadTokens,
		entry.CostInput, entry.CostOutput,
	)
	return err
}

// QueryByUser returns aggregated usage for a specific user.
func (s *UsageStore) QueryByUser(ctx context.Context, userID string, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
	where := []string{"user_id = $1"}
	args := []interface{}{userID}
	return s.query(ctx, where, args, opts)
}

// QueryByUserPersonal returns aggregated usage for a user's own (BYOK) providers only.
func (s *UsageStore) QueryByUserPersonal(ctx context.Context, userID string, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
	where := []string{
		"user_id = $1",
		"provider_config_id IN (SELECT id FROM provider_configs WHERE scope = 'personal' AND owner_id = $1)",
	}
	args := []interface{}{userID}
	return s.query(ctx, where, args, opts)
}

// QueryByTeam returns aggregated usage for all members of a team (admin view).
func (s *UsageStore) QueryByTeam(ctx context.Context, teamID string, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
	where := []string{"user_id IN (SELECT user_id FROM team_members WHERE team_id = $1)"}
	args := []interface{}{teamID}
	return s.query(ctx, where, args, opts)
}

// QueryByTeamProviders returns usage against providers owned by this team.
// Used by team admins to see how their team's API keys are being consumed.
func (s *UsageStore) QueryByTeamProviders(ctx context.Context, teamID string, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
	where := []string{"provider_config_id IN (SELECT id FROM provider_configs WHERE scope = 'team' AND owner_id = $1)"}
	args := []interface{}{teamID}
	return s.query(ctx, where, args, opts)
}

// QueryByModel returns aggregated usage grouped by model.
func (s *UsageStore) QueryByModel(ctx context.Context, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
	return s.query(ctx, nil, nil, opts)
}

// GetTotals returns aggregate totals (admin view — excludes BYOK by default).
func (s *UsageStore) GetTotals(ctx context.Context, opts store.UsageQueryOptions) (*models.UsageTotals, error) {
	where, args := s.buildFilters(nil, nil, opts)

	q := fmt.Sprintf(`
		SELECT
			COUNT(*),
			COALESCE(SUM(prompt_tokens), 0),
			COALESCE(SUM(completion_tokens), 0),
			COALESCE(SUM(COALESCE(cost_input, 0) + COALESCE(cost_output, 0)), 0)
		FROM usage_log
		WHERE %s
	`, strings.Join(where, " AND "))

	var t models.UsageTotals
	err := DB.QueryRowContext(ctx, q, args...).Scan(
		&t.Requests, &t.InputTokens, &t.OutputTokens, &t.TotalCost,
	)
	return &t, err
}

// GetTeamProviderTotals returns aggregate totals for providers owned by a team.
func (s *UsageStore) GetTeamProviderTotals(ctx context.Context, teamID string, opts store.UsageQueryOptions) (*models.UsageTotals, error) {
	baseWhere := []string{"provider_config_id IN (SELECT id FROM provider_configs WHERE scope = 'team' AND owner_id = $1)"}
	baseArgs := []interface{}{teamID}
	where, args := s.buildFilters(baseWhere, baseArgs, opts)

	q := fmt.Sprintf(`
		SELECT
			COUNT(*),
			COALESCE(SUM(prompt_tokens), 0),
			COALESCE(SUM(completion_tokens), 0),
			COALESCE(SUM(COALESCE(cost_input, 0) + COALESCE(cost_output, 0)), 0)
		FROM usage_log
		WHERE %s
	`, strings.Join(where, " AND "))

	var t models.UsageTotals
	err := DB.QueryRowContext(ctx, q, args...).Scan(
		&t.Requests, &t.InputTokens, &t.OutputTokens, &t.TotalCost,
	)
	return &t, err
}

// GetPersonalTotals returns aggregate totals for a user's own provider usage.
// Only includes usage against personal (BYOK) providers — global provider
// costs are the org's responsibility, not the user's.
func (s *UsageStore) GetPersonalTotals(ctx context.Context, userID string, opts store.UsageQueryOptions) (*models.UsageTotals, error) {
	baseWhere := []string{
		"user_id = $1",
		"provider_config_id IN (SELECT id FROM provider_configs WHERE scope = 'personal' AND owner_id = $1)",
	}
	baseArgs := []interface{}{userID}
	opts.ExcludeBYOK = false
	where, args := s.buildFilters(baseWhere, baseArgs, opts)

	q := fmt.Sprintf(`
		SELECT
			COUNT(*),
			COALESCE(SUM(prompt_tokens), 0),
			COALESCE(SUM(completion_tokens), 0),
			COALESCE(SUM(COALESCE(cost_input, 0) + COALESCE(cost_output, 0)), 0)
		FROM usage_log
		WHERE %s
	`, strings.Join(where, " AND "))

	var t models.UsageTotals
	err := DB.QueryRowContext(ctx, q, args...).Scan(
		&t.Requests, &t.InputTokens, &t.OutputTokens, &t.TotalCost,
	)
	return &t, err
}

// CountRecentByRole counts how many usage entries a user has for a given
// role within the specified duration. Used for rate limiting utility calls.
func (s *UsageStore) CountRecentByRole(ctx context.Context, userID, role string, since time.Time) (int, error) {
	var count int
	err := DB.QueryRowContext(ctx, `
		SELECT COUNT(*) FROM usage_log
		WHERE user_id = $1 AND role = $2 AND created_at >= $3
	`, userID, role, since).Scan(&count)
	return count, err
}

// ── Internal ───────────────────────────────

func (s *UsageStore) buildFilters(baseWhere []string, baseArgs []interface{}, opts store.UsageQueryOptions) ([]string, []interface{}) {
	where := append([]string{}, baseWhere...)
	args := append([]interface{}{}, baseArgs...)
	idx := len(args) + 1

	if opts.ExcludeBYOK {
		where = append(where, fmt.Sprintf("provider_scope != $%d", idx))
		args = append(args, "personal")
		idx++
	}
	if opts.Since != nil {
		where = append(where, fmt.Sprintf("created_at >= $%d", idx))
		args = append(args, *opts.Since)
		idx++
	}
	if opts.Until != nil {
		where = append(where, fmt.Sprintf("created_at < $%d", idx))
		args = append(args, *opts.Until)
		idx++
	}

	if len(where) == 0 {
		where = append(where, "1=1")
	}
	return where, args
}

func (s *UsageStore) query(ctx context.Context, baseWhere []string, baseArgs []interface{}, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
	where, args := s.buildFilters(baseWhere, baseArgs, opts)

	// Determine GROUP BY expression and label
	groupExpr := "model_id"
	labelExpr := "model_id"
	switch opts.GroupBy {
	case "day":
		groupExpr = "DATE(created_at)"
		labelExpr = "DATE(created_at)::text"
	case "user":
		groupExpr = "user_id"
		labelExpr = "user_id::text"
	case "provider":
		groupExpr = "provider_config_id"
		labelExpr = "provider_config_id::text"
	case "model":
		groupExpr = "model_id"
		labelExpr = "model_id"
	default:
		groupExpr = "model_id"
		labelExpr = "model_id"
	}

	limit := opts.Limit
	if limit <= 0 {
		limit = 100
	}

	q := fmt.Sprintf(`
		SELECT
			%s AS group_key,
			%s AS label,
			COUNT(*) AS requests,
			COALESCE(SUM(prompt_tokens), 0) AS input_tokens,
			COALESCE(SUM(completion_tokens), 0) AS output_tokens,
			COALESCE(SUM(COALESCE(cost_input, 0) + COALESCE(cost_output, 0)), 0) AS total_cost
		FROM usage_log
		WHERE %s
		GROUP BY %s
		ORDER BY total_cost DESC
		LIMIT %d
	`, groupExpr, labelExpr, strings.Join(where, " AND "), groupExpr, limit)

	rows, err := DB.QueryContext(ctx, q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var results []models.UsageAggregate
	for rows.Next() {
		var a models.UsageAggregate
		if err := rows.Scan(&a.GroupKey, &a.Label, &a.Requests, &a.InputTokens, &a.OutputTokens, &a.TotalCost); err != nil {
			return nil, err
		}
		results = append(results, a)
	}
	return results, rows.Err()
}