sub2api/backend/internal/repository/global_model_pricing_repo.go

package repository

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

	"github.com/Wei-Shaw/sub2api/internal/service"
)

// globalModelPricingRepository 实现 service.GlobalModelPricingRepository。
// 直接 SQL 访问 global_model_pricing 表（无需 ent codegen）。
type globalModelPricingRepository struct {
	sql sqlExecutor
}

func NewGlobalModelPricingRepository(sqlDB *sql.DB) service.GlobalModelPricingRepository {
	return &globalModelPricingRepository{sql: sqlDB}
}

const globalModelPricingColumns = `id, model_pattern, billing_mode, input_price, output_price, cache_write_price, cache_write_1h_price, cache_read_price, per_request_price, enabled, created_at, updated_at`

func scanGlobalModelPricingRow(scan func(dest ...any) error) (*service.GlobalModelPrice, error) {
	var p service.GlobalModelPrice
	var input, output, cw, cw1h, cr, pr sql.NullFloat64
	var created, updated sql.NullTime
	if err := scan(&p.ID, &p.ModelPattern, &p.BillingMode,
		&input, &output, &cw, &cw1h, &cr, &pr,
		&p.Enabled, &created, &updated); err != nil {
		return nil, err
	}
	if input.Valid {
		p.InputPrice = &input.Float64
	}
	if output.Valid {
		p.OutputPrice = &output.Float64
	}
	if cw.Valid {
		p.CacheWritePrice = &cw.Float64
	}
	if cw1h.Valid {
		p.CacheWrite1hPrice = &cw1h.Float64
	}
	if cr.Valid {
		p.CacheReadPrice = &cr.Float64
	}
	if pr.Valid {
		p.PerRequestPrice = &pr.Float64
	}
	if created.Valid {
		p.CreatedAt = created.Time
	}
	if updated.Valid {
		p.UpdatedAt = updated.Time
	}
	return &p, nil
}

// ListEnabled 返回全部启用的全局定价（service 层做快照缓存）。
func (r *globalModelPricingRepository) ListEnabled(ctx context.Context) ([]service.GlobalModelPrice, error) {
	rows, err := r.sql.QueryContext(ctx,
		`SELECT `+globalModelPricingColumns+` FROM global_model_pricing WHERE enabled = TRUE ORDER BY id`)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()
	out := make([]service.GlobalModelPrice, 0)
	for rows.Next() {
		p, err := scanGlobalModelPricingRow(rows.Scan)
		if err != nil {
			return nil, err
		}
		out = append(out, *p)
	}
	return out, rows.Err()
}

// ListAll 返回全部（含禁用，供管理端列表）。
func (r *globalModelPricingRepository) ListAll(ctx context.Context) ([]service.GlobalModelPrice, error) {
	rows, err := r.sql.QueryContext(ctx,
		`SELECT `+globalModelPricingColumns+` FROM global_model_pricing ORDER BY id`)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()
	out := make([]service.GlobalModelPrice, 0)
	for rows.Next() {
		p, err := scanGlobalModelPricingRow(rows.Scan)
		if err != nil {
			return nil, err
		}
		out = append(out, *p)
	}
	return out, rows.Err()
}

func nullFloat64(v *float64) any {
	if v == nil {
		return nil
	}
	return *v
}

