path: root/internal/billing/service.go

package billing

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"strings"
	"time"

	"aigw/internal/domain"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

const microsPerUnit = int64(1_000_000)

type Service struct {
	db                     *pgxpool.Pool
	currency               string
	defaultMaxOutputTokens int64
	minTopUpMinor          int64
	maxTopUpMinor          int64
	stripeEnabled          bool
	stripeWebhookSecret    string
	stripeSuccessURL       string
	stripeCancelURL        string
	integrationIdentifier  string
	createStripeCheckout   stripeCheckoutCreator
}

func New(ctx context.Context, options Options) (*Service, error) {
	db, err := pgxpool.New(ctx, options.DatabaseURL)
	if err != nil {
		return nil, fmt.Errorf("configure billing PostgreSQL: %w", err)
	}
	if err := db.Ping(ctx); err != nil {
		db.Close()
		return nil, fmt.Errorf("connect billing PostgreSQL: %w", err)
	}
	service := &Service{
		db: db, currency: options.Currency, defaultMaxOutputTokens: options.DefaultMaxOutputTokens,
		minTopUpMinor: options.MinTopUpMinor, maxTopUpMinor: options.MaxTopUpMinor,
		stripeEnabled: options.StripeEnabled, stripeWebhookSecret: options.StripeWebhookSecret,
		stripeSuccessURL: options.StripeSuccessURL, stripeCancelURL: options.StripeCancelURL,
		integrationIdentifier: "aigw_balance_" + randomLetters(8),
	}
	if options.StripeEnabled {
		service.createStripeCheckout = newStripeCheckoutCreator(options.StripeAPIKey)
	}
	return service, nil
}

func (s *Service) Close() {
	s.db.Close()
}

func (s *Service) StripeEnabled() bool {
	return s.stripeEnabled
}

func (s *Service) Currency() string {
	return s.currency
}

func (s *Service) Authorize(ctx context.Context, input Authorization) error {
	if input.RequestID == "" || input.Principal.TenantID == "" || input.Principal.ProjectID == "" || input.Principal.KeyID == "" {
		return errors.New("billing authorization identity is incomplete")
	}
	reserved, err := reservationCost(input.Model, input.Body, s.defaultMaxOutputTokens)
	if err != nil {
		return err
	}
	tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted})
	if err != nil {
		return fmt.Errorf("begin billing authorization: %w", err)
	}
	defer tx.Rollback(ctx)
	if _, err := tx.Exec(ctx, `
		INSERT INTO tenant_wallets (tenant_id, currency) VALUES ($1, $2)
		ON CONFLICT (tenant_id) DO NOTHING`, input.Principal.TenantID, s.currency); err != nil {
		return fmt.Errorf("ensure tenant wallet: %w", err)
	}
	var currency string
	var balance, held int64
	if err := tx.QueryRow(ctx, `
		SELECT currency, balance_micros, reserved_micros FROM tenant_wallets
		WHERE tenant_id = $1 FOR UPDATE`, input.Principal.TenantID).Scan(&currency, &balance, &held); err != nil {
		return fmt.Errorf("lock tenant wallet: %w", err)
	}
	if currency != s.currency {
		return fmt.Errorf("tenant wallet currency %s does not match billing currency %s", currency, s.currency)
	}
	if input.Policy.MonthlySpendMicros > 0 {
		period := time.Date(time.Now().UTC().Year(), time.Now().UTC().Month(), 1, 0, 0, 0, 0, time.UTC)
		nextPeriod := period.AddDate(0, 1, 0)
		var used, pending int64
		if err := tx.QueryRow(ctx, `SELECT
			COALESCE((SELECT cost_micros FROM usage_monthly_rollups WHERE project_id=$1 AND period_start=$2),0),
			COALESCE((SELECT sum(reserved_micros) FROM billing_reservations WHERE project_id=$1 AND status='pending' AND created_at >= $2 AND created_at < $3),0)`,
			input.Principal.ProjectID, period, nextPeriod).Scan(&used, &pending); err != nil {
			return fmt.Errorf("read monthly spend quota: %w", err)
		}
		if reserved > input.Policy.MonthlySpendMicros || used > input.Policy.MonthlySpendMicros-reserved || pending > input.Policy.MonthlySpendMicros-used-reserved {
			return ErrQuotaExceeded
		}
	}
	if balance-held < reserved {
		return ErrInsufficientBalance
	}
	if _, err := tx.Exec(ctx, `
		INSERT INTO billing_reservations (
			request_id, tenant_id, project_id, key_id, public_model, currency, reserved_micros,
			input_price_micros_per_million, output_price_micros_per_million,
			cache_read_price_micros_per_million, cache_write_price_micros_per_million)
		VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11)`,
		input.RequestID, input.Principal.TenantID, input.Principal.ProjectID, input.Principal.KeyID,
		input.Model.ID, s.currency, reserved, input.Model.InputPriceMicrosPerMillion,
		input.Model.OutputPriceMicrosPerMillion, input.Model.CacheReadPriceMicrosPerMillion,
		input.Model.CacheWritePriceMicrosPerMillion); err != nil {
		return fmt.Errorf("create billing reservation: %w", err)
	}
	if _, err := tx.Exec(ctx, `
		UPDATE tenant_wallets SET reserved_micros = reserved_micros + $2, updated_at = now()
		WHERE tenant_id = $1`, input.Principal.TenantID, reserved); err != nil {
		return fmt.Errorf("reserve tenant balance: %w", err)
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("commit billing authorization: %w", err)
	}
	return nil
}

