Veyra/main.go

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"git.sundynix.cn/Blizzard/Veyra/internal/logging"
	"git.sundynix.cn/Blizzard/Veyra/internal/market"
	"git.sundynix.cn/Blizzard/Veyra/internal/research"
	"go.uber.org/zap"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/redis/go-redis/v9"

	"git.sundynix.cn/Blizzard/Veyra/internal/config"
	"git.sundynix.cn/Blizzard/Veyra/internal/database"
	"git.sundynix.cn/Blizzard/Veyra/internal/handler"
	"git.sundynix.cn/Blizzard/Veyra/internal/ratelimit"
	"git.sundynix.cn/Blizzard/Veyra/internal/repository"
	"git.sundynix.cn/Blizzard/Veyra/internal/secret"
	"git.sundynix.cn/Blizzard/Veyra/internal/service"
)

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	if err := run(ctx, os.Args[1:]); err != nil && !errors.Is(err, flag.ErrHelp) {
		fmt.Fprintln(os.Stderr, "Veyra stopped:", err)
		os.Exit(1)
	}
}

// One executable owns service startup and explicit maintenance operations.
func run(ctx context.Context, args []string) (err error) {
	flags := flag.NewFlagSet("veyra", flag.ContinueOnError)
	envFile := flags.String("env-file", "configs/.env", "environment file (empty disables file loading)")
	command := flags.String("command", "serve", "serve, db-create, db-migrate, user-create, or admin-set")
	previousLogin := flags.String("previous-login", "", "current administrator login for admin-set (defaults to ADMIN_LOGIN)")
	role := flags.String("role", "ADMIN", "ADMIN or USER for user-create")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if flags.NArg() != 0 {
		return fmt.Errorf("unexpected positional arguments")
	}
	switch *command {
	case "serve", "db-create", "db-migrate", "user-create", "admin-set":
	default:
		return fmt.Errorf("unsupported command")
	}
	if *envFile != "" {
		err := config.LoadEnvFile(*envFile)
		// Deployments can use environment variables without a local .env.
		if err != nil && !(*envFile == "configs/.env" && errors.Is(err, os.ErrNotExist)) {
			return err
		}
	}
	c, err := config.Load()
	if err != nil {
		return err
	}
	logger, closeLogs, err := logging.New(c.Logging)
	if err != nil {
		return err
	}
	defer closeLogs()
	defer func() {
		if err != nil {
			logger.Error("command failed", zap.String("command", *command), zap.Error(err))
		}
	}()
	switch *command {
	case "db-create", "db-migrate":
		return maintainDatabase(ctx, c, *command, logger)
	case "admin-set":
		return setAdministrator(ctx, c, *previousLogin, logger)
	case "user-create":
		return provisionUser(ctx, c, *role, logger)
	default:
		return serve(ctx, c, logger)
	}
}

func serve(ctx context.Context, c config.Config, logger *zap.Logger) error {
	var postgres service.Probe
	var authRepo service.AuthRepository
	var accountRepo service.AccountRepository
	var providerRepo service.ProviderRepository
	var candidateRepo service.CandidateRepository
	var planRepo service.TradePlanRepository
	var fillRepo service.TradeFillRepository
	var notificationRepo service.NotificationRepository
	var strategyRepo service.StrategyRepository
	var backtestRepo service.BacktestRepository
	var marketRepo service.MarketRepository
	var screeningRepo service.ScreeningRepository
	vault, err := secret.NewVault(c.SecretEncryptionKey)
	if err != nil {
		return err
	}
	if c.PGUser != "" {
		if err := c.ValidateDatabase(); err != nil {
			return err
		}
		pool, err := pgxpool.New(ctx, c.DatabaseURL(c.PGDatabase))
		if err != nil {
			return errors.New("invalid PostgreSQL configuration")
		}
		defer pool.Close()
		postgres = pool.Ping
		authRepo = repository.NewIdentityRepository(pool)
		accountRepo = repository.NewAccountRepository(pool)
		providerRepo = repository.NewProviderRepository(pool)
		candidateRepo = repository.NewCandidateRepository(pool)
		planRepo = repository.NewTradePlanRepository(pool)
		fillRepo = repository.NewTradeFillRepository(pool)
		notificationRepo = repository.NewNotificationRepository(pool)
		strategyRepo = repository.NewStrategyRepository(pool)
		backtestRepo = repository.NewBacktestRepository(pool)
		marketRepo = repository.NewMarketRepository(pool)
		screeningRepo = repository.NewScreeningRepository(pool)
	}
	redisClient := redis.NewClient(&redis.Options{Addr: c.RedisAddr, Username: c.RedisUsername, Password: c.RedisPassword, DB: c.RedisDB, DialTimeout: time.Second, ReadTimeout: time.Second, WriteTimeout: time.Second, MaxRetries: -1, ContextTimeoutEnabled: true})
	defer redisClient.Close()
	health := service.NewHealthService(postgres, func(ctx context.Context) error { return redisClient.Ping(ctx).Err() })
	auth, err := service.NewAuthService(authRepo)
	if err != nil {
		return err
	}
	development := c.Environment == "local" || c.Environment == "test"
	marketService := service.NewMarketService(marketRepo, providerRepo, market.New(), development)
	screenings := service.NewScreeningService(screeningRepo, marketService, development)
	notifications := service.NewNotificationService(notificationRepo, development)
	backtests := service.NewBacktestService(backtestRepo, research.New(c.PythonExecutable), development, marketService)
	workerContext, workerCancel := context.WithCancel(ctx)
	workerDone := make(chan struct{})
	go func() { defer close(workerDone); notifications.Run(workerContext, logger) }()
	researchDone := make(chan struct{})
	go func() { defer close(researchDone); backtests.Run(workerContext, logger) }()
	marketDone := make(chan struct{})
	go func() { defer close(marketDone); marketService.Run(workerContext, logger) }()
	screeningDone := make(chan struct{})
	go func() { defer close(screeningDone); screenings.Run(workerContext, logger) }()
	defer func() { workerCancel(); <-workerDone; <-researchDone; <-marketDone; <-screeningDone }()
	identityOptions := handler.IdentityOptions{Screenings: screenings, Market: marketService, Backtests: backtests, Strategies: service.NewStrategyService(strategyRepo, development), Notifications: notifications, Auth: auth, Fills: service.NewTradeFillService(fillRepo, development), Plans: service.NewTradePlanService(planRepo, development), Candidates: service.NewCandidateService(candidateRepo, development), Providers: service.NewProviderService(providerRepo, vault), Accounts: service.NewAccountService(accountRepo), Limiter: ratelimit.NewRedis(redisClient), DevelopmentLogin: development, SecureCookie: !development || strings.HasPrefix(c.PublicOrigin, "https://"), PublicOrigin: c.PublicOrigin}
	server := &http.Server{Addr: c.HTTPAddr, Handler: handler.NewRouter(health, logger, identityOptions), ReadHeaderTimeout: 5 * time.Second, ReadTimeout: 10 * time.Second, WriteTimeout: 10 * time.Second, IdleTimeout: 60 * time.Second, MaxHeaderBytes: 1 << 20}
	result := make(chan error, 1)
	go func() { result <- server.ListenAndServe() }()
	logger.Info("API starting", zap.String("address", c.HTTPAddr), zap.String("environment", c.Environment))
	select {
	case err := <-result:
		if errors.Is(err, http.ErrServerClosed) {
			return nil
		}
		return err
	case <-ctx.Done():
		shutdown, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		return server.Shutdown(shutdown)
	}
}

