descrybe/apps/api/internal/aiaudit/backfill.go

package aiaudit

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

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

// Recovering cost for products processed before capture existed.
//
// processed_products.gpt_response keeps the provider's usage block for some rows
// (model + prompt/completion tokens). That is enough to reconstruct spend per
// company and day, so history is not simply lost. It is explicitly best effort:
//
//   - Only rows that still carry a usage block are recoverable. gpt_response is
//     overwritten on every reprocess and absent on older/failed rows.
//   - There is no per-user attribution: processed_products.user_id is null on these
//     rows, so backfilled spend lands under "Unattributed" in the by-user view.
//   - The day is processed_products.updated_at, which is the LAST write to the row,
//     not necessarily when the call was made.
//   - Cached-token discounts cannot be recovered (the field did not exist), so a
//     backfilled call is priced as if none of its input was cached — an
//     over-estimate, never an under-estimate.
//
// Rows are written with source='backfill' so they never collide with live capture
// and a re-run replaces rather than doubles them.

// BackfillResult reports what a backfill pass recovered.
type BackfillResult struct {
	ProductsScanned  int    `json:"products_scanned"`
	ProductsUsed     int    `json:"products_used"`
	ProductsNoUsage  int    `json:"products_without_usage"`
	CallsRecovered   int    `json:"calls_recovered"`
	PromptTokens     int64  `json:"prompt_tokens"`
	CompletionTokens int64  `json:"completion_tokens"`
	CostMicros       int64  `json:"cost_micros"`
	RowsWritten      int    `json:"rows_written"`
	RowsDeleted      int64  `json:"rows_deleted"`
	Cutoff           string `json:"cutoff"`
	DryRun           bool   `json:"dry_run"`
}

// CostUSD renders the recovered spend for CLI output.
func (r BackfillResult) CostUSD() float64 {
	return float64(r.CostMicros) / float64(MicrosPerUSD)
}

type backfillKey struct {
	day       time.Time
	companyID uuid.UUID
	model     string
	role      string
}

type backfillAgg struct {
	calls            int
	promptTokens     int64
	completionTokens int64
	totalTokens      int64
	costMicros       int64
}

// LiveCaptureStart returns the earliest day live capture recorded, or zero when
// nothing has been captured yet. The backfill uses it as a default cutoff so a
// product processed after capture began is never counted twice.
func LiveCaptureStart(ctx context.Context, pool *pgxpool.Pool) (time.Time, error) {
	if pool == nil {
		return time.Time{}, nil
	}
	var day *time.Time
	err := pool.QueryRow(ctx, `
		SELECT min(day) FROM ai_usage_daily WHERE source = $1`, SourceLive).Scan(&day)
	if err != nil {
		return time.Time{}, err
	}
	if day == nil {
		return time.Time{}, nil
	}
	return *day, nil
}

