pangolin/server/internal/usage/store.go

package usage

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

	dbx "github.com/wangjia/pangolin/server/internal/db"
)

// Unlimited is the sentinel returned by CheckFreeConnect for plans with no
// daily-minute cap (pro / team).  Callers (e.g. the connect endpoint) treat it
// as "no minute limit; use the default paid credential TTL".
const Unlimited = -1

// dateLayout is the canonical UTC day format used as the usage_daily.date key
// and in the API response.
const dateLayout = "2006-01-02"

// Plan is the minimal subset of the plans row needed for quota decisions.
type Plan struct {
	Code         string        // free | pro | team
	DailyMinutes sql.NullInt64 // free=10; NULL = unlimited
	AdGate       bool          // free=true: daily rewarded-ad unlock required
}

// DailyUsage mirrors a usage_daily row.  It carries only aggregate byte and
// minute counts plus the ad-unlock timestamp — never destinations or DNS, per
// the no-log policy.
type DailyUsage struct {
	Date         time.Time
	BytesUp      uint64
	BytesDown    uint64
	MinutesUsed  int
	AdUnlockedAt sql.NullTime
}

// Store wraps a *sql.DB and exposes the usage_daily / plans / users queries the
// usage package needs.
type Store struct {
	db      *sql.DB
	dialect dbx.Dialect
}

// NewStore creates a Store backed by the given connection pool (MySQL or SQLite).
func NewStore(db *sql.DB) *Store { return &Store{db: db, dialect: dbx.DialectForDB(db)} }

