moteva/server/internal/svc/service_context.go

// Code scaffolded by goctl. Safe to edit.
// goctl 1.10.1

package svc

import (
	"context"
	"fmt"
	"os"
	"strings"
	"time"

	"img_infinite_canvas/internal/application"
	"img_infinite_canvas/internal/config"
	"img_infinite_canvas/internal/domain/design"
	"img_infinite_canvas/internal/infrastructure/backgroundremoval"
	cacheinfra "img_infinite_canvas/internal/infrastructure/cache"
	"img_infinite_canvas/internal/infrastructure/memory"
	"img_infinite_canvas/internal/infrastructure/objectrecognition"
	"img_infinite_canvas/internal/infrastructure/objectstorage"
	"img_infinite_canvas/internal/infrastructure/piagent"
	"img_infinite_canvas/internal/infrastructure/postgres"
	"img_infinite_canvas/internal/infrastructure/textextraction"
	"img_infinite_canvas/internal/infrastructure/websearch"
	authmodule "img_infinite_canvas/internal/modules/auth"
	jobmodule "img_infinite_canvas/internal/modules/job"
	realtimemodule "img_infinite_canvas/internal/modules/realtime"

	"github.com/zeromicro/go-zero/core/logx"
)

type ServiceContext struct {
	Config        config.Config
	DesignService *application.DesignService
	AuthService   *authmodule.Service
	RealtimeBus   realtimemodule.Bus
	jobQueue      design.JobQueue
	jobWorker     backgroundWorker
	outboxWorker  backgroundWorker
}

type backgroundWorker interface {
	Start() error
	Run() error
	Shutdown()
}

func NewServiceContext(c config.Config) *ServiceContext {
	repo := newProjectRepository(c)
	cacheStore := newCacheStore(c)
	authService := newAuthService(c, cacheStore)
	repo = withCache(c, repo, cacheStore)
	realtimeBus := newRealtimeBus(c)
	repo = withRealtime(repo, realtimeBus)
	creativeAgent := newCreativeAgent(c)
	agentRunner := newAgentRunner(c)
	assets := newAssetStorage(c)
	researcher := newResearcher(c)
	textExtractor := newTextExtractor(c)
	service := application.NewDesignService(repo, agentRunner, assets, researcher, textExtractor, creativeAgent)
	service.SetBackgroundRemover(newBackgroundRemover(c))
	service.SetObjectRecognizer(newObjectRecognizer(c))
	if mockupAnalyzer, ok := creativeAgent.(design.MockupModelAnalyzer); ok {
		service.SetMockupModelAnalyzer(mockupAnalyzer)
	}
	service.SetTextExtractionCache(cacheStore, textExtractionCacheTTL(c))
	queue, worker := newJobQueue(c, service)
	service.SetJobQueue(queue)
	outboxWorker := newOutboxPublisher(c, realtimeBus)
	if err := service.EnsureSeedData(context.Background()); err != nil {
		logx.Must(err)
	}

	return &ServiceContext{
		Config:        c,
		DesignService: service,
		AuthService:   authService,
		RealtimeBus:   realtimeBus,
		jobQueue:      queue,
		jobWorker:     worker,
		outboxWorker:  outboxWorker,
	}
}

func (s *ServiceContext) StartJobWorker() error {
	if s == nil || s.jobWorker == nil {
		return nil
	}
	return s.jobWorker.Start()
}

func (s *ServiceContext) HasJobWorker() bool {
	return s != nil && s.jobWorker != nil
}

func (s *ServiceContext) RunJobWorker() error {
	if s == nil || s.jobWorker == nil {
		return nil
	}
	return s.jobWorker.Run()
}

func (s *ServiceContext) ShutdownJobWorker() {
	if s == nil || s.jobWorker == nil {
		return
	}
	s.jobWorker.Shutdown()
}

func (s *ServiceContext) StartOutboxPublisher() error {
	if s == nil || s.outboxWorker == nil {
		return nil
	}
	return s.outboxWorker.Start()
}

func (s *ServiceContext) ShutdownOutboxPublisher() {
	if s == nil || s.outboxWorker == nil {
		return
	}
	s.outboxWorker.Shutdown()
}

