NianAIGC/backend/internal/orchestration/settlement.go

package orchestration

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"time"

	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing"
	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs"
)

type SettlementLedger interface {
	Settle(context.Context, billing.SettlementRequest) (billing.WalletPosting, error)
}

type SettlementStateWriter interface {
	WriteBilling(context.Context, string, json.RawMessage) error
}

type FencedSettlementStateWriter interface {
	WriteBillingFenced(context.Context, string, json.RawMessage, jobs.Status, string) error
}

// SettlementProcessor adds Seedance actual-usage settlement to a jobs
// processor. Other providers and non-successful advances pass through.
type SettlementProcessor struct {
	inner  jobs.Processor
	ledger SettlementLedger
	state  SettlementStateWriter
	now    func() time.Time
}

func NewSettlementProcessor(inner jobs.Processor, ledger SettlementLedger, state SettlementStateWriter, now func() time.Time) *SettlementProcessor {
	if now == nil {
		now = time.Now
	}
	return &SettlementProcessor{inner: inner, ledger: ledger, state: state, now: now}
}

func (p *SettlementProcessor) Advance(ctx context.Context, job jobs.Job) (jobs.Job, error) {
	if p == nil || p.inner == nil {
		return jobs.Job{}, errors.New("advance generation job")
	}
	advanced, err := p.inner.Advance(ctx, job)
	if err != nil || advanced.Provider != "seedance" || advanced.Status != jobs.StatusSucceeded || len(advanced.Billing) == 0 {
		return advanced, err
	}

	var charge billingState
	if json.Unmarshal(advanced.Billing, &charge) != nil {
		return jobs.Job{}, errors.New("settle seedance generation charge")
	}
	status, _ := charge.raw["settlementStatus"].(string)
	if status == "settled" || status == "estimated" {
		return advanced, nil
	}
	var use usageContext
	if len(advanced.UsageContext) != 0 && json.Unmarshal(advanced.UsageContext, &use) != nil {
		return jobs.Job{}, errors.New("settle seedance generation charge")
	}
	quotaExempt := charge.QuotaExempt || use.Source == "platform" && use.Role == "super_admin"
	chargeReady := charge.Status == "charged" || quotaExempt && charge.Status == "not_charged"
	if !chargeReady || !quotaExempt && use.OrganizationID == "" {
		return advanced, nil
	}

	completionTokens := seedanceCompletionTokens(advanced.ResponsePayload)
	if completionTokens <= 0 {
		charge.raw["settlementStatus"] = "estimated"
		charge.raw["settlementReason"] = "provider_usage_unavailable"
		charge.raw["settledAt"] = p.now().UTC().Format(time.RFC3339Nano)
		return p.write(ctx, advanced, charge.raw)
	}

	resolution := seedanceResolution(charge.raw)
	inputVideo := seedanceRequestHasInputVideo(advanced.RequestPayload)
	actualAmount, err := billing.CalculateSeedanceActualAmountFen(billing.SeedanceActualAmountInput{
		Resolution: resolution, InputVideo: inputVideo, CompletionTokens: completionTokens,
		MarkupMultiplier: numberOrZero(charge.raw["markupMultiplier"]),
	})
	if err != nil {
		return jobs.Job{}, errors.New("settle seedance generation charge")
	}
	delta := actualAmount - charge.AmountFen
	settledAt := p.now().UTC()
	var settlementLedgerID string
	if delta != 0 && !quotaExempt {
		if p.ledger == nil {
			return jobs.Job{}, errors.New("settle seedance generation charge")
		}
		description := capabilityLabel(advanced.Capability) + "实际用量差额退回"
		if delta > 0 {
			description = capabilityLabel(advanced.Capability) + "实际用量补扣"
		}
		reservedAmount := charge.AmountFen
		if value, ok := integer(charge.raw["reservedAmountFen"]); ok {
			reservedAmount = value
		}
		posting, postErr := p.ledger.Settle(ctx, billing.SettlementRequest{
			OrganizationID: use.OrganizationID, AccountID: use.AccountID, JobID: advanced.ID,
			DeltaFen: delta, Description: description,
			Metadata: map[string]any{
				"operation": "seedance_actual_settlement", "reservedAmountFen": reservedAmount,
				"chargedAmountFen": charge.AmountFen, "actualAmountFen": actualAmount,
				"completionTokens": completionTokens, "inputVideo": inputVideo, "resolution": resolution,
			},
		})
		if postErr != nil {
			return jobs.Job{}, errors.New("settle seedance generation charge")
		}
		settlementLedgerID = posting.LedgerID
		if !posting.CreatedAt.IsZero() {
			settledAt = posting.CreatedAt.UTC()
		}
	}

	charge.raw["amountFen"] = actualAmount
	charge.raw["settlementStatus"] = "settled"
	charge.raw["settledAt"] = settledAt.Format(time.RFC3339Nano)
	if settlementLedgerID != "" {
		charge.raw["settlementLedgerEntryId"] = settlementLedgerID
	} else {
		delete(charge.raw, "settlementLedgerEntryId")
	}
	charge.raw["providerUsage"] = map[string]any{
		"completionTokens": completionTokens, "resolution": resolution, "inputVideo": inputVideo,
		"tokenPriceFenPerMillion": billing.SeedanceTokenPriceFenPerMillion(resolution, inputVideo),
	}
	return p.write(ctx, advanced, charge.raw)
}