// BackfillFromProcessedProducts reconstructs cost history from stored provider
// responses for products last written before cutoff. Passing a zero cutoff scans
// everything, which is only safe when live capture has recorded nothing yet.
func BackfillFromProcessedProducts(ctx context.Context, pool *pgxpool.Pool, cutoff time.Time, dryRun bool) (BackfillResult, error) {
	out := BackfillResult{DryRun: dryRun}
	if pool == nil {
		return out, fmt.Errorf("nil pool")
	}
	if cutoff.IsZero() {
		// Far future = no bound.
		cutoff = time.Now().UTC().AddDate(100, 0, 0)
	}
	out.Cutoff = cutoff.Format("2006-01-02")

	prices, err := loadPrices(ctx, pool)
	if err != nil {
		return out, err
	}

	rows, err := pool.Query(ctx, `
		SELECT company_id, updated_at, COALESCE(gpt_response, '{}'::jsonb)
		FROM processed_products
		WHERE gpt_response IS NOT NULL
		  AND updated_at < $1
		ORDER BY updated_at`, cutoff)
	if err != nil {
		return out, err
	}
	defer rows.Close()

	agg := map[backfillKey]*backfillAgg{}
	for rows.Next() {
		var companyID uuid.UUID
		var updatedAt time.Time
		var raw []byte
		if err := rows.Scan(&companyID, &updatedAt, &raw); err != nil {
			return out, err
		}
		out.ProductsScanned++
		calls := extractUsageCalls(raw)
		if len(calls) == 0 {
			out.ProductsNoUsage++
			continue
		}
		out.ProductsUsed++
		day := updatedAt.UTC().Truncate(24 * time.Hour)
		for _, c := range calls {
			price := priceFromMap(prices, c.Model)
			// Cached tokens are unknown for history — price all input as fresh.
			cost := CostMicros(price, c.PromptTokens, 0, c.CompletionTokens)
			k := backfillKey{day: day, companyID: companyID, model: NormalizeModel(c.Model), role: c.Role}
			a := agg[k]
			if a == nil {
				a = &backfillAgg{}
				agg[k] = a
			}
			a.calls++
			a.promptTokens += int64(c.PromptTokens)
			a.completionTokens += int64(c.CompletionTokens)
			a.totalTokens += int64(c.TotalTokens)
			a.costMicros += cost

			out.CallsRecovered++
			out.PromptTokens += int64(c.PromptTokens)
			out.CompletionTokens += int64(c.CompletionTokens)
			out.CostMicros += cost
		}
	}
	if err := rows.Err(); err != nil {
		return out, err
	}
	out.RowsWritten = len(agg)
	if dryRun {
		return out, nil
	}

	tx, err := pool.Begin(ctx)
	if err != nil {
		return out, err
	}
	defer func() { _ = tx.Rollback(ctx) }()

	// Replace rather than add: a re-run must not double the recovered history.
	tag, err := tx.Exec(ctx, `DELETE FROM ai_usage_daily WHERE source = $1`, SourceBackfill)
	if err != nil {
		return out, err
	}
	out.RowsDeleted = tag.RowsAffected()

	for k, a := range agg {
		if _, err := tx.Exec(ctx, `
			INSERT INTO ai_usage_daily (
				day, company_id, user_id, model, role, source,
				calls, failed_calls, prompt_tokens, cached_prompt_tokens,
				completion_tokens, total_tokens, cost_micros, updated_at)
			VALUES ($1, $2, '00000000-0000-0000-0000-000000000000', $3, $4, $5,
				$6, 0, $7, 0, $8, $9, $10, now())`,
			k.day, k.companyID, k.model, k.role, SourceBackfill,
			a.calls, a.promptTokens, a.completionTokens, a.totalTokens, a.costMicros); err != nil {
			return out, err
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return out, err
	}
	return out, nil
}

func loadPrices(ctx context.Context, pool *pgxpool.Pool) (map[string]ModelPrice, error) {
	out := map[string]ModelPrice{}
	rows, err := pool.Query(ctx, `
		SELECT model, input_micros_per_mtok, cached_input_micros_per_mtok, output_micros_per_mtok
		FROM ai_model_prices`)
	if err != nil {
		return out, err
	}
	defer rows.Close()
	for rows.Next() {
		var model string
		var in, cached, o int64
		if err := rows.Scan(&model, &in, &cached, &o); err != nil {
			return out, err
		}
		out[NormalizeModel(model)] = ModelPrice{InputPerMTok: in, CachedInputPerMTok: cached, OutputPerMTok: o}
	}
	return out, rows.Err()
}

func priceFromMap(prices map[string]ModelPrice, model string) ModelPrice {
	key := NormalizeModel(model)
	if p, ok := prices[key]; ok {
		return p
	}
	if p, ok := DefaultModelPrice(key); ok {
		return p
	}
	return ModelPrice{}
}

// recoveredCall is one provider response found inside gpt_response.
type recoveredCall struct {
	Model            string
	Role             string
	PromptTokens     int
	CompletionTokens int
	TotalTokens      int
}

// extractUsageCalls walks a stored gpt_response and returns every provider usage
// block it holds. The shape varies by pipeline version, so this looks for any
// object with a "usage" child rather than assuming one fixed path — older rows
// nest it differently and a rigid path would silently recover nothing.
func extractUsageCalls(raw []byte) []recoveredCall {
	var doc any
	if err := json.Unmarshal(raw, &doc); err != nil {
		return nil
	}
	var out []recoveredCall
	walkUsage(doc, "", &out)
	return out
}

func walkUsage(node any, stepHint string, out *[]recoveredCall) {
	switch v := node.(type) {
	case map[string]any:
		if s, ok := v["step"].(string); ok && s != "" {
			stepHint = s
		}
		if usage, ok := v["usage"].(map[string]any); ok {
			call := recoveredCall{
				Model:            stringField(v, "model"),
				Role:             roleFromStep(stepHint),
				PromptTokens:     intField(usage, "prompt_tokens"),
				CompletionTokens: intField(usage, "completion_tokens"),
				TotalTokens:      intField(usage, "total_tokens"),
			}
			if call.PromptTokens > 0 || call.CompletionTokens > 0 {
				*out = append(*out, call)
			}
		}
		for _, child := range v {
			walkUsage(child, stepHint, out)
		}
	case []any:
		for _, child := range v {
			walkUsage(child, stepHint, out)
		}
	}
}

func roleFromStep(step string) string {
	switch step {
	case "categorize":
		return RoleCategorize
	case "ai_enhance", "":
		return RoleProcessing
	default:
		return RoleProcessing
	}
}

func stringField(m map[string]any, key string) string {
	if s, ok := m[key].(string); ok {
		return s
	}
	return ""
}

func intField(m map[string]any, key string) int {
	switch v := m[key].(type) {
	case float64:
		return int(v)
	case int:
		return v
	case json.Number:
		n, _ := v.Int64()
		return int(n)
	}
	return 0
}