func (s *Service) Settle(ctx context.Context, event domain.UsageEvent) error {
	tx, err := s.db.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted})
	if err != nil {
		return fmt.Errorf("begin usage settlement: %w", err)
	}
	defer tx.Rollback(ctx)

	var tenantID, projectID, keyID, modelID, currency, status string
	var reserved, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64
	err = tx.QueryRow(ctx, `
		SELECT tenant_id::text, project_id::text, key_id::text, public_model, currency, reserved_micros, status,
		       input_price_micros_per_million, output_price_micros_per_million,
		       cache_read_price_micros_per_million, cache_write_price_micros_per_million
		FROM billing_reservations WHERE request_id = $1 FOR UPDATE`, event.RequestID,
	).Scan(&tenantID, &projectID, &keyID, &modelID, &currency, &reserved, &status,
		&inputPrice, &outputPrice, &cacheReadPrice, &cacheWritePrice)
	if err != nil {
		return fmt.Errorf("load billing reservation: %w", err)
	}
	if status != "pending" {
		return tx.Commit(ctx)
	}

	actualCost := int64(0)
	if event.StatusCode >= 200 && event.StatusCode < 300 {
		actualCost, err = usageCost(event.Usage, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice)
		if err != nil {
			return err
		}
	}
	var balance, held int64
	if err := tx.QueryRow(ctx, `SELECT balance_micros, reserved_micros FROM tenant_wallets WHERE tenant_id = $1 FOR UPDATE`, tenantID).Scan(&balance, &held); err != nil {
		return fmt.Errorf("lock wallet for settlement: %w", err)
	}
	charged, err := collectibleCharge(actualCost, balance, held, reserved)
	if err != nil {
		return err
	}
	uncollected := actualCost - charged
	newBalance := balance - charged
	if _, err := tx.Exec(ctx, `
		UPDATE tenant_wallets SET balance_micros = $2, reserved_micros = reserved_micros - $3, updated_at = now()
		WHERE tenant_id = $1`, tenantID, newBalance, reserved); err != nil {
		return fmt.Errorf("settle tenant wallet: %w", err)
	}
	reservationStatus := "released"
	if actualCost > 0 {
		reservationStatus = "settled"
	}
	if _, err := tx.Exec(ctx, `
		UPDATE billing_reservations
		SET status = $2, actual_cost_micros = $3, charged_micros = $4, uncollected_micros = $5, settled_at = now()
		WHERE request_id = $1`, event.RequestID, reservationStatus, actualCost, charged, uncollected); err != nil {
		return fmt.Errorf("update billing reservation: %w", err)
	}
	usageAlreadyRecorded := false
	if err := tx.QueryRow(ctx, `SELECT true FROM usage_events WHERE request_id=$1 FOR UPDATE`, event.RequestID).Scan(&usageAlreadyRecorded); err != nil && !errors.Is(err, pgx.ErrNoRows) {
		return fmt.Errorf("lock existing usage event: %w", err)
	}
	if usageAlreadyRecorded {
		if _, err := tx.Exec(ctx, `UPDATE usage_events SET cost_micros=$2, charged_micros=$3, uncollected_micros=$4 WHERE request_id=$1`,
			event.RequestID, actualCost, charged, uncollected); err != nil {
			return fmt.Errorf("apply usage charge: %w", err)
		}
	} else if _, err := tx.Exec(ctx, `
		INSERT INTO usage_events (
			request_id, tenant_id, project_id, key_id, public_model, provider_id, upstream_model,
			protocol, stream, status_code, success, error_type, attempts, started_at, duration_ms,
			input_tokens, output_tokens, total_tokens, cache_creation_input_tokens, cache_read_input_tokens,
			cost_micros, charged_micros, uncollected_micros)
		VALUES ($1,$2,$3,$4,$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23)
		ON CONFLICT (request_id) DO NOTHING`,
		event.RequestID, tenantID, projectID, keyID, modelID, event.ProviderID, event.UpstreamModel,
		string(event.Protocol), event.Stream, event.StatusCode, event.Success, event.ErrorType, event.Attempts,
		event.StartedAt, event.DurationMS, event.Usage.InputTokens, event.Usage.OutputTokens,
		event.Usage.TotalTokens, event.Usage.CacheCreationInputTokens, event.Usage.CacheReadInputTokens,
		actualCost, charged, uncollected); err != nil {
		return fmt.Errorf("persist usage event: %w", err)
	}
	if charged > 0 {
		if _, err := tx.Exec(ctx, `
			INSERT INTO billing_ledger (tenant_id, project_id, currency, amount_micros, balance_after_micros, kind, source_type, source_id, description)
			VALUES ($1,$2,$3,$4,$5,'usage','request',$6,$7)
			ON CONFLICT (source_type, source_id) DO NOTHING`, tenantID, projectID, currency, -charged, newBalance, event.RequestID, modelID); err != nil {
			return fmt.Errorf("write usage ledger entry: %w", err)
		}
	}
	period := time.Date(event.StartedAt.UTC().Year(), event.StartedAt.UTC().Month(), 1, 0, 0, 0, 0, time.UTC)
	if usageAlreadyRecorded {
		if _, err := tx.Exec(ctx, `UPDATE usage_monthly_rollups SET cost_micros=cost_micros+$3,
			charged_micros=charged_micros+$4, uncollected_micros=uncollected_micros+$5, updated_at=now()
			WHERE project_id=$1 AND period_start=$2`, projectID, period, actualCost, charged, uncollected); err != nil {
			return fmt.Errorf("apply usage rollup charge: %w", err)
		}
	} else if _, err := tx.Exec(ctx, `
		INSERT INTO usage_monthly_rollups
		(period_start, tenant_id, project_id, request_count, successful_requests, input_tokens, output_tokens, total_tokens, cost_micros, charged_micros, uncollected_micros)
		VALUES ($1,$2,$3,1,$4,$5,$6,$7,$8,$9,$10)
		ON CONFLICT (project_id, period_start) DO UPDATE SET request_count=usage_monthly_rollups.request_count+1,
			successful_requests=usage_monthly_rollups.successful_requests+EXCLUDED.successful_requests,
			input_tokens=usage_monthly_rollups.input_tokens+EXCLUDED.input_tokens,
			output_tokens=usage_monthly_rollups.output_tokens+EXCLUDED.output_tokens,
			total_tokens=usage_monthly_rollups.total_tokens+EXCLUDED.total_tokens,
			cost_micros=usage_monthly_rollups.cost_micros+EXCLUDED.cost_micros,
			charged_micros=usage_monthly_rollups.charged_micros+EXCLUDED.charged_micros,
			uncollected_micros=usage_monthly_rollups.uncollected_micros+EXCLUDED.uncollected_micros,
			updated_at=now()`, period, tenantID, projectID, boolToInt(event.Success), event.Usage.InputTokens,
		event.Usage.OutputTokens, event.Usage.TotalTokens, actualCost, charged, uncollected); err != nil {
		return fmt.Errorf("update usage monthly rollup: %w", err)
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("commit usage settlement: %w", err)
	}
	return nil
}