func (s *ServiceContext) RealtimeFallbackPoll() time.Duration {
	if s == nil {
		return 10 * time.Second
	}
	interval := time.Duration(s.Config.Realtime.FallbackPollMillis) * time.Millisecond
	if interval <= 0 {
		return 10 * time.Second
	}
	return interval
}

func (s *ServiceContext) Close() {
	s.ShutdownJobWorker()
	s.ShutdownOutboxPublisher()
	if closer, ok := s.jobQueue.(interface{ Close() error }); ok {
		if err := closer.Close(); err != nil {
			logx.Errorf("close job queue failed: %v", err)
		}
	}
	if s.RealtimeBus != nil {
		if err := s.RealtimeBus.Close(); err != nil {
			logx.Errorf("close realtime bus failed: %v", err)
		}
	}
	if s.AuthService != nil {
		if err := s.AuthService.Close(); err != nil {
			logx.Errorf("close auth service failed: %v", err)
		}
	}
}

func newAuthService(c config.Config, cacheStore cacheinfra.Store) *authmodule.Service {
	store := newAuthStore(c)
	opts := c.Auth
	codeTTL := time.Duration(opts.CodeTTLSeconds) * time.Second
	codeStore := authmodule.NewCacheCodeStore(cacheStore, codeTTL)
	humanVerifier := newAuthHumanVerifier(c)
	return authmodule.NewServiceWithCodeStoreEmailAndHumanVerifier(store, codeStore, newAuthEmailSender(c), humanVerifier, authmodule.Config{
		Region:            opts.Region,
		RegionMode:        opts.RegionMode,
		TokenSecret:       opts.TokenSecret,
		TokenTTL:          time.Duration(opts.TokenTTLSeconds) * time.Second,
		RefreshTokenTTL:   time.Duration(opts.RefreshTTLSeconds) * time.Second,
		CodeTTL:           time.Duration(opts.CodeTTLSeconds) * time.Second,
		DevReturnCode:     opts.DevReturnCode,
		EmailFromName:     opts.Email.FromName,
		EmailFromEmail:    opts.Email.FromEmail,
		BrandName:         opts.Email.BrandName,
		TurnstileRequired: opts.Turnstile.Enabled,
		TurnstileSiteKey:  opts.Turnstile.SiteKey,
		GoogleClientID:    opts.Google.ClientID,
		GoogleTokenInfo:   opts.Google.TokenInfoURL,
		GoogleTimeout:     time.Duration(opts.Google.TimeoutSeconds) * time.Second,
		WechatAppID:       opts.Wechat.AppID,
		WechatAppSecret:   opts.Wechat.AppSecret,
		WechatRedirect:    opts.Wechat.RedirectURI,
		WechatScope:       opts.Wechat.Scope,
		WechatAuthURL:     opts.Wechat.AuthURL,
		WechatTokenURL:    opts.Wechat.TokenURL,
		WechatTimeout:     time.Duration(opts.Wechat.TimeoutSeconds) * time.Second,
	})
}

func newAuthHumanVerifier(c config.Config) authmodule.HumanVerifier {
	opts := c.Auth.Turnstile
	if !opts.Enabled {
		return nil
	}
	secretKey := strings.TrimSpace(opts.SecretKey)
	if envName := strings.TrimSpace(opts.SecretKeyEnv); envName != "" {
		if envValue := strings.TrimSpace(os.Getenv(envName)); envValue != "" {
			secretKey = envValue
		}
	}
	verifier, err := authmodule.NewTurnstileVerifier(authmodule.TurnstileVerifierOptions{
		SecretKey:     secretKey,
		SiteVerifyURL: opts.SiteVerifyURL,
		Timeout:       time.Duration(opts.TimeoutSeconds) * time.Second,
	})
	if err != nil {
		logx.Must(err)
	}
	return verifier
}

