NianAIGC/backend/internal/localstore/store.go

// Package localstore provides process-local development persistence.
// It is intentionally non-durable and must not be used as a production store.
package localstore

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"
	"sort"
	"sync"
	"time"

	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/administration"
	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/assets"
	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing"
	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs"
	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/templates"
	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/usage"
)

const (
	DemoAccountID      = "demo-merchant"
	DemoOrganizationID = "org-demo"
	DemoPhone          = "13800000000"
	DemoPassword       = "demo-password"
)

type Option func(*Store)

func WithClock(now func() time.Time) Option {
	return func(s *Store) {
		if now != nil {
			s.now = now
		}
	}
}

type Store struct {
	mu            sync.RWMutex
	now           func() time.Time
	accounts      map[string]administration.Account
	organizations map[string]administration.Organization
	assets        map[string]assets.Asset
	jobs          map[string]jobs.Job
	templates     map[string]templates.Template
	usageEvents   map[string]usage.Event
	usageJobIDs   map[string]string
	priceRules    map[string]billing.PriceRule
	wallets       map[string]billing.Wallet
	ledger        []billing.LedgerEntry
	postings      map[string]walletRecord
}

type walletRecord struct {
	params  billing.WalletPostParams
	posting billing.WalletPosting
}

func New(options ...Option) *Store {
	s := &Store{now: time.Now, accounts: map[string]administration.Account{}, organizations: map[string]administration.Organization{}, assets: map[string]assets.Asset{}, jobs: map[string]jobs.Job{}, templates: map[string]templates.Template{}, usageEvents: map[string]usage.Event{}, usageJobIDs: map[string]string{}, priceRules: map[string]billing.PriceRule{}, wallets: map[string]billing.Wallet{}, postings: map[string]walletRecord{}}
	for _, option := range options {
		option(s)
	}
	s.seedDemo()
	return s
}

func NewStore(options ...Option) *Store { return New(options...) }

func (s *Store) seedDemo() {
	now := s.now().UTC()
	password, err := administration.HashPassword(DemoPassword)
	if err != nil {
		panic(fmt.Sprintf("localstore: hash demo password: %v", err))
	}
	s.organizations[DemoOrganizationID] = administration.Organization{ID: DemoOrganizationID, Name: "演示组织", Status: administration.StatusActive, ArchiveOwnerID: DemoAccountID, CreatedAt: now, UpdatedAt: now}
	s.accounts[DemoAccountID] = administration.Account{ID: DemoAccountID, Phone: DemoPhone, DisplayName: "智念用户", Role: administration.RoleSuperAdmin, OrganizationID: DemoOrganizationID, Status: administration.StatusActive, PasswordHash: password.Hash, PasswordSalt: password.Salt, SessionVersion: 1, CreatedAt: now, UpdatedAt: now}
}

func cloneAsset(a assets.Asset) assets.Asset {
	a.Tags = append([]string(nil), a.Tags...)
	a.Metadata = cloneMap(a.Metadata)
	return a
}
func cloneJob(j jobs.Job) jobs.Job {
	j.InputAssetIDs = append([]string(nil), j.InputAssetIDs...)
	j.InputURLs = append([]string(nil), j.InputURLs...)
	j.OutputAssetIDs = append([]string(nil), j.OutputAssetIDs...)
	j.RequestPayload = append(json.RawMessage(nil), j.RequestPayload...)
	j.ResponsePayload = append(json.RawMessage(nil), j.ResponsePayload...)
	j.WebhookLastStatus = append(json.RawMessage(nil), j.WebhookLastStatus...)
	j.UsageContext = append(json.RawMessage(nil), j.UsageContext...)
	j.Billing = append(json.RawMessage(nil), j.Billing...)
	if j.Error != nil {
		e := *j.Error
		j.Error = &e
	}
	if j.ProviderDispatchStartedAt != nil {
		value := *j.ProviderDispatchStartedAt
		j.ProviderDispatchStartedAt = &value
	}
	if j.DispatchReadyAt != nil {
		value := *j.DispatchReadyAt
		j.DispatchReadyAt = &value
	}
	if j.FinalizedAt != nil {
		value := *j.FinalizedAt
		j.FinalizedAt = &value
	}
	return j
}
func cloneMap(m map[string]any) map[string]any {
	if m == nil {
		return nil
	}
	raw, err := json.Marshal(m)
	if err != nil {
		out := make(map[string]any, len(m))
		for k, v := range m {
			out[k] = v
		}
		return out
	}
	var out map[string]any
	_ = json.Unmarshal(raw, &out)
	return out
}
func cloneRule(r billing.PriceRule) billing.PriceRule {
	r.Conditions = billing.Conditions(cloneMap(map[string]any(r.Conditions)))
	r.Source = cloneMap(r.Source)
	raw, _ := json.Marshal(r.Dimensions)
	_ = json.Unmarshal(raw, &r.Dimensions)
	return r
}

