pangolin/server/internal/codes/store.go

package codes

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

// PlanCode represents a plan tier.
type PlanCode string

const (
	PlanFree PlanCode = "free"
	PlanPro  PlanCode = "pro"
	PlanTeam PlanCode = "team"
)

// planTier returns a numeric tier for comparison (higher = better).
func planTier(p PlanCode) int {
	switch p {
	case PlanTeam:
		return 3
	case PlanPro:
		return 2
	default:
		return 1
	}
}

// BatchChannel is the source channel for a code batch.
type BatchChannel string

const (
	ChannelStore  BatchChannel = "store"
	ChannelTG     BatchChannel = "tg"
	ChannelLine   BatchChannel = "line"
	ChannelManual BatchChannel = "manual"
)

// CodeRow mirrors the `codes` DB row (hash only; no plaintext).
type CodeRow struct {
	ID           int64
	CodeHash     string // SHA-256 hex of canonical plaintext
	PlanID       int64
	PlanCode     PlanCode
	DurationDays int
	BatchID      int64
	Status       string // unused | redeemed | void
	RedeemedBy   sql.NullInt64
	RedeemedAt   sql.NullTime
}

// SubscriptionRow mirrors the `subscriptions` DB row.
type SubscriptionRow struct {
	ID           int64
	UserID       int64
	PlanID       int64
	PlanCode     PlanCode
	ExpiresAt    time.Time
	Source       string
}

// BatchRow mirrors the `code_batches` DB row.
type BatchRow struct {
	ID        int64
	Channel   BatchChannel
	CreatedBy string
	Note      sql.NullString
	CreatedAt time.Time
}

// Store wraps a *sql.DB and exposes all database operations needed by the
// codes package.  Every method that modifies data is safe to call inside or
// outside an explicit transaction; methods that accept a *sql.Tx run within
// that transaction.
type Store struct {
	db *sql.DB
}

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

// --------------------------------------------------------------------------
// Batch and code creation (used by Service.CreateBatch and webhook)
// --------------------------------------------------------------------------

// CreateBatch inserts a new code_batches row and returns its ID.
func (s *Store) CreateBatch(ctx context.Context, channel BatchChannel, createdBy, note string) (int64, error) {
	res, err := s.db.ExecContext(ctx,
		`INSERT INTO code_batches (channel, created_by, note, created_at)
		 VALUES (?, ?, NULLIF(?, ''), UTC_TIMESTAMP(6))`,
		string(channel), createdBy, note)
	if err != nil {
		return 0, fmt.Errorf("store.CreateBatch: %w", err)
	}
	id, err := res.LastInsertId()
	if err != nil {
		return 0, fmt.Errorf("store.CreateBatch last id: %w", err)
	}
	return id, nil
}