func newAuthEmailSender(c config.Config) authmodule.EmailSender {
	opts := c.Auth.Email
	switch strings.ToLower(strings.TrimSpace(opts.Driver)) {
	case "", "none", "disabled":
		return nil
	case "smtp":
		sender, err := authmodule.NewSMTPSender(authmodule.SMTPOptions{
			Host:               opts.Host,
			Port:               opts.Port,
			Username:           opts.Username,
			Password:           opts.Password,
			UseTLS:             opts.UseTLS,
			StartTLS:           opts.StartTLS,
			InsecureSkipVerify: opts.InsecureSkipVerify,
			Timeout:            time.Duration(opts.TimeoutSeconds) * time.Second,
		})
		if err != nil {
			logx.Must(err)
		}
		return sender
	default:
		logx.Must(fmt.Errorf("unsupported auth email driver: %s", opts.Driver))
		return nil
	}
}

func newAuthStore(c config.Config) authmodule.Store {
	switch c.Storage.Driver {
	case "", "memory":
		return authmodule.NewMemoryStore()
	case "postgres":
		if c.Storage.DataSource == "" {
			logx.Must(application.ErrMissingDataSource)
		}
		store, err := authmodule.NewPostgresStore(context.Background(), c.Storage.DataSource)
		if err != nil {
			logx.Must(err)
		}
		return store
	default:
		logx.Must(application.ErrUnsupportedStorage)
		return authmodule.NewMemoryStore()
	}
}

func newResearcher(c config.Config) design.Researcher {
	opts := c.Agent.Research
	switch opts.Driver {
	case "", "duckduckgo":
		return websearch.NewDuckDuckGoResearcher(websearch.DuckDuckGoOptions{
			TimeoutSeconds: opts.TimeoutSeconds,
			MaxResults:     opts.MaxResults,
		})
	case "none", "disabled":
		return nil
	default:
		logx.Must(application.ErrUnsupportedResearch)
		return nil
	}
}

func newAgentRunner(c config.Config) design.AgentRunner {
	base := application.NewDeterministicAgent()
	creativeAgent := newCreativeAgent(c)
	switch c.Agent.Driver {
	case "", "pi":
		return piagent.NewDesignAgent(base, newImageGenerator(c), creativeAgent)
	case "deterministic":
		return base
	default:
		logx.Must(application.ErrUnsupportedAgent)
		return base
	}
}

func newCreativeAgent(c config.Config) design.CreativeAgent {
	opts := c.Agent.Planner
	switch opts.Driver {
	case "", "vision":
		apiKey := resolveImageAPIKey(c.Agent.Image)
		return piagent.NewCreativeAgent(piagent.CreativeAgentOptions{
			BaseURL:        c.Agent.Image.BaseURL,
			APIKey:         apiKey,
			Model:          opts.Model,
			TimeoutSeconds: opts.TimeoutSeconds,
		})
	case "none", "disabled":
		return nil
	default:
		logx.Must(application.ErrUnsupportedAgent)
		return nil
	}
}

func newImageGenerator(c config.Config) piagent.ImageGenerator {
	opts := c.Agent.Image
	apiKey := resolveImageAPIKey(opts)
	client := piagent.NewImageClient(piagent.ImageClientOptions{
		BaseURL:        opts.BaseURL,
		APIKey:         apiKey,
		TimeoutSeconds: opts.TimeoutSeconds,
	})
	if client == nil {
		return nil
	}
	return client
}

func newTextExtractor(c config.Config) design.TextExtractor {
	opts := c.Agent.Text
	switch opts.Driver {
	case "", "vision":
		apiKey := resolveImageAPIKey(c.Agent.Image)
		return textextraction.NewTextExtractionClient(textextraction.TextExtractionClientOptions{
			BaseURL:        c.Agent.Image.BaseURL,
			APIKey:         apiKey,
			Model:          opts.Model,
			TimeoutSeconds: opts.TimeoutSeconds,
		})
	case "none", "disabled":
		return nil
	default:
		logx.Must(application.ErrUnsupportedAgent)
		return nil
	}
}

func newObjectRecognizer(c config.Config) design.ObjectRecognizer {
	opts := c.Agent.Object
	switch opts.Driver {
	case "", "vision":
		apiKey := resolveImageAPIKey(c.Agent.Image)
		return objectrecognition.NewClient(objectrecognition.ClientOptions{
			BaseURL:        c.Agent.Image.BaseURL,
			APIKey:         apiKey,
			Model:          opts.Model,
			TimeoutSeconds: opts.TimeoutSeconds,
		})
	case "none", "disabled":
		return nil
	default:
		logx.Must(application.ErrUnsupportedAgent)
		return nil
	}
}

