batchaccounting

type AggregateLogEmitter interface {
	EmitBatchAggregateLog(ctx context.Context, entry *logstore.Log)
}

type AggregateLogStore interface {
	CreateIfNotExists(ctx context.Context, entry *logstore.Log) error
	// FindByID reads the aggregate cost row back, keyed by the same deterministic
	// id CreateIfNotExists wrote it under. Used only to mirror an already-settled
	// batch's price onto a caller that did not win this call's settlement claim
	// (see the !claimed branch in AccountBatchResults) — never to drive a write.
	FindByID(ctx context.Context, id string) (*logstore.Log, error)
}

type BatchJobStore interface {
	UpsertBatchJob(ctx context.Context, job *cstables.TableBatchJob) error
	GetBatchJob(ctx context.Context, jobID string) (*cstables.TableBatchJob, error)
	// ClaimBatchJob fences the job on runnerID. allowUnpriceable relaxes the
	// terminal-state guard to admit "unpriceable" jobs: that state means "stop
	// polling", not "refuse money", so a caller that actually holds results may
	// re-drive one. Only settlement paths with results in hand pass true — the
	// sweeper's poll loop must keep skipping unpriceable jobs (it does, via
	// ListDueBatchJobs). "accounted" is terminal forever either way.
	ClaimBatchJob(ctx context.Context, jobID, runnerID string, staleBefore time.Time, allowUnpriceable bool) (bool, error)
	MarkBatchJobAggregateLogWritten(ctx context.Context, jobID, runnerID string) error
	MarkBatchJobGovernanceReported(ctx context.Context, jobID, runnerID string) error
	CompleteBatchJob(ctx context.Context, jobID, runnerID string) error
	MarkBatchJobUnpriceable(ctx context.Context, jobID, runnerID, reason string, err error) error
	FailBatchJob(ctx context.Context, jobID, runnerID string, err error) error
}

type BatchModelUsage struct {
	Model      string
	Cost       float64
	TokensUsed int64
}

type BatchResultFetcher interface {
	RetrieveBatch(ctx context.Context, job *cstables.TableBatchJob) (*schemas.BifrostBatchRetrieveResponse, error)
	FetchBatchResults(ctx context.Context, job *cstables.TableBatchJob) (*schemas.BifrostBatchResultsResponse, error)
}

type BatchUsageReport struct {
	RequestID    string
	Provider     schemas.ModelProvider
	Model        string
	Cost         float64
	TokensUsed   int64
	BudgetIDs    []string
	RateLimitIDs []string
	UserID       string
	VirtualKeyID string
	ModelUsage   []BatchModelUsage
}

type ModelBreakdown = schemas.BatchModelBreakdown

type PricingManager interface {
	CalculateBatchCostDetailsForUsage(usage *schemas.BifrostLLMUsage, provider schemas.ModelProvider, model string, requestType schemas.RequestType, scopes *modelcatalog.PricingLookupScopes) modelcatalog.BatchCostDetails
}

type Request struct {
	Provider      schemas.ModelProvider
	BatchID       string
	FallbackModel string
	Endpoint      schemas.BatchEndpoint
	Results       []schemas.BatchResultItem
	ParseErrors   []schemas.BatchError
	RequestCounts *schemas.BatchRequestCounts
	BatchJob      *cstables.TableBatchJob
	BaseLog       *logstore.Log
	SourceLog     *logstore.Log
	Emitter       AggregateLogEmitter
	UsageReporter UsageReporter
	ClaimedBy     string
	Scopes        *modelcatalog.PricingLookupScopes
	Now           time.Time
}

type Summary struct {
	JobID           string                    `json:"job_id"`
	LogID           string                    `json:"log_id"`
	Provider        schemas.ModelProvider     `json:"provider"`
	BatchID         string                    `json:"batch_id"`
	Cost            float64                   `json:"cost"`
	Usage           schemas.BifrostLLMUsage   `json:"usage"`
	ModelBreakdowns map[string]ModelBreakdown `json:"model_breakdowns"`
	PricedCount     int                       `json:"priced_count"`
	UnpricedCount   int                       `json:"unpriced_count"`
	FailedCount     int                       `json:"failed_count"`
	// UnpricedModel names the model behind wholly-unpriced usage when it is
	// unambiguous, so the logged row is attributable (and therefore backfillable).
	// Empty when the usage spans several models or carried no model at all.
	UnpricedModel string `json:"unpriced_model,omitempty"`
	// Status is the provider's batch lifecycle status as of this settlement attempt
	// (from req.BatchJob.ProviderStatus when claimed, or mirrored from the
	// already-written aggregate row otherwise). See BifrostBatchDebug.Status.
	Status string `json:"status,omitempty"`
	// Complete reports that every usage-bearing row priced, so Cost is the batch's
	// real total. When false the batch consumed tokens we could not price and Cost
	// is only the known part — it must never be persisted as the row's cost, and
	// governance must not be told the batch settled for it.
	Complete          bool   `json:"complete"`
	Accounted         bool   `json:"accounted"`
	Claimed           bool   `json:"claimed"`
	UnpriceableReason string `json:"unpriceable_reason,omitempty"`
}

type SweepStore interface {
	BatchJobStore
	ListDueBatchJobs(ctx context.Context, provider string, now time.Time, limit int) ([]*cstables.TableBatchJob, error)
}

type Sweeper struct {
	// contains filtered or unexported fields
}

type SweeperConfig struct {
	Interval   time.Duration
	Limit      int
	ClaimedBy  string
	Provider   schemas.ModelProvider
	Scopes     *modelcatalog.PricingLookupScopes
	KVStore    schemas.KVStore
	KVLeaseTTL time.Duration
	Logger     schemas.Logger
}

type UsageReporter interface {
	ReportBatchUsage(ctx context.Context, usage BatchUsageReport) error
}