// CreateCode inserts a single codes row.  It returns ErrDuplicate if a row
// with the same code_hash already exists (UNIQUE constraint violation).
func (s *Store) CreateCode(ctx context.Context, codeHash string, planID int64, durationDays int, batchID int64) error {
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO codes (code_hash, plan_id, duration_days, batch_id, status)
		 VALUES (?, ?, ?, ?, 'unused')`,
		codeHash, planID, durationDays, batchID)
	if err != nil {
		if isDuplicateKey(err) {
			return ErrDuplicate
		}
		return fmt.Errorf("store.CreateCode: %w", err)
	}
	return nil
}

// --------------------------------------------------------------------------
// Redemption (used by Service.Redeem inside a transaction)
// --------------------------------------------------------------------------

// FindCodeByHashForUpdate looks up a codes row by its SHA-256 hash using
// SELECT … FOR UPDATE so that the row is locked for the duration of the
// surrounding transaction.  It also fetches the plan code from the plans
// table in the same query.
func (s *Store) FindCodeByHashForUpdate(ctx context.Context, tx *sql.Tx, hash string) (*CodeRow, error) {
	row := tx.QueryRowContext(ctx,
		`SELECT c.id, c.code_hash, c.plan_id, p.code, c.duration_days,
		        c.batch_id, c.status, c.redeemed_by, c.redeemed_at
		 FROM codes c
		 JOIN plans p ON p.id = c.plan_id
		 WHERE c.code_hash = ?
		 FOR UPDATE`,
		hash)

	var cr CodeRow
	if err := row.Scan(
		&cr.ID, &cr.CodeHash, &cr.PlanID, &cr.PlanCode, &cr.DurationDays,
		&cr.BatchID, &cr.Status, &cr.RedeemedBy, &cr.RedeemedAt,
	); err == sql.ErrNoRows {
		return nil, nil
	} else if err != nil {
		return nil, fmt.Errorf("store.FindCodeByHashForUpdate: %w", err)
	}
	return &cr, nil
}

// MarkRedeemed updates a codes row to status='redeemed' within tx.
func (s *Store) MarkRedeemed(ctx context.Context, tx *sql.Tx, codeID, userID int64) error {
	_, err := tx.ExecContext(ctx,
		`UPDATE codes SET status='redeemed', redeemed_by=?, redeemed_at=UTC_TIMESTAMP(6)
		 WHERE id=? AND status='unused'`,
		userID, codeID)
	if err != nil {
		return fmt.Errorf("store.MarkRedeemed: %w", err)
	}
	return nil
}

// --------------------------------------------------------------------------
// Subscription helpers (used inside redeem transaction)
// --------------------------------------------------------------------------

// GetPlanID returns the primary key of plans.code.
func (s *Store) GetPlanID(ctx context.Context, code PlanCode) (int64, error) {
	var id int64
	err := s.db.QueryRowContext(ctx, `SELECT id FROM plans WHERE code=?`, string(code)).Scan(&id)
	if err != nil {
		return 0, fmt.Errorf("store.GetPlanID(%s): %w", code, err)
	}
	return id, nil
}

// GetActiveSubscriptions returns all non-expired subscriptions for userID,
// ordered by expires_at DESC.  Called inside the redeem transaction.
func (s *Store) GetActiveSubscriptions(ctx context.Context, tx *sql.Tx, userID int64) ([]SubscriptionRow, error) {
	rows, err := tx.QueryContext(ctx,
		`SELECT s.id, s.user_id, s.plan_id, p.code, s.expires_at, s.source
		 FROM subscriptions s
		 JOIN plans p ON p.id = s.plan_id
		 WHERE s.user_id=? AND s.expires_at > UTC_TIMESTAMP(6)
		 ORDER BY s.expires_at DESC`,
		userID)
	if err != nil {
		return nil, fmt.Errorf("store.GetActiveSubscriptions: %w", err)
	}
	defer rows.Close()

	var subs []SubscriptionRow
	for rows.Next() {
		var sr SubscriptionRow
		if err := rows.Scan(&sr.ID, &sr.UserID, &sr.PlanID, &sr.PlanCode, &sr.ExpiresAt, &sr.Source); err != nil {
			return nil, fmt.Errorf("store.GetActiveSubscriptions scan: %w", err)
		}
		subs = append(subs, sr)
	}
	return subs, rows.Err()
}

// ExtendSubscription sets expires_at to max(current_expires_at, now) + duration.
// Called inside the redeem transaction.
func (s *Store) ExtendSubscription(ctx context.Context, tx *sql.Tx, subID int64, durationDays int) error {
	_, err := tx.ExecContext(ctx,
		`UPDATE subscriptions
		 SET expires_at = DATE_ADD(
		       GREATEST(expires_at, UTC_TIMESTAMP(6)),
		       INTERVAL ? DAY)
		 WHERE id=?`,
		durationDays, subID)
	if err != nil {
		return fmt.Errorf("store.ExtendSubscription: %w", err)
	}
	return nil
}

// CreateSubscription inserts a new subscriptions row.  expires_at is set to
// max(now, latest_expires_at_for_same_plan) + duration_days.
// latestSamePlan may be zero-value if the user has no existing sub for that plan.
func (s *Store) CreateSubscription(
	ctx context.Context,
	tx *sql.Tx,
	userID, planID int64,
	durationDays int,
	latestSamePlan time.Time,
) (int64, error) {
	now := time.Now().UTC()
	base := now
	if latestSamePlan.After(now) {
		base = latestSamePlan
	}
	expiresAt := base.AddDate(0, 0, durationDays)

	res, err := tx.ExecContext(ctx,
		`INSERT INTO subscriptions (user_id, plan_id, expires_at, source, created_at)
		 VALUES (?, ?, ?, 'code', UTC_TIMESTAMP(6))`,
		userID, planID, expiresAt)
	if err != nil {
		return 0, fmt.Errorf("store.CreateSubscription: %w", err)
	}
	id, _ := res.LastInsertId()
	return id, nil
}

// --------------------------------------------------------------------------
// Audit log
// --------------------------------------------------------------------------

// WriteAuditLog inserts an audit_log row.  actor and target are string
// identifiers; meta is a JSON object (may be nil/empty).
func (s *Store) WriteAuditLog(ctx context.Context, tx *sql.Tx, actor, action, target, metaJSON string) error {
	var err error
	if metaJSON == "" {
		metaJSON = "null"
	}
	if tx != nil {
		_, err = tx.ExecContext(ctx,
			`INSERT INTO audit_log (actor, action, target, meta, at) VALUES (?, ?, ?, ?, UTC_TIMESTAMP(6))`,
			actor, action, target, metaJSON)
	} else {
		_, err = s.db.ExecContext(ctx,
			`INSERT INTO audit_log (actor, action, target, meta, at) VALUES (?, ?, ?, ?, UTC_TIMESTAMP(6))`,
			actor, action, target, metaJSON)
	}
	if err != nil {
		return fmt.Errorf("store.WriteAuditLog: %w", err)
	}
	return nil
}

// --------------------------------------------------------------------------
// Webhook-specific lookup
// --------------------------------------------------------------------------

// CodeExistsByHash returns true if a codes row with the given hash already exists.
func (s *Store) CodeExistsByHash(ctx context.Context, hash string) (bool, error) {
	var count int
	err := s.db.QueryRowContext(ctx, `SELECT COUNT(1) FROM codes WHERE code_hash=?`, hash).Scan(&count)
	if err != nil {
		return false, fmt.Errorf("store.CodeExistsByHash: %w", err)
	}
	return count > 0, nil
}

// BeginTx starts a new transaction at Read Committed isolation level.
// The SELECT … FOR UPDATE in FindCodeByHashForUpdate provides the necessary
// row-level exclusivity; Serializable is intentionally avoided to reduce
// deadlock risk under concurrent redemptions.
func (s *Store) BeginTx(ctx context.Context) (*sql.Tx, error) {
	return s.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelReadCommitted})
}

// --------------------------------------------------------------------------
// helper
// --------------------------------------------------------------------------

// isDuplicateKey returns true if err is a MySQL duplicate-key error (1062).
func isDuplicateKey(err error) bool {
	if err == nil {
		return false
	}
	// go-sql-driver wraps MySQL errors; the error number is accessible via
	// the mysql.MySQLError type.  We do a string match to avoid importing
	// the mysql package here.
	msg := err.Error()
	return strings.Contains(msg, "Duplicate entry") ||
		strings.Contains(msg, "1062")
}