gooseek/backend/internal/usage/repository.go

package usage

import (
	"context"
	"database/sql"
	"time"
)

type Repository struct {
	db *sql.DB
}

func NewRepository(db *sql.DB) *Repository {
	return &Repository{db: db}
}

func (r *Repository) RunMigrations(ctx context.Context) error {
	migrations := []string{
		`CREATE TABLE IF NOT EXISTS usage_metrics (
			id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
			user_id UUID NOT NULL,
			date DATE NOT NULL,
			tier VARCHAR(50) NOT NULL DEFAULT 'free',
			api_requests INT DEFAULT 0,
			llm_requests INT DEFAULT 0,
			llm_tokens INT DEFAULT 0,
			search_requests INT DEFAULT 0,
			storage_used BIGINT DEFAULT 0,
			created_at TIMESTAMPTZ DEFAULT NOW(),
			updated_at TIMESTAMPTZ DEFAULT NOW(),
			UNIQUE(user_id, date)
		)`,
		`CREATE INDEX IF NOT EXISTS idx_usage_user_date ON usage_metrics(user_id, date DESC)`,
		`CREATE INDEX IF NOT EXISTS idx_usage_date ON usage_metrics(date)`,
	}

	for _, m := range migrations {
		if _, err := r.db.ExecContext(ctx, m); err != nil {
			return err
		}
	}
	return nil
}

func (r *Repository) GetTodayUsage(ctx context.Context, userID string) (*UsageMetric, error) {
	today := time.Now().Format("2006-01-02")
	return r.GetUsageByDate(ctx, userID, today)
}

func (r *Repository) GetUsageByDate(ctx context.Context, userID, date string) (*UsageMetric, error) {
	query := `
		SELECT id, user_id, date, tier, api_requests, llm_requests, llm_tokens, 
		       search_requests, storage_used, created_at, updated_at
		FROM usage_metrics 
		WHERE user_id = $1 AND date = $2
	`

	var m UsageMetric
	var dateVal time.Time
	err := r.db.QueryRowContext(ctx, query, userID, date).Scan(
		&m.ID, &m.UserID, &dateVal, &m.Tier,
		&m.APIRequests, &m.LLMRequests, &m.LLMTokens,
		&m.SearchReqs, &m.StorageUsed, &m.CreatedAt, &m.UpdatedAt,
	)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	m.Date = dateVal.Format("2006-01-02")
	return &m, nil
}

func (r *Repository) IncrementAPIRequests(ctx context.Context, userID, tier string) error {
	today := time.Now().Format("2006-01-02")
	query := `
		INSERT INTO usage_metrics (user_id, date, tier, api_requests)
		VALUES ($1, $2, $3, 1)
		ON CONFLICT (user_id, date) 
		DO UPDATE SET api_requests = usage_metrics.api_requests + 1, updated_at = NOW()
	`
	_, err := r.db.ExecContext(ctx, query, userID, today, tier)
	return err
}

func (r *Repository) IncrementLLMUsage(ctx context.Context, userID, tier string, tokens int) error {
	today := time.Now().Format("2006-01-02")
	query := `
		INSERT INTO usage_metrics (user_id, date, tier, llm_requests, llm_tokens)
		VALUES ($1, $2, $3, 1, $4)
		ON CONFLICT (user_id, date) 
		DO UPDATE SET 
			llm_requests = usage_metrics.llm_requests + 1,
			llm_tokens = usage_metrics.llm_tokens + $4,
			updated_at = NOW()
	`
	_, err := r.db.ExecContext(ctx, query, userID, today, tier, tokens)
	return err
}

func (r *Repository) IncrementSearchRequests(ctx context.Context, userID, tier string) error {
	today := time.Now().Format("2006-01-02")
	query := `
		INSERT INTO usage_metrics (user_id, date, tier, search_requests)
		VALUES ($1, $2, $3, 1)
		ON CONFLICT (user_id, date) 
		DO UPDATE SET search_requests = usage_metrics.search_requests + 1, updated_at = NOW()
	`
	_, err := r.db.ExecContext(ctx, query, userID, today, tier)
	return err
}

func (r *Repository) UpdateStorageUsed(ctx context.Context, userID, tier string, bytes int64) error {
	today := time.Now().Format("2006-01-02")
	query := `
		INSERT INTO usage_metrics (user_id, date, tier, storage_used)
		VALUES ($1, $2, $3, $4)
		ON CONFLICT (user_id, date) 
		DO UPDATE SET storage_used = $4, updated_at = NOW()
	`
	_, err := r.db.ExecContext(ctx, query, userID, today, tier, bytes)
	return err
}

func (r *Repository) GetUsageHistory(ctx context.Context, userID string, days int) ([]*UsageMetric, error) {
	query := `
		SELECT id, user_id, date, tier, api_requests, llm_requests, llm_tokens, 
		       search_requests, storage_used, created_at, updated_at
		FROM usage_metrics 
		WHERE user_id = $1 AND date >= CURRENT_DATE - $2::int
		ORDER BY date DESC
	`

	rows, err := r.db.QueryContext(ctx, query, userID, days)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var metrics []*UsageMetric
	for rows.Next() {
		var m UsageMetric
		var dateVal time.Time
		if err := rows.Scan(
			&m.ID, &m.UserID, &dateVal, &m.Tier,
			&m.APIRequests, &m.LLMRequests, &m.LLMTokens,
			&m.SearchReqs, &m.StorageUsed, &m.CreatedAt, &m.UpdatedAt,
		); err != nil {
			return nil, err
		}
		m.Date = dateVal.Format("2006-01-02")
		metrics = append(metrics, &m)
	}

	return metrics, nil
}

func (r *Repository) GetUserEmail(ctx context.Context, userID string) (string, error) {
	var emailAddr string
	err := r.db.QueryRowContext(ctx,
		"SELECT email FROM auth_users WHERE id = $1", userID,
	).Scan(&emailAddr)
	if err != nil {
		return "", err
	}
	return emailAddr, nil
}

func (r *Repository) CheckLLMLimits(ctx context.Context, userID, tier string) (bool, string) {
	usage, err := r.GetTodayUsage(ctx, userID)
	if err != nil {
		return true, ""
	}

	limits := GetLimits(tier)

	if usage != nil {
		if usage.LLMRequests >= limits.LLMRequestsPerDay {
			return false, "Daily LLM request limit exceeded"
		}
		if usage.LLMTokens >= limits.LLMTokensPerDay {
			return false, "Daily LLM token limit exceeded"
		}
	}

	return true, ""
}