// AggregateUsage accumulates one usage delta into usage_daily for (userID, day)
// using INSERT … ON DUPLICATE KEY UPDATE so concurrent reports from multiple
// nodes add up correctly.  day is truncated to its UTC calendar date.
func (s *Store) AggregateUsage(ctx context.Context, userID int64, day time.Time, bytesUp, bytesDown uint64, minutes int) error {
	d := day.UTC().Format(dateLayout)
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO usage_daily (user_id, date, bytes_up, bytes_down, minutes_used)
		 VALUES (?, ?, ?, ?, ?)
		 `+s.dialect.Upsert([]string{"user_id", "date"},
			"bytes_up = bytes_up + EXCLUDED.bytes_up",
			"bytes_down = bytes_down + EXCLUDED.bytes_down",
			"minutes_used = minutes_used + EXCLUDED.minutes_used"),
		userID, d, bytesUp, bytesDown, minutes)
	if err != nil {
		return fmt.Errorf("store.AggregateUsage: %w", err)
	}
	return nil
}

// LookupUserIDByDPUUID resolves a data-plane credential UUID to the owning
// user_id.  Returns (0, nil) when no user matches (e.g. a rotated/stale dp_uuid)
// so the caller can decide to silently drop the report.
func (s *Store) LookupUserIDByDPUUID(ctx context.Context, dpUUID string) (int64, error) {
	var id int64
	err := s.db.QueryRowContext(ctx,
		`SELECT id FROM users WHERE dp_uuid = ? LIMIT 1`, dpUUID).Scan(&id)
	if err == sql.ErrNoRows {
		return 0, nil
	}
	if err != nil {
		return 0, fmt.Errorf("store.LookupUserIDByDPUUID: %w", err)
	}
	return id, nil
}

// GetUsageRange returns the usage_daily rows for userID with date in
// [from, to] (inclusive, UTC dates), ordered ascending.  Missing days are NOT
// filled here; zero-filling is the handler's responsibility.
func (s *Store) GetUsageRange(ctx context.Context, userID int64, from, to time.Time) ([]DailyUsage, error) {
	rows, err := s.db.QueryContext(ctx,
		`SELECT date, bytes_up, bytes_down, minutes_used, ad_unlocked_at
		 FROM usage_daily
		 WHERE user_id = ? AND date BETWEEN ? AND ?
		 ORDER BY date ASC`,
		userID, from.UTC().Format(dateLayout), to.UTC().Format(dateLayout))
	if err != nil {
		return nil, fmt.Errorf("store.GetUsageRange: %w", err)
	}
	defer rows.Close()

	var out []DailyUsage
	for rows.Next() {
		var u DailyUsage
		if err := rows.Scan(&u.Date, &u.BytesUp, &u.BytesDown, &u.MinutesUsed, &u.AdUnlockedAt); err != nil {
			return nil, fmt.Errorf("store.GetUsageRange scan: %w", err)
		}
		out = append(out, u)
	}
	return out, rows.Err()
}

// GetDay returns the usage_daily row for (userID, day), or nil if none exists.
func (s *Store) GetDay(ctx context.Context, userID int64, day time.Time) (*DailyUsage, error) {
	var u DailyUsage
	err := s.db.QueryRowContext(ctx,
		`SELECT date, bytes_up, bytes_down, minutes_used, ad_unlocked_at
		 FROM usage_daily
		 WHERE user_id = ? AND date = ?`,
		userID, day.UTC().Format(dateLayout)).
		Scan(&u.Date, &u.BytesUp, &u.BytesDown, &u.MinutesUsed, &u.AdUnlockedAt)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("store.GetDay: %w", err)
	}
	return &u, nil
}

// MarkAdUnlocked sets usage_daily.ad_unlocked_at for (userID, day) to now if it
// is not already set.  It returns alreadyUnlocked=true when the day was already
// unlocked (making a second call idempotent).  Runs inside a transaction with
// SELECT … FOR UPDATE to be safe under concurrent unlock attempts.
func (s *Store) MarkAdUnlocked(ctx context.Context, userID int64, day time.Time) (alreadyUnlocked bool, err error) {
	d := day.UTC().Format(dateLayout)
	now := time.Now().UTC()

	tx, err := s.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelReadCommitted})
	if err != nil {
		return false, fmt.Errorf("store.MarkAdUnlocked begin: %w", err)
	}
	committed := false
	defer func() {
		if !committed {
			_ = tx.Rollback()
		}
	}()

	var unlockedAt sql.NullTime
	row := tx.QueryRowContext(ctx,
		`SELECT ad_unlocked_at FROM usage_daily WHERE user_id = ? AND date = ? `+s.dialect.LockForUpdate(),
		userID, d)
	switch err := row.Scan(&unlockedAt); err {
	case nil:
		if unlockedAt.Valid {
			// Already unlocked today → idempotent success.
			if cErr := tx.Commit(); cErr != nil {
				return false, fmt.Errorf("store.MarkAdUnlocked commit: %w", cErr)
			}
			committed = true
			return true, nil
		}
		if _, uErr := tx.ExecContext(ctx,
			`UPDATE usage_daily SET ad_unlocked_at = ?
			 WHERE user_id = ? AND date = ? AND ad_unlocked_at IS NULL`,
			now, userID, d); uErr != nil {
			return false, fmt.Errorf("store.MarkAdUnlocked update: %w", uErr)
		}
	case sql.ErrNoRows:
		if _, iErr := tx.ExecContext(ctx,
			`INSERT INTO usage_daily (user_id, date, ad_unlocked_at)
			 VALUES (?, ?, ?)`,
			userID, d, now); iErr != nil {
			return false, fmt.Errorf("store.MarkAdUnlocked insert: %w", iErr)
		}
	default:
		return false, fmt.Errorf("store.MarkAdUnlocked select: %w", err)
	}

	if cErr := tx.Commit(); cErr != nil {
		return false, fmt.Errorf("store.MarkAdUnlocked commit: %w", cErr)
	}
	committed = true
	return false, nil
}

// EffectivePlan returns the plan that currently governs userID: the highest
// tier among the user's non-expired subscriptions, falling back to the free
// plan when the user has no active subscription (trial-expired free state).
func (s *Store) EffectivePlan(ctx context.Context, userID int64) (*Plan, error) {
	var p Plan
	err := s.db.QueryRowContext(ctx,
		`SELECT p.code, p.daily_minutes, p.ad_gate
		 FROM subscriptions s
		 JOIN plans p ON p.id = s.plan_id
		 WHERE s.user_id = ? AND s.expires_at > ?
		 ORDER BY CASE p.code WHEN 'team' THEN 0 WHEN 'pro' THEN 1 WHEN 'free' THEN 2 ELSE 3 END,
		          s.expires_at DESC
		 LIMIT 1`,
		userID, time.Now().UTC()).Scan(&p.Code, &p.DailyMinutes, &p.AdGate)
	if err == nil {
		return &p, nil
	}
	if err != sql.ErrNoRows {
		return nil, fmt.Errorf("store.EffectivePlan: %w", err)
	}

	// No active subscription → free plan.
	if err := s.db.QueryRowContext(ctx,
		`SELECT code, daily_minutes, ad_gate FROM plans WHERE code = 'free'`).
		Scan(&p.Code, &p.DailyMinutes, &p.AdGate); err != nil {
		return nil, fmt.Errorf("store.EffectivePlan free fallback: %w", err)
	}
	return &p, nil
}