func (r *globalModelPricingRepository) Create(ctx context.Context, in service.GlobalModelPriceInput) (*service.GlobalModelPrice, error) {
	pattern := strings.ToLower(strings.TrimSpace(in.ModelPattern))
	mode := strings.TrimSpace(string(in.BillingMode))
	if mode == "" {
		mode = string(service.BillingModeToken)
	}
	rows, err := r.sql.QueryContext(ctx, `
		INSERT INTO global_model_pricing
			(model_pattern, billing_mode, input_price, output_price, cache_write_price,
			 cache_write_1h_price, cache_read_price, per_request_price, enabled, created_at, updated_at)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,NOW(),NOW())
		RETURNING `+globalModelPricingColumns,
		pattern, mode,
		nullFloat64(in.InputPrice), nullFloat64(in.OutputPrice),
		nullFloat64(in.CacheWritePrice), nullFloat64(in.CacheWrite1hPrice),
		nullFloat64(in.CacheReadPrice), nullFloat64(in.PerRequestPrice),
		in.Enabled)
	if err != nil {
		if isUniqueConstraintViolation(err) {
			return nil, fmt.Errorf("create global model pricing %q: %w", pattern, service.ErrGlobalModelPricingDuplicate)
		}
		return nil, fmt.Errorf("create global model pricing: %w", err)
	}
	defer func() { _ = rows.Close() }()
	if !rows.Next() {
		return nil, fmt.Errorf("create global model pricing: no row returned")
	}
	p, err := scanGlobalModelPricingRow(rows.Scan)
	if err != nil {
		return nil, err
	}
	return p, rows.Err()
}

func (r *globalModelPricingRepository) GetByID(ctx context.Context, id int64) (*service.GlobalModelPrice, error) {
	rows, err := r.sql.QueryContext(ctx,
		`SELECT `+globalModelPricingColumns+` FROM global_model_pricing WHERE id = $1`, id)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()
	if !rows.Next() {
		return nil, sql.ErrNoRows
	}
	p, err := scanGlobalModelPricingRow(rows.Scan)
	if err != nil {
		return nil, err
	}
	return p, rows.Err()
}

func (r *globalModelPricingRepository) Update(ctx context.Context, id int64, in service.GlobalModelPriceInput) (*service.GlobalModelPrice, error) {
	pattern := strings.ToLower(strings.TrimSpace(in.ModelPattern))
	mode := strings.TrimSpace(string(in.BillingMode))
	if mode == "" {
		mode = string(service.BillingModeToken)
	}
	rows, err := r.sql.QueryContext(ctx, `
		UPDATE global_model_pricing SET
			model_pattern = $2, billing_mode = $3,
			input_price = $4, output_price = $5,
			cache_write_price = $6, cache_write_1h_price = $7, cache_read_price = $8,
			per_request_price = $9, enabled = $10, updated_at = NOW()
		WHERE id = $1
		RETURNING `+globalModelPricingColumns, id, pattern, mode,
		nullFloat64(in.InputPrice), nullFloat64(in.OutputPrice),
		nullFloat64(in.CacheWritePrice), nullFloat64(in.CacheWrite1hPrice),
		nullFloat64(in.CacheReadPrice), nullFloat64(in.PerRequestPrice),
		in.Enabled)
	if err != nil {
		if isUniqueConstraintViolation(err) {
			return nil, fmt.Errorf("update global model pricing %q: %w", pattern, service.ErrGlobalModelPricingDuplicate)
		}
		return nil, fmt.Errorf("update global model pricing: %w", err)
	}
	defer func() { _ = rows.Close() }()
	if !rows.Next() {
		return nil, sql.ErrNoRows
	}
	p, err := scanGlobalModelPricingRow(rows.Scan)
	if err != nil {
		return nil, err
	}
	return p, rows.Err()
}

func (r *globalModelPricingRepository) Delete(ctx context.Context, id int64) error {
	rows, err := r.sql.QueryContext(ctx, `DELETE FROM global_model_pricing WHERE id = $1 RETURNING id`, id)
	if err != nil {
		return err
	}
	defer func() { _ = rows.Close() }()
	if !rows.Next() {
		return sql.ErrNoRows
	}
	return rows.Err()
}

func (r *globalModelPricingRepository) SetEnabled(ctx context.Context, id int64, enabled bool) error {
	rows, err := r.sql.QueryContext(ctx,
		`UPDATE global_model_pricing SET enabled = $2, updated_at = NOW() WHERE id = $1 RETURNING id`, id, enabled)
	if err != nil {
		return err
	}
	defer func() { _ = rows.Close() }()
	if !rows.Next() {
		return sql.ErrNoRows
	}
	return rows.Err()
}