// Assets.
func (s *Store) ListOwner(_ context.Context, owner string) ([]assets.Asset, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := []assets.Asset{}
	for _, a := range s.assets {
		if a.OwnerID == owner {
			out = append(out, cloneAsset(a))
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt.After(out[j].CreatedAt) })
	return out, nil
}
func (s *Store) GetOwner(_ context.Context, owner, id string) (assets.Asset, bool, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	a, ok := s.assets[id]
	if !ok || a.OwnerID != owner {
		return assets.Asset{}, false, nil
	}
	return cloneAsset(a), true, nil
}
func (s *Store) GetOwnerByStoragePath(_ context.Context, owner, path string) (assets.Asset, bool, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, a := range s.assets {
		if a.OwnerID == owner && a.StoragePath == path {
			return cloneAsset(a), true, nil
		}
	}
	return assets.Asset{}, false, nil
}
func (s *Store) ListPublic(_ context.Context, owner, client string, limit int) ([]assets.Asset, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	allowed := s.publicAssetIDs(owner, client, limit)
	tag := assets.ClientTag(client)
	out := []assets.Asset{}
	for _, a := range s.assets {
		if a.OwnerID == owner && (contains(a.Tags, tag) || allowed[a.ID]) {
			out = append(out, cloneAsset(a))
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt.After(out[j].CreatedAt) })
	return out, nil
}
func (s *Store) GetPublic(_ context.Context, owner, client, id string, limit int) (assets.Asset, bool, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	a, ok := s.assets[id]
	if !ok || a.OwnerID != owner {
		return assets.Asset{}, false, nil
	}
	if !contains(a.Tags, assets.ClientTag(client)) && !s.publicAssetIDs(owner, client, limit)[id] {
		return assets.Asset{}, false, nil
	}
	return cloneAsset(a), true, nil
}
func (s *Store) publicAssetIDs(owner, client string, limit int) map[string]bool {
	candidates := []jobs.Job{}
	for _, j := range s.jobs {
		if j.OwnerID == owner && j.ExternalClientID == client {
			candidates = append(candidates, j)
		}
	}
	sort.Slice(candidates, func(i, j int) bool { return candidates[i].CreatedAt.After(candidates[j].CreatedAt) })
	if limit <= 0 {
		limit = 200
	}
	if len(candidates) > limit {
		candidates = candidates[:limit]
	}
	out := map[string]bool{}
	for _, j := range candidates {
		for _, id := range append(append([]string{}, j.InputAssetIDs...), j.OutputAssetIDs...) {
			out[id] = true
		}
	}
	return out
}
func (s *Store) Create(_ context.Context, a assets.Asset) (assets.Asset, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.assets[a.ID]; ok {
		return assets.Asset{}, fmt.Errorf("asset already exists: %s", a.ID)
	}
	if a.Tags == nil {
		a.Tags = []string{}
	}
	if a.Metadata == nil {
		a.Metadata = map[string]any{}
	}
	s.assets[a.ID] = cloneAsset(a)
	return cloneAsset(a), nil
}
func (s *Store) DeleteOwner(_ context.Context, owner, id string) (assets.Asset, bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	a, ok := s.assets[id]
	if !ok || a.OwnerID != owner {
		return assets.Asset{}, false, nil
	}
	delete(s.assets, id)
	return cloneAsset(a), true, nil
}