func boolToInt(value bool) int {
	if value {
		return 1
	}
	return 0
}

func reservationCost(model domain.Model, body []byte, defaultMaxOutput int64) (int64, error) {
	maxOutput := defaultMaxOutput
	var limits struct {
		MaxTokens           int64 `json:"max_tokens"`
		MaxCompletionTokens int64 `json:"max_completion_tokens"`
		MaxOutputTokens     int64 `json:"max_output_tokens"`
	}
	if json.Unmarshal(body, &limits) == nil {
		explicitMax := int64(0)
		for _, value := range []int64{limits.MaxTokens, limits.MaxCompletionTokens, limits.MaxOutputTokens} {
			if value > explicitMax {
				explicitMax = value
			}
		}
		if explicitMax > 0 {
			maxOutput = explicitMax
		}
	}
	cacheReservePrice := model.CacheReadPriceMicrosPerMillion
	if model.CacheWritePriceMicrosPerMillion > cacheReservePrice {
		cacheReservePrice = model.CacheWritePriceMicrosPerMillion
	}
	return calculateCost(int64(len(body)), maxOutput, 0, int64(len(body)),
		model.InputPriceMicrosPerMillion, model.OutputPriceMicrosPerMillion,
		0, cacheReservePrice)
}

func usageCost(usage domain.Usage, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64) (int64, error) {
	return calculateCost(usage.InputTokens, usage.OutputTokens, usage.CacheReadInputTokens,
		usage.CacheCreationInputTokens, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice)
}

