story/internal/billing/usage.go

package billing

import (
	"context"
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	"data-terminal-cloud/internal/deviceauth"
	"data-terminal-cloud/internal/identity"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

var (
	ErrUsageConflict     = errors.New("usage event key was reused with different content")
	ErrPriceRateNotFound = errors.New("published price rate was not found for usage")
	ErrQuotaExceeded     = errors.New("usage quota was exceeded")
)

var quantityPattern = regexp.MustCompile(`^(0|[1-9][0-9]{0,17})(\.[0-9]{1,6})?$`)

type UsageMetricInput struct {
	UsageCategory string `json:"usage_category"`
	MetricCode    string `json:"metric_code"`
	Quantity      string `json:"quantity"`
	Unit          string `json:"unit"`
}

type RecordedUsage struct {
	UsageEventID string `json:"usage_event_id"`
	CostLineID   string `json:"cost_line_id"`
	CostMicros   int64  `json:"cost_micros"`
	Currency     string `json:"currency"`
}

type CostSummaryLine struct {
	UsageCategory string `json:"usage_category"`
	MetricCode    string `json:"metric_code"`
	Quantity      string `json:"quantity"`
	Unit          string `json:"unit"`
	CostMicros    int64  `json:"cost_micros"`
	Currency      string `json:"currency"`
}

type CostSummary struct {
	From       time.Time         `json:"from"`
	To         time.Time         `json:"to"`
	CostMicros int64             `json:"cost_micros"`
	Currency   string            `json:"currency"`
	Lines      []CostSummaryLine `json:"lines"`
}

func RecordProviderUsageTx(ctx context.Context, tx pgx.Tx, tenantID, conversationID, providerEventID, provider, sourceEventKey string, occurredAt time.Time, metrics []UsageMetricInput) ([]RecordedUsage, error) {
	provider = strings.ToLower(strings.TrimSpace(provider))
	sourceEventKey = strings.TrimSpace(sourceEventKey)
	if provider == "" || sourceEventKey == "" || occurredAt.IsZero() || len(metrics) == 0 || len(metrics) > 100 {
		return nil, ErrInvalid
	}
	seen := map[string]struct{}{}
	for index := range metrics {
		metrics[index].UsageCategory = strings.ToUpper(strings.TrimSpace(metrics[index].UsageCategory))
		metrics[index].MetricCode = strings.ToLower(strings.TrimSpace(metrics[index].MetricCode))
		metrics[index].Quantity = strings.TrimSpace(metrics[index].Quantity)
		metrics[index].Unit = strings.ToLower(strings.TrimSpace(metrics[index].Unit))
		key := metrics[index].UsageCategory + ":" + metrics[index].MetricCode + ":" + metrics[index].Unit
		if !validCategory(metrics[index].UsageCategory) || !codePattern.MatchString(metrics[index].MetricCode) || !regexp.MustCompile(`^[a-z][a-z0-9_.-]{1,39}$`).MatchString(metrics[index].Unit) || !quantityPattern.MatchString(metrics[index].Quantity) {
			return nil, ErrInvalid
		}
		if _, exists := seen[key]; exists {
			return nil, ErrInvalid
		}
		seen[key] = struct{}{}
	}
	result := make([]RecordedUsage, 0, len(metrics))
	for _, metric := range metrics {
		usagePublicID := uuid.NewString()
		var usageInternalID string
		tag, err := tx.Exec(ctx, `
			INSERT INTO conversation_usage_event(public_id,tenant_id,conversation_id,provider_event_id,provider,source_event_key,usage_category,metric_code,quantity,unit,occurred_at)
			VALUES($1::uuid,$2::uuid,$3::uuid,$4::uuid,$5,$6,$7,$8,$9::numeric,$10,$11)
			ON CONFLICT(tenant_id,provider,source_event_key,usage_category,metric_code,unit) DO NOTHING`,
			usagePublicID, tenantID, conversationID, providerEventID, provider, sourceEventKey, metric.UsageCategory, metric.MetricCode, metric.Quantity, metric.Unit, occurredAt.UTC())
		if err != nil {
			return nil, err
		}
		if tag.RowsAffected() == 0 {
			var quantityMatches, occurredAtMatches bool
			if err = tx.QueryRow(ctx, `SELECT id::text,quantity=$7::numeric,occurred_at=$8 FROM conversation_usage_event WHERE tenant_id=$1::uuid AND provider=$2 AND source_event_key=$3 AND usage_category=$4 AND metric_code=$5 AND unit=$6`, tenantID, provider, sourceEventKey, metric.UsageCategory, metric.MetricCode, metric.Unit, metric.Quantity, occurredAt.UTC()).Scan(&usageInternalID, &quantityMatches, &occurredAtMatches); err != nil {
				return nil, err
			}
			if !quantityMatches || !occurredAtMatches {
				return nil, ErrUsageConflict
			}
			var recorded RecordedUsage
			err = tx.QueryRow(ctx, `SELECT cue.public_id::text,ccl.public_id::text,ccl.cost_micros,ccl.currency FROM conversation_usage_event cue JOIN conversation_cost_line ccl ON ccl.usage_event_id=cue.id WHERE cue.id=$1::uuid`, usageInternalID).Scan(&recorded.UsageEventID, &recorded.CostLineID, &recorded.CostMicros, &recorded.Currency)
			if err != nil {
				return nil, err
			}
			result = append(result, recorded)
			continue
		}
		if err = tx.QueryRow(ctx, `SELECT id::text FROM conversation_usage_event WHERE public_id=$1::uuid`, usagePublicID).Scan(&usageInternalID); err != nil {
			return nil, err
		}
		var priceBookID, priceRateID, currency string
		var unitPriceMicros int64
		err = tx.QueryRow(ctx, `
			SELECT pb.id::text,pr.id::text,pr.unit_price_micros,pb.currency
			FROM billing_price_book pb JOIN billing_price_rate pr ON pr.price_book_id=pb.id
			WHERE pb.tenant_id=$1::uuid AND pb.provider=$2 AND pb.status IN ('PUBLISHED','RETIRED')
			  AND pb.published_at IS NOT NULL AND pb.effective_from<=$3
			  AND pr.usage_category=$4 AND pr.metric_code=$5 AND pr.unit=$6
			ORDER BY pb.effective_from DESC,pb.version_no DESC LIMIT 1`, tenantID, provider, occurredAt.UTC(), metric.UsageCategory, metric.MetricCode, metric.Unit).Scan(&priceBookID, &priceRateID, &unitPriceMicros, &currency)
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrPriceRateNotFound
		}
		if err != nil {
			return nil, err
		}
		costPublicID := uuid.NewString()
		var costMicros int64
		if err = tx.QueryRow(ctx, `INSERT INTO conversation_cost_line(public_id,tenant_id,conversation_id,usage_event_id,price_book_id,price_rate_id,quantity,unit,unit_price_micros,cost_micros,currency) VALUES($1::uuid,$2::uuid,$3::uuid,$4::uuid,$5::uuid,$6::uuid,$7::numeric,$8,$9::bigint,round($7::numeric*$9::numeric)::bigint,$10) RETURNING cost_micros`, costPublicID, tenantID, conversationID, usageInternalID, priceBookID, priceRateID, metric.Quantity, metric.Unit, unitPriceMicros, currency).Scan(&costMicros); err != nil {
			return nil, err
		}
		result = append(result, RecordedUsage{UsageEventID: usagePublicID, CostLineID: costPublicID, CostMicros: costMicros, Currency: currency})
	}
	return result, nil
}