func maintainDatabase(ctx context.Context, c config.Config, command string, logger *zap.Logger) error {
	ctx, cancel := context.WithTimeout(ctx, 60*time.Second)
	defer cancel()
	switch command {
	case "db-create":
		created, err := database.Create(ctx, c)
		if err != nil {
			return err
		}
		logger.Info("database ready", zap.String("database", c.PGDatabase), zap.Bool("created", created))
	case "db-migrate":
		pool, err := database.Open(ctx, c)
		if err != nil {
			return err
		}
		defer pool.Close()
		if err := database.Migrate(ctx, pool); err != nil {
			return err
		}
		logger.Info("migrations applied", zap.String("database", c.PGDatabase))
	default:
		return fmt.Errorf("command must be db-create or db-migrate")
	}
	return nil
}

func provisionUser(ctx context.Context, c config.Config, role string, logger *zap.Logger) error {
	if c.Environment != "local" && c.Environment != "test" {
		return fmt.Errorf("development provisioning is disabled outside local/test environments")
	}
	var login, password, name string
	switch role {
	case "ADMIN":
		login = os.Getenv("ADMIN_LOGIN")
		password = os.Getenv("ADMIN_PASSWORD")
		name = "管理员"
	case "USER":
		login = os.Getenv("DEV_USER_LOGIN")
		password = os.Getenv("DEV_USER_PASSWORD")
		name = "开发测试用户"
	default:
		return fmt.Errorf("role must be ADMIN or USER")
	}
	if login == "" || password == "" {
		return fmt.Errorf("configure the selected identity login and password in the local environment file")
	}
	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()
	pool, err := database.Open(ctx, c)
	if err != nil {
		return err
	}
	defer pool.Close()
	auth, err := service.NewAuthService(repository.NewIdentityRepository(pool))
	if err != nil {
		return err
	}
	user, err := auth.ProvisionUser(ctx, login, password, name, role)
	if err != nil {
		return err
	}
	logger.Info("identity provisioned", zap.String("user_id", user.ID), zap.String("role", role))
	return nil
}

func setAdministrator(ctx context.Context, c config.Config, previousLogin string, logger *zap.Logger) error {
	login, password := os.Getenv("ADMIN_LOGIN"), os.Getenv("ADMIN_PASSWORD")
	if previousLogin == "" {
		previousLogin = login
	}
	if login == "" || password == "" {
		return fmt.Errorf("configure administrator credentials in the local environment file")
	}
	if c.Environment != "local" && c.Environment != "test" {
		return fmt.Errorf("development administrator setup is disabled outside local/test")
	}
	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()
	pool, err := database.Open(ctx, c)
	if err != nil {
		return err
	}
	defer pool.Close()
	updater := service.NewAdminCredentialService(repository.NewIdentityRepository(pool), true)
	userID, err := updater.Set(ctx, previousLogin, login, password)
	if err != nil {
		return err
	}
	logger.Info("administrator credential updated", zap.String("user_id", userID))
	return nil
}