func newBackgroundRemover(c config.Config) design.BackgroundRemover {
	opts := c.Agent.Background
	switch opts.Driver {
	case "", "onnxruntime":
		return backgroundremoval.NewONNXRuntimeRemover(backgroundremoval.ONNXRuntimeOptions{
			ModelPath:         opts.ModelPath,
			SharedLibraryPath: opts.SharedLibraryPath,
			ExecutionProvider: opts.ExecutionProvider,
			InputSize:         opts.InputSize,
			TimeoutSeconds:    opts.TimeoutSeconds,
			InputName:         opts.InputName,
			OutputName:        opts.OutputName,
		})
	case "none", "disabled":
		return nil
	default:
		logx.Must(application.ErrUnsupportedAgent)
		return nil
	}
}

func resolveImageAPIKey(opts config.ImageGenerationConfig) string {
	if opts.APIKeyEnv != "" {
		if value := strings.TrimSpace(os.Getenv(opts.APIKeyEnv)); value != "" {
			return value
		}
	}
	return strings.TrimSpace(opts.APIKey)
}

func newAssetStorage(c config.Config) design.AssetStorage {
	opts := c.ObjectStorage
	switch opts.Driver {
	case "", "memory":
		return objectstorage.NewMemoryStorage(opts.Bucket, opts.PublicBase, opts.ExpiresIn)
	case "minio", "s3", "aws", "r2":
		storage, err := objectstorage.NewS3Storage(context.Background(), objectstorage.S3Options{
			Driver:     opts.Driver,
			Bucket:     opts.Bucket,
			Endpoint:   opts.Endpoint,
			Region:     opts.Region,
			AccessKey:  opts.AccessKey,
			SecretKey:  opts.SecretKey,
			PublicBase: opts.PublicBase,
			UseSSL:     opts.UseSSL,
			ExpiresIn:  opts.ExpiresIn,
		})
		if err != nil {
			logx.Must(err)
		}
		return storage
	case "aliyun":
		storage, err := objectstorage.NewAliyunStorage(context.Background(), objectstorage.AliyunOptions{
			Bucket:     opts.Bucket,
			Endpoint:   opts.Endpoint,
			AccessKey:  opts.AccessKey,
			SecretKey:  opts.SecretKey,
			PublicBase: opts.PublicBase,
			ExpiresIn:  opts.ExpiresIn,
		})
		if err != nil {
			logx.Must(err)
		}
		return storage
	default:
		logx.Must(application.ErrUnsupportedObjectStorage)
		return objectstorage.NewMemoryStorage(opts.Bucket, opts.PublicBase, opts.ExpiresIn)
	}
}

func withCache(c config.Config, repo design.Repository, store cacheinfra.Store) design.Repository {
	ttl := time.Duration(c.Cache.TTLSeconds) * time.Second
	if ttl <= 0 {
		ttl = 5 * time.Minute
	}
	return cacheinfra.NewProjectRepository(repo, store, ttl)
}

func withRealtime(repo design.Repository, bus realtimemodule.Bus) design.Repository {
	if bus == nil {
		return repo
	}
	return realtimemodule.NewProjectRepository(repo, bus)
}

func textExtractionCacheTTL(c config.Config) time.Duration {
	ttl := time.Duration(c.Agent.Text.CacheTTLSeconds) * time.Second
	if ttl <= 0 {
		return 10 * time.Minute
	}
	return ttl
}

func newCacheStore(c config.Config) cacheinfra.Store {
	switch c.Cache.Driver {
	case "", "memory":
		return cacheinfra.NewMemoryStore()
	case "redis":
		store, err := cacheinfra.NewRedisStore(context.Background(), cacheinfra.RedisOptions{
			Addr:     c.Cache.Addr,
			Password: c.Cache.Password,
			DB:       c.Cache.DB,
		})
		if err != nil {
			logx.Must(err)
		}
		return store
	default:
		logx.Must(cacheinfra.ErrUnsupportedDriver)
		return cacheinfra.NewMemoryStore()
	}
}