func (s *Service) CostSummary(ctx context.Context, principal identity.Principal, tenantPublicID string, from, to time.Time) (CostSummary, error) {
	currentActor, err := s.resolveActor(ctx, principal, tenantPublicID, "ops_viewer", "ops_operator", "ops_approver", "ops_admin")
	if err != nil {
		return CostSummary{}, err
	}
	from = from.UTC()
	to = to.UTC()
	if from.IsZero() || !to.After(from) || to.Sub(from) > 366*24*time.Hour || to.After(time.Now().UTC().Add(5*time.Minute)) {
		return CostSummary{}, ErrInvalid
	}
	rows, err := s.pool.Query(ctx, `
		SELECT cue.usage_category,cue.metric_code,SUM(cue.quantity)::text,cue.unit,SUM(ccl.cost_micros)::bigint,ccl.currency
		FROM conversation_usage_event cue JOIN conversation_cost_line ccl ON ccl.usage_event_id=cue.id
		WHERE cue.tenant_id=$1::uuid AND cue.occurred_at>=$2 AND cue.occurred_at<$3
		GROUP BY cue.usage_category,cue.metric_code,cue.unit,ccl.currency
		ORDER BY cue.usage_category,cue.metric_code,cue.unit,ccl.currency`, currentActor.tenantID, from, to)
	if err != nil {
		return CostSummary{}, err
	}
	defer rows.Close()
	result := CostSummary{From: from, To: to, Lines: []CostSummaryLine{}}
	currencies := map[string]struct{}{}
	for rows.Next() {
		var item CostSummaryLine
		if err = rows.Scan(&item.UsageCategory, &item.MetricCode, &item.Quantity, &item.Unit, &item.CostMicros, &item.Currency); err != nil {
			return CostSummary{}, err
		}
		result.CostMicros += item.CostMicros
		currencies[item.Currency] = struct{}{}
		result.Lines = append(result.Lines, item)
	}
	if err = rows.Err(); err != nil {
		return CostSummary{}, err
	}
	if len(currencies) == 1 {
		for currency := range currencies {
			result.Currency = currency
		}
	} else if len(currencies) > 1 {
		return CostSummary{}, fmt.Errorf("cost summary contains multiple currencies")
	}
	return result, nil
}