func calculateCost(input, output, cacheRead, cacheWrite, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice int64) (int64, error) {
	values := []int64{input, output, cacheRead, cacheWrite, inputPrice, outputPrice, cacheReadPrice, cacheWritePrice}
	for _, value := range values {
		if value < 0 {
			return 0, errors.New("billing values cannot be negative")
		}
	}
	total := new(big.Int)
	for _, pair := range [][2]int64{{input, inputPrice}, {output, outputPrice}, {cacheRead, cacheReadPrice}, {cacheWrite, cacheWritePrice}} {
		total.Add(total, new(big.Int).Mul(big.NewInt(pair[0]), big.NewInt(pair[1])))
	}
	if total.Sign() == 0 {
		return 0, nil
	}
	total.Add(total, big.NewInt(microsPerUnit-1))
	total.Div(total, big.NewInt(microsPerUnit))
	if !total.IsInt64() {
		return 0, errors.New("calculated charge exceeds supported range")
	}
	return total.Int64(), nil
}

func collectibleCharge(actualCost, balance, held, reservation int64) (int64, error) {
	if actualCost < 0 || balance < 0 || held < 0 || reservation < 0 || held > balance || reservation > held {
		return 0, errors.New("wallet reservation invariant violated")
	}
	spendable := balance - (held - reservation)
	if actualCost > spendable {
		return spendable, nil
	}
	return actualCost, nil
}

func randomHex(bytes int) string {
	buffer := make([]byte, bytes)
	if _, err := rand.Read(buffer); err != nil {
		return fmt.Sprintf("%08x", time.Now().UnixNano())[:bytes*2]
	}
	return hex.EncodeToString(buffer)
}

func randomLetters(length int) string {
	const letters = "abcdefghijklmnopqrstuvwxyz"
	buffer := make([]byte, length)
	if _, err := rand.Read(buffer); err != nil {
		return strings.Repeat("a", length)
	}
	for i := range buffer {
		buffer[i] = letters[int(buffer[i])%len(letters)]
	}
	return string(buffer)
}

func normalizeDescription(value string) string {
	value = strings.TrimSpace(value)
	if len(value) > 240 {
		value = value[:240]
	}
	return value
}