func (p *SettlementProcessor) write(ctx context.Context, job jobs.Job, snapshot map[string]any) (jobs.Job, error) {
	if p.state == nil {
		return jobs.Job{}, errors.New("persist seedance generation settlement")
	}
	encoded, err := json.Marshal(snapshot)
	if err != nil {
		return jobs.Job{}, errors.New("persist seedance generation settlement")
	}
	var writeErr error
	if fenced, ok := p.state.(FencedSettlementStateWriter); ok && job.LockedBy != "" {
		writeErr = fenced.WriteBillingFenced(ctx, job.ID, encoded, job.Status, job.LockedBy)
	} else {
		writeErr = p.state.WriteBilling(ctx, job.ID, encoded)
	}
	if writeErr != nil {
		return jobs.Job{}, errors.New("persist seedance generation settlement")
	}
	job.Billing = encoded
	return job, nil
}

func seedanceCompletionTokens(payload json.RawMessage) int64 {
	var response struct {
		Usage map[string]any `json:"usage"`
	}
	if json.Unmarshal(payload, &response) != nil {
		return 0
	}
	value, _ := integer(response.Usage["completionTokens"])
	return value
}

func seedanceResolution(snapshot map[string]any) string {
	parameters, _ := snapshot["parameters"].(map[string]any)
	resolution, _ := parameters["resolution"].(string)
	resolution = strings.ToLower(strings.TrimSpace(resolution))
	if resolution == "480p" || resolution == "1080p" || resolution == "4k" {
		return resolution
	}
	return "720p"
}

func seedanceRequestHasInputVideo(payload json.RawMessage) bool {
	var value any
	if json.Unmarshal(payload, &value) != nil {
		return false
	}
	return hasVideoMaterial(value)
}

func hasVideoMaterial(value any) bool {
	switch typed := value.(type) {
	case map[string]any:
		if materialType, _ := typed["type"].(string); strings.EqualFold(strings.TrimSpace(materialType), "video") {
			return true
		}
		for _, child := range typed {
			if hasVideoMaterial(child) {
				return true
			}
		}
	case []any:
		for _, child := range typed {
			if hasVideoMaterial(child) {
				return true
			}
		}
	}
	return false
}

func numberOrZero(value any) float64 {
	if number, ok := value.(float64); ok {
		return number
	}
	return 0
}

func integer(value any) (int64, bool) {
	number, ok := value.(float64)
	if !ok || number <= 0 {
		return 0, false
	}
	return int64(number), true
}

var _ jobs.Processor = (*SettlementProcessor)(nil)