type BudgetSnapshot struct {
	TenantPolicyID       *string
	DevicePolicyID       *string
	BillingTimezone      string
	Currency             string
	DailyLimitMicros     int64
	MonthlyLimitMicros   int64
	MaxCallDurationSec   int
	InProgressAction     string
	DailySpendAtCreate   int64
	MonthlySpendAtCreate int64
}

type quotaRow struct {
	id, timezone, currency, action string
	daily, monthly                 int64
	concurrent, duration           int
}

func CheckQuotaTx(ctx context.Context, tx pgx.Tx, identity deviceauth.Identity, now time.Time) (*BudgetSnapshot, error) {
	// New and legacy call admissions share the same tenant concurrency boundary.
	if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1,19))`, identity.TenantID); err != nil {
		return nil, err
	}
	tenantPolicy, err := readActiveQuota(ctx, tx, identity.TenantID, "TENANT", nil)
	if err != nil && !errors.Is(err, pgx.ErrNoRows) {
		return nil, err
	}
	hasTenantPolicy := err == nil
	var deviceID = identity.ID
	devicePolicy, deviceErr := readActiveQuota(ctx, tx, identity.TenantID, "DEVICE", &deviceID)
	if deviceErr != nil && !errors.Is(deviceErr, pgx.ErrNoRows) {
		return nil, deviceErr
	}
	hasDevicePolicy := deviceErr == nil
	if !hasTenantPolicy && !hasDevicePolicy {
		return nil, nil
	}
	policies := []quotaRow{}
	if hasTenantPolicy {
		policies = append(policies, tenantPolicy)
	}
	if hasDevicePolicy {
		policies = append(policies, devicePolicy)
	}
	if len(policies) == 2 && (policies[0].timezone != policies[1].timezone || policies[0].currency != policies[1].currency) {
		return nil, fmt.Errorf("tenant and device quota policy boundaries are inconsistent")
	}
	effective := policies[0]
	for _, policy := range policies {
		if policy.daily < effective.daily {
			effective.daily = policy.daily
		}
		if policy.monthly < effective.monthly {
			effective.monthly = policy.monthly
		}
		if policy.duration < effective.duration {
			effective.duration = policy.duration
		}
		if policy.action == "STOP_IMMEDIATELY" {
			effective.action = policy.action
		}
	}
	location, err := time.LoadLocation(effective.timezone)
	if err != nil {
		return nil, err
	}
	localNow := now.In(location)
	dayStart := time.Date(localNow.Year(), localNow.Month(), localNow.Day(), 0, 0, 0, 0, location).UTC()
	monthStart := time.Date(localNow.Year(), localNow.Month(), 1, 0, 0, 0, 0, location).UTC()
	var tenantDaily, tenantMonthly, deviceDaily, deviceMonthly int64
	if err = tx.QueryRow(ctx, `SELECT COALESCE(SUM(ccl.cost_micros) FILTER(WHERE cue.occurred_at>=$2),0)::bigint,COALESCE(SUM(ccl.cost_micros) FILTER(WHERE cue.occurred_at>=$3),0)::bigint FROM conversation_cost_line ccl JOIN conversation_usage_event cue ON cue.id=ccl.usage_event_id WHERE ccl.tenant_id=$1::uuid AND ccl.currency=$4`, identity.TenantID, dayStart, monthStart, effective.currency).Scan(&tenantDaily, &tenantMonthly); err != nil {
		return nil, err
	}
	if err = tx.QueryRow(ctx, `SELECT COALESCE(SUM(ccl.cost_micros) FILTER(WHERE cue.occurred_at>=$3),0)::bigint,COALESCE(SUM(ccl.cost_micros) FILTER(WHERE cue.occurred_at>=$4),0)::bigint FROM conversation_cost_line ccl JOIN conversation_usage_event cue ON cue.id=ccl.usage_event_id JOIN conversation c ON c.id=ccl.conversation_id WHERE ccl.tenant_id=$1::uuid AND c.device_id=$2::uuid AND ccl.currency=$5`, identity.TenantID, identity.ID, dayStart, monthStart, effective.currency).Scan(&deviceDaily, &deviceMonthly); err != nil {
		return nil, err
	}
	if hasTenantPolicy && (tenantDaily >= tenantPolicy.daily || tenantMonthly >= tenantPolicy.monthly) {
		return nil, ErrQuotaExceeded
	}
	if hasDevicePolicy && (deviceDaily >= devicePolicy.daily || deviceMonthly >= devicePolicy.monthly) {
		return nil, ErrQuotaExceeded
	}
	for _, policy := range policies {
		var active int
		query := `SELECT count(*) FROM (SELECT tenant_id,device_id FROM conversation WHERE status IN ('PREPARING','ROOM_ISSUED','STARTING_AGENT','ACTIVE','STOPPING') UNION ALL SELECT tenant_id,device_id FROM call_session WHERE manifest_status IN ('ISSUED','REVOKED') AND NOT(local_audio_off AND sdk_stopped AND remote_state IN ('NOT_STARTED','STOPPED_OBSERVED'))) admitted WHERE tenant_id=$1::uuid`
		args := []any{identity.TenantID}
		if hasDevicePolicy && policy.id == devicePolicy.id {
			query += ` AND device_id=$2::uuid`
			args = append(args, identity.ID)
		}
		if err = tx.QueryRow(ctx, query, args...).Scan(&active); err != nil {
			return nil, err
		}
		if active >= policy.concurrent {
			return nil, ErrQuotaExceeded
		}
	}
	snapshot := &BudgetSnapshot{BillingTimezone: effective.timezone, Currency: effective.currency, DailyLimitMicros: effective.daily, MonthlyLimitMicros: effective.monthly, MaxCallDurationSec: effective.duration, InProgressAction: effective.action, DailySpendAtCreate: tenantDaily, MonthlySpendAtCreate: tenantMonthly}
	if hasTenantPolicy {
		snapshot.TenantPolicyID = &tenantPolicy.id
	}
	if hasDevicePolicy {
		snapshot.DevicePolicyID = &devicePolicy.id
		if deviceDaily > snapshot.DailySpendAtCreate {
			snapshot.DailySpendAtCreate = deviceDaily
		}
		if deviceMonthly > snapshot.MonthlySpendAtCreate {
			snapshot.MonthlySpendAtCreate = deviceMonthly
		}
	}
	return snapshot, nil
}

func FreezeBudgetTx(ctx context.Context, tx pgx.Tx, conversationID, tenantID string, snapshot *BudgetSnapshot) error {
	if snapshot == nil {
		return nil
	}
	_, err := tx.Exec(ctx, `INSERT INTO conversation_usage_budget(conversation_id,tenant_id,tenant_policy_id,device_policy_id,billing_timezone,currency,daily_limit_micros,monthly_limit_micros,max_call_duration_sec,in_progress_action,daily_spend_at_create,monthly_spend_at_create) VALUES($1::uuid,$2::uuid,$3::uuid,$4::uuid,$5,$6,$7,$8,$9,$10,$11,$12)`, conversationID, tenantID, snapshot.TenantPolicyID, snapshot.DevicePolicyID, snapshot.BillingTimezone, snapshot.Currency, snapshot.DailyLimitMicros, snapshot.MonthlyLimitMicros, snapshot.MaxCallDurationSec, snapshot.InProgressAction, snapshot.DailySpendAtCreate, snapshot.MonthlySpendAtCreate)
	return err
}

func readActiveQuota(ctx context.Context, tx pgx.Tx, tenantID, scope string, deviceID *string) (quotaRow, error) {
	var item quotaRow
	err := tx.QueryRow(ctx, `SELECT id::text,billing_timezone,currency,daily_limit_micros,monthly_limit_micros,max_concurrent_calls,max_call_duration_sec,in_progress_action FROM usage_quota_policy WHERE tenant_id=$1::uuid AND scope_type=$2 AND status='ACTIVE' AND (($2='TENANT') OR device_id=$3::uuid)`, tenantID, scope, deviceID).Scan(&item.id, &item.timezone, &item.currency, &item.daily, &item.monthly, &item.concurrent, &item.duration, &item.action)
	return item, err
}