// Jobs.
func (s *Store) ListJobs(_ context.Context, f jobs.ListFilter) ([]jobs.Job, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := []jobs.Job{}
	for _, j := range s.jobs {
		if !f.Scope.Owns(j) || f.Status != "" && j.Status != f.Status || f.Capability != "" && j.Capability != f.Capability || f.Before != nil && !j.CreatedAt.Before(*f.Before) {
			continue
		}
		out = append(out, cloneJob(j))
	}
	sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt.After(out[j].CreatedAt) })
	if f.Limit > 0 && len(out) > f.Limit {
		out = out[:f.Limit]
	}
	return out, nil
}
func (s *Store) FindJob(_ context.Context, scope jobs.Scope, id string) (jobs.Job, bool, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	j, ok := s.jobs[id]
	if !ok || !scope.Owns(j) {
		return jobs.Job{}, false, nil
	}
	return cloneJob(j), true, nil
}
func (s *Store) FindIdempotentJob(_ context.Context, scope jobs.Scope, key string) (jobs.Job, bool, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, j := range s.jobs {
		if scope.Owns(j) && j.IdempotencyKey == key {
			return cloneJob(j), true, nil
		}
	}
	return jobs.Job{}, false, nil
}
func (s *Store) CreateJob(_ context.Context, j jobs.Job) (jobs.Job, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.jobs[j.ID]; ok {
		return jobs.Job{}, fmt.Errorf("generation job already exists: %s", j.ID)
	}
	if j.ExternalClientID != "" && j.IdempotencyKey != "" {
		for _, v := range s.jobs {
			if v.OwnerID == j.OwnerID && v.ExternalClientID == j.ExternalClientID && v.IdempotencyKey == j.IdempotencyKey {
				return jobs.Job{}, jobs.ErrUniqueIdempotency
			}
		}
	}
	s.jobs[j.ID] = cloneJob(j)
	return cloneJob(j), nil
}
func (s *Store) UpdateJob(_ context.Context, id string, p jobs.Patch) (jobs.Job, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	j, ok := s.jobs[id]
	if !ok {
		return jobs.Job{}, fmt.Errorf("generation job not found: %s", id)
	}
	if len(p.ExpectedStatuses) > 0 && !containsStatus(p.ExpectedStatuses, j.Status) {
		return jobs.Job{}, jobs.ErrTransitionConflict
	}
	if p.ExpectedLockedBy != nil && j.LockedBy != *p.ExpectedLockedBy {
		return jobs.Job{}, jobs.ErrTransitionConflict
	}
	if p.Status != nil {
		j.Status = *p.Status
	}
	if p.Error != nil {
		e := *p.Error
		j.Error = &e
	}
	if p.ClearError {
		j.Error = nil
	}
	if p.Attempts != nil {
		j.Attempts = *p.Attempts
	}
	if p.ScheduledAt != nil {
		j.ScheduledAt = *p.ScheduledAt
	}
	if p.CompletedAt != nil {
		v := *p.CompletedAt
		j.CompletedAt = &v
	}
	if p.ProviderTaskID != nil {
		j.ProviderTaskID = *p.ProviderTaskID
	}
	if p.ProviderDispatchStartedAt != nil {
		started := *p.ProviderDispatchStartedAt
		j.ProviderDispatchStartedAt = &started
	}
	if p.ClearProviderDispatch {
		j.ProviderDispatchStartedAt = nil
	}
	if p.SetResponsePayload {
		j.ResponsePayload = append([]byte(nil), p.ResponsePayload...)
	}
	if p.DispatchReadyAt != nil {
		ready := *p.DispatchReadyAt
		j.DispatchReadyAt = &ready
	}
	if p.FinalizedAt != nil {
		finalized := *p.FinalizedAt
		j.FinalizedAt = &finalized
	}
	if p.ClearProviderTaskID {
		j.ProviderTaskID = ""
	}
	if p.ClearLease {
		j.LockedAt = nil
		j.LockedBy = ""
	}
	if p.WebhookAttempts != nil {
		j.WebhookAttempts = *p.WebhookAttempts
	}
	if p.SetWebhookStatus {
		j.WebhookLastStatus = append([]byte(nil), p.WebhookLastStatus...)
	}
	j.UpdatedAt = s.now().UTC()
	s.jobs[id] = cloneJob(j)
	return cloneJob(j), nil
}
func (s *Store) DeleteJob(_ context.Context, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.jobs[id]; !ok {
		return fmt.Errorf("generation job not found: %s", id)
	}
	delete(s.jobs, id)
	return nil
}
func (s *Store) ClaimJobs(_ context.Context, worker string, limit, timeoutSeconds int) ([]jobs.Job, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if limit < 1 {
		limit = 1
	}
	if limit > 20 {
		limit = 20
	}
	if timeoutSeconds <= 0 {
		timeoutSeconds = 300
	}
	now := s.now().UTC()
	eligible := []jobs.Job{}
	for _, j := range s.jobs {
		if j.DispatchReadyAt == nil || j.FinalizedAt != nil || j.ScheduledAt.After(now) {
			continue
		}
		leaseAvailable := j.LockedAt == nil || j.LockedAt.Add(time.Duration(timeoutSeconds)*time.Second).Before(now)
		if leaseAvailable && (j.Status == jobs.StatusQueued || j.Status == jobs.StatusRunning || j.Status.Terminal()) {
			eligible = append(eligible, j)
		}
	}
	sort.Slice(eligible, func(i, j int) bool {
		if eligible[i].Priority != eligible[j].Priority {
			return eligible[i].Priority > eligible[j].Priority
		}
		if !eligible[i].ScheduledAt.Equal(eligible[j].ScheduledAt) {
			return eligible[i].ScheduledAt.Before(eligible[j].ScheduledAt)
		}
		return eligible[i].CreatedAt.Before(eligible[j].CreatedAt)
	})
	if len(eligible) > limit {
		eligible = eligible[:limit]
	}
	out := make([]jobs.Job, 0, len(eligible))
	for _, j := range eligible {
		j.LockedBy = worker
		j.LockedAt = timePtr(now)
		if j.StartedAt == nil {
			j.StartedAt = timePtr(now)
		}
		j.UpdatedAt = now
		s.jobs[j.ID] = cloneJob(j)
		out = append(out, cloneJob(j))
	}
	return out, nil
}
func (s *Store) WriteBilling(_ context.Context, id string, value json.RawMessage) error {
	if len(value) == 0 || !json.Valid(value) {
		return fmt.Errorf("write generation billing: invalid JSON")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	j, ok := s.jobs[id]
	if !ok {
		return fmt.Errorf("write generation billing: generation job not found")
	}
	j.Billing = append([]byte(nil), value...)
	j.UpdatedAt = s.now().UTC()
	s.jobs[id] = j
	return nil
}

func (s *Store) WriteBillingFenced(_ context.Context, id string, value json.RawMessage, status jobs.Status, lockedBy string) error {
	if len(value) == 0 || !json.Valid(value) {
		return fmt.Errorf("write generation billing: invalid JSON")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	j, ok := s.jobs[id]
	if !ok || j.Status != status || j.LockedBy != lockedBy {
		return jobs.ErrTransitionConflict
	}
	j.Billing = append([]byte(nil), value...)
	j.UpdatedAt = s.now().UTC()
	s.jobs[id] = cloneJob(j)
	return nil
}

func (s *Store) ActivateCreation(_ context.Context, id string, value json.RawMessage) error {
	if len(value) == 0 || !json.Valid(value) {
		return fmt.Errorf("activate generation creation: invalid JSON")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	j, ok := s.jobs[id]
	if !ok {
		return fmt.Errorf("activate generation creation: generation job not found")
	}
	if j.DispatchReadyAt != nil {
		if !jsonEqual(j.Billing, value) {
			return jobs.ErrTransitionConflict
		}
		return nil
	}
	j.Billing = append([]byte(nil), value...)
	ready := s.now().UTC()
	j.DispatchReadyAt = &ready
	j.UpdatedAt = s.now().UTC()
	s.jobs[id] = cloneJob(j)
	return nil
}

func jsonEqual(a, b json.RawMessage) bool {
	var left, right any
	return json.Unmarshal(a, &left) == nil && json.Unmarshal(b, &right) == nil && reflect.DeepEqual(left, right)
}
func (s *Store) WriteOutputAssetIDs(_ context.Context, id string, values []string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	j, ok := s.jobs[id]
	if !ok {
		return fmt.Errorf("write generation output assets: generation job not found")
	}
	j.OutputAssetIDs = append([]string{}, values...)
	j.UpdatedAt = s.now().UTC()
	s.jobs[id] = j
	return nil
}

func (s *Store) WriteOutputAssetIDsFenced(_ context.Context, id string, values []string, status jobs.Status, lockedBy string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	j, ok := s.jobs[id]
	if !ok || j.Status != status || j.LockedBy != lockedBy {
		return jobs.ErrTransitionConflict
	}
	j.OutputAssetIDs = append([]string{}, values...)
	j.UpdatedAt = s.now().UTC()
	s.jobs[id] = cloneJob(j)
	return nil
}
func (s *Store) FailCreation(_ context.Context, j jobs.Job) error {
	if j.Status != jobs.StatusFailed || j.Error == nil || len(j.Billing) == 0 || !json.Valid(j.Billing) {
		return fmt.Errorf("fail generation creation: invalid terminal state")
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	current, ok := s.jobs[j.ID]
	if !ok {
		return fmt.Errorf("fail generation creation: generation job not found")
	}
	current.Status = j.Status
	e := *j.Error
	current.Error = &e
	current.Billing = append([]byte(nil), j.Billing...)
	now := s.now().UTC()
	if current.CompletedAt == nil {
		current.CompletedAt = &now
	}
	current.LockedAt = nil
	current.LockedBy = ""
	current.UpdatedAt = now
	s.jobs[j.ID] = current
	return nil
}

func timePtr(v time.Time) *time.Time { return &v }
func contains(values []string, want string) bool {
	for _, v := range values {
		if v == want {
			return true
		}
	}
	return false
}

func containsStatus(values []jobs.Status, want jobs.Status) bool {
	for _, value := range values {
		if value == want {
			return true
		}
	}
	return false
}