func newRealtimeBus(c config.Config) realtimemodule.Bus {
	opts := c.Realtime
	driver := strings.ToLower(strings.TrimSpace(opts.Driver))
	switch driver {
	case "", "memory":
		return realtimemodule.NewMemoryBus()
	case "redis":
		if strings.TrimSpace(opts.Addr) == "" {
			opts.Addr = c.Cache.Addr
		}
		if strings.TrimSpace(opts.Password) == "" {
			opts.Password = c.Cache.Password
		}
		bus, err := realtimemodule.NewRedisBus(context.Background(), realtimemodule.RedisOptions{
			Addr:     opts.Addr,
			Password: opts.Password,
			DB:       opts.DB,
			Prefix:   opts.ChannelPrefix,
		})
		if err != nil {
			logx.Must(err)
		}
		return bus
	case "none", "disabled":
		return nil
	default:
		logx.Must(fmt.Errorf("unsupported realtime driver: %s", opts.Driver))
		return nil
	}
}

func newOutboxPublisher(c config.Config, bus realtimemodule.Bus) backgroundWorker {
	opts := c.Outbox
	driver := strings.ToLower(strings.TrimSpace(opts.Driver))
	switch driver {
	case "", "memory", "none", "disabled":
		return nil
	case "postgres":
		if strings.TrimSpace(opts.DataSource) == "" {
			opts.DataSource = c.Storage.DataSource
		}
		publisher, err := realtimemodule.NewOutboxPublisher(context.Background(), realtimemodule.OutboxOptions{
			DataSource:     opts.DataSource,
			IntervalMillis: opts.PollMillis,
			BatchSize:      opts.BatchSize,
		}, bus)
		if err != nil {
			logx.Must(err)
		}
		return publisher
	default:
		logx.Must(fmt.Errorf("unsupported outbox driver: %s", opts.Driver))
		return nil
	}
}

func newJobQueue(c config.Config, processor design.JobProcessor) (design.JobQueue, backgroundWorker) {
	opts := c.JobQueue
	driver := strings.ToLower(strings.TrimSpace(opts.Driver))
	switch driver {
	case "", "memory", "inline", "none", "disabled":
		return nil, nil
	case "asynq", "redis":
		if strings.TrimSpace(opts.Addr) == "" {
			opts.Addr = c.Cache.Addr
		}
		if strings.TrimSpace(opts.Password) == "" {
			opts.Password = c.Cache.Password
		}
		if strings.TrimSpace(opts.Addr) == "" {
			opts.Addr = "localhost:6379"
		}
		asynqOpts := jobmodule.AsynqOptions{
			Addr:                   opts.Addr,
			Password:               opts.Password,
			DB:                     opts.DB,
			Queue:                  opts.Queue,
			Concurrency:            opts.Concurrency,
			MaxRetry:               opts.MaxRetry,
			TimeoutSeconds:         opts.TimeoutSeconds,
			ShutdownTimeoutSeconds: opts.ShutdownTimeoutSeconds,
		}
		queue, err := jobmodule.NewAsynqQueue(context.Background(), asynqOpts)
		if err != nil {
			logx.Must(err)
		}
		worker, err := jobmodule.NewAsynqWorker(asynqOpts, processor)
		if err != nil {
			logx.Must(err)
		}
		return queue, worker
	default:
		logx.Must(fmt.Errorf("unsupported job queue driver: %s", opts.Driver))
		return nil, nil
	}
}

func newProjectRepository(c config.Config) design.Repository {
	switch c.Storage.Driver {
	case "", "memory":
		return memory.NewProjectRepository()
	case "postgres":
		if c.Storage.DataSource == "" {
			logx.Must(application.ErrMissingDataSource)
		}
		repo, err := postgres.NewProjectRepository(context.Background(), c.Storage.DataSource)
		if err != nil {
			logx.Must(err)
		}
		return repo
	default:
		logx.Must(application.ErrUnsupportedStorage)
		return memory.NewProjectRepository()
	}
}