main.go

package main

import (
	"bytes"
	"context"
	"embed"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/MAX-API-Next/MAX-API/common"
	"github.com/MAX-API-Next/MAX-API/constant"
	"github.com/MAX-API-Next/MAX-API/controller"
	"github.com/MAX-API-Next/MAX-API/i18n"
	"github.com/MAX-API-Next/MAX-API/logger"
	"github.com/MAX-API-Next/MAX-API/middleware"
	"github.com/MAX-API-Next/MAX-API/model"
	"github.com/MAX-API-Next/MAX-API/oauth"
	perfmetrics "github.com/MAX-API-Next/MAX-API/pkg/perf_metrics"
	"github.com/MAX-API-Next/MAX-API/relay"
	"github.com/MAX-API-Next/MAX-API/router"
	"github.com/MAX-API-Next/MAX-API/service"
	_ "github.com/MAX-API-Next/MAX-API/setting/performance_setting"
	"github.com/MAX-API-Next/MAX-API/setting/ratio_setting"
	"github.com/MAX-API-Next/MAX-API/setting/system_setting"

	"github.com/bytedance/gopkg/util/gopool"
	"github.com/gin-contrib/sessions"
	"github.com/gin-contrib/sessions/cookie"
	"github.com/gin-gonic/gin"
	"github.com/joho/godotenv"

	_ "net/http/pprof"
)

//go:embed web/default/dist
var buildFS embed.FS

//go:embed web/default/dist/index.html
var indexPage []byte

func main() {
	startTime := time.Now()

	err := InitResources()
	if err != nil {
		common.FatalLog("failed to initialize resources: " + err.Error())
		return
	}

	common.SysLog("MAX API " + common.Version + " started")
	if os.Getenv("GIN_MODE") != "debug" {
		gin.SetMode(gin.ReleaseMode)
	}
	if common.DebugEnabled {
		common.SysLog("running in debug mode")
	}

	defer func() {
		err := model.CloseDB()
		if err != nil {
			common.FatalLog("failed to close database: " + err.Error())
		}
	}()
	backgroundCtx, stopBackground := context.WithCancel(context.Background())
	defer stopBackground()

	// Recover durable balance settlements after process restarts or transient DB failures.
	model.StartBillingSettlementTaskRunner()

	if common.RedisEnabled {
		// for compatibility with old versions
		common.MemoryCacheEnabled = true
		model.StartCacheInvalidationTaskRunner()
	}
	if common.MemoryCacheEnabled {
		common.SysLog("memory cache enabled")
		common.SysLog(fmt.Sprintf("sync frequency: %d seconds", common.SyncFrequency))

		// Add panic recovery and retry for InitChannelCache
		func() {
			defer func() {
				if r := recover(); r != nil {
					common.SysLog(fmt.Sprintf("InitChannelCache panic: %v, retrying once", r))
					// Retry once
					_, _, fixErr := model.FixAbility()
					if fixErr != nil {
						common.FatalLog(fmt.Sprintf("InitChannelCache failed: %s", fixErr.Error()))
					}
				}
			}()
			model.InitChannelCache()
		}()

		go model.SyncChannelCache(common.SyncFrequency)
	}

	// 热更新配置
	go model.SyncOptions(common.SyncFrequency)

	// 数据看板
	go model.UpdateQuotaData()

	if os.Getenv("CHANNEL_UPDATE_FREQUENCY") != "" {
		frequency, err := strconv.Atoi(os.Getenv("CHANNEL_UPDATE_FREQUENCY"))
		if err != nil {
			common.FatalLog("failed to parse CHANNEL_UPDATE_FREQUENCY: " + err.Error())
		}
		go controller.AutomaticallyUpdateChannels(frequency)
	}

	go controller.AutomaticallyTestChannels()

	// Codex credential auto-refresh check every 10 minutes, refresh when expires within 1 day
	service.StartCodexCredentialAutoRefreshTask()

	// Subscription quota reset task (daily/weekly/monthly/custom)
	service.StartSubscriptionQuotaResetTask()

	// System instance heartbeat for multi-node operations visibility
	service.StartSystemInstanceReporter()

	// Proactively detect sustained system resource pressure and notify admins.
	service.StartSmartOpsAlertMonitor()

	// Persistent system maintenance task runner
	service.StartSystemTaskRunner()
	authFlowCleanupDone := service.StartAuthFlowCleanupWithContext(backgroundCtx)

	// Wire task polling adaptor factory (breaks service -> relay import cycle)
	service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor {
		a := relay.GetTaskAdaptor(platform)
		if a == nil {
			return nil
		}
		return a
	}

	// Channel upstream model update check task
	controller.StartChannelUpstreamModelUpdateTask()

	if common.IsMasterNode && constant.UpdateTask {
		gopool.Go(func() {
			controller.UpdateMidjourneyTaskBulk()
		})
		gopool.Go(func() {
			controller.UpdateTaskBulk()
		})
	}
	if os.Getenv("BATCH_UPDATE_ENABLED") == "true" {
		common.BatchUpdateEnabled = true
		common.SysLog("batch update enabled with interval " + strconv.Itoa(common.BatchUpdateInterval) + "s")
		model.InitBatchUpdater()
	}

	var cpuMonitorDone <-chan struct{}
	if os.Getenv("ENABLE_PPROF") == "true" {
		gopool.Go(func() {
			log.Println(http.ListenAndServe("0.0.0.0:8005", nil))
		})
		done := make(chan struct{})
		cpuMonitorDone = done
		go func() {
			defer close(done)
			common.MonitorContext(backgroundCtx)
		}()
		common.SysLog("pprof enabled")
	}

	err = common.StartPyroScope()
	if err != nil {
		common.SysError(fmt.Sprintf("start pyroscope error : %v", err))
	}

	// Initialize HTTP server
	server := gin.New()
	if err := server.SetTrustedProxies(common.TrustedProxies); err != nil {
		common.FatalLog("failed to configure trusted proxies: " + err.Error())
	}
	server.Use(gin.CustomRecovery(func(c *gin.Context, err any) {
		common.SysLog(fmt.Sprintf("panic detected: %v", err))
		c.JSON(http.StatusInternalServerError, gin.H{
			"error": gin.H{
				"message": fmt.Sprintf("Panic detected, error: %v. Please submit a issue here: https://github.com/MAX-API-Next/MAX-API", err),
				"type":    "max_api_panic",
			},
		})
	}))
	// This will cause SSE not to work!!!
	//server.Use(gzip.Gzip(gzip.DefaultCompression))
	server.Use(middleware.RequestId())
	server.Use(middleware.PoweredBy())
	server.Use(middleware.I18n())
	middleware.SetUpLogger(server)
	// Initialize session store
	if err := configureSessionCookieSecure(system_setting.ServerAddress); err != nil {
		common.FatalLog(err.Error())
	}
	store := cookie.NewStore([]byte(common.SessionSecret))
	store.Options(sessionCookieOptions())
	server.Use(sessions.Sessions("session", store))

	InjectUmamiAnalytics()
	InjectGoogleAnalytics()

	// 设置路由
	router.SetRouter(server, router.ThemeAssets{
		DefaultBuildFS:   buildFS,
		DefaultIndexPage: indexPage,
	})
	var port = os.Getenv("PORT")
	if port == "" {
		port = strconv.Itoa(*common.Port)
	}

	srv := &http.Server{
		Addr:    ":" + port,
		Handler: server,
	}

	ln, err := net.Listen("tcp", srv.Addr)
	if err != nil {
		common.FatalLog("failed to bind HTTP server: " + err.Error())
	}

	go func() {
		if err := srv.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
			common.FatalLog("failed to start HTTP server: " + err.Error())
		}
	}()

	common.LogStartupSuccess(startTime, port)

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	sig := <-quit
	common.SysLog(fmt.Sprintf("received signal: %v, shutting down...", sig))

	shutdownTimeout := time.Duration(common.GetEnvOrDefault("SHUTDOWN_TIMEOUT_SECONDS", 120)) * time.Second
	ctx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
	defer cancel()
	stopBackground()
	shutdownHTTPServer(ctx, srv)
	waitForBackgroundRunner(ctx, "auth flow cleanup", authFlowCleanupDone)
	waitForBackgroundRunner(ctx, "CPU profile monitor", cpuMonitorDone)
	if err := common.StopSystemMonitor(ctx); err != nil {
		common.SysError(fmt.Sprintf("timed out stopping system monitor: %v", err))
	}
	if err := service.StopSmartOpsAlertMonitor(ctx); err != nil {
		common.SysError(fmt.Sprintf("timed out stopping smart ops alert monitor: %v", err))
	}
	if err := middleware.StopInMemoryRateLimiter(ctx); err != nil {
		common.SysError(fmt.Sprintf("timed out stopping in-memory rate limiter: %v", err))
	}
	if err := model.StopBillingSettlementTaskRunner(ctx); err != nil {
		common.SysError(fmt.Sprintf("timed out stopping billing settlement runner: %v", err))
	}
	if common.RedisEnabled {
		if err := model.StopCacheInvalidationTaskRunner(ctx); err != nil {
			common.SysError(fmt.Sprintf("timed out stopping cache invalidation runner: %v", err))
		}
	}
	if common.DataExportEnabled {
		saveTimeout := time.Duration(common.GetEnvOrDefault("QUOTA_DATA_CACHE_SAVE_TIMEOUT_SECONDS", 30)) * time.Second
		if !runWithTimeout(saveTimeout, model.WaitPendingLogQuotaData) {
			common.SysError(fmt.Sprintf("timed out waiting for pending quota data export after %s", saveTimeout))
		}
		saveCtx, saveCancel := context.WithTimeout(context.Background(), saveTimeout)
		saveErr := runWithContext(saveCtx, model.SaveQuotaDataCache)
		saveCancel()
		if saveErr != nil {
			if errors.Is(saveErr, context.DeadlineExceeded) || errors.Is(saveErr, context.Canceled) {
				common.SysError(fmt.Sprintf("timed out waiting for quota data cache save after %s", saveTimeout))
			} else {
				common.SysError(fmt.Sprintf("failed to save quota data cache during shutdown: %v", saveErr))
			}
		}
	}
	common.SysLog("server exited")
}

func sessionCookieOptions() sessions.Options {
	return sessions.Options{
		Path:     "/",
		MaxAge:   2592000, // 30 days
		HttpOnly: true,
		Secure:   common.SessionCookieSecure,
		// OAuth providers return through a cross-site top-level GET redirect.
		SameSite: http.SameSiteLaxMode,
	}
}

func configureSessionCookieSecure(serverAddress string) error {
	value, explicitlyConfigured := os.LookupEnv("SESSION_COOKIE_SECURE")
	value = strings.TrimSpace(value)
	if explicitlyConfigured && value != "" {
		secure, err := strconv.ParseBool(value)
		if err != nil {
			return fmt.Errorf("invalid SESSION_COOKIE_SECURE value %q: %w", value, err)
		}
		common.SessionCookieSecure = secure
		return nil
	}

	common.SessionCookieSecure = common.SessionCookieSecureForServerAddress(serverAddress)
	if !common.SessionCookieSecure {
		common.SysLog("SESSION_COOKIE_SECURE is not set and ServerAddress is not HTTPS; session cookies may traverse plain HTTP. Set SESSION_COOKIE_SECURE=true for HTTPS deployments")
	}
	return nil
}

func shutdownHTTPServer(ctx context.Context, srv *http.Server) {
	if err := srv.Shutdown(ctx); err != nil {
		common.SysError(fmt.Sprintf("server forced to shutdown: %v", err))
		if closeErr := srv.Close(); closeErr != nil && !errors.Is(closeErr, http.ErrServerClosed) {
			common.SysError(fmt.Sprintf("server close after forced shutdown failed: %v", closeErr))
		}
	}
}

func waitForBackgroundRunner(ctx context.Context, name string, done <-chan struct{}) {
	if done == nil {
		return
	}
	select {
	case <-done:
		return
	default:
	}
	select {
	case <-done:
	case <-ctx.Done():
		common.SysError(fmt.Sprintf("timed out stopping %s: %v", name, ctx.Err()))
	}
}

func runWithTimeout(timeout time.Duration, fn func()) bool {
	done := make(chan struct{})
	go func() {
		defer close(done)
		defer func() {
			if r := recover(); r != nil {
				common.SysError(fmt.Sprintf("runWithTimeout: recovered panic: %v", r))
			}
		}()
		fn()
	}()

	timer := time.NewTimer(timeout)
	defer timer.Stop()

	select {
	case <-done:
		return true
	case <-timer.C:
		return false
	}
}

func runWithContext(ctx context.Context, fn func(context.Context) error) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("runWithContext: recovered panic: %v", r)
		}
	}()
	return fn(ctx)
}

func InjectUmamiAnalytics() {
	analyticsInjectBuilder := &strings.Builder{}
	if os.Getenv("UMAMI_WEBSITE_ID") != "" {
		umamiSiteID := os.Getenv("UMAMI_WEBSITE_ID")
		umamiScriptURL := os.Getenv("UMAMI_SCRIPT_URL")
		if umamiScriptURL == "" {
			umamiScriptURL = "https://analytics.umami.is/script.js"
		}
		analyticsInjectBuilder.WriteString("<script defer src=\"")
		analyticsInjectBuilder.WriteString(umamiScriptURL)
		analyticsInjectBuilder.WriteString("\" data-website-id=\"")
		analyticsInjectBuilder.WriteString(umamiSiteID)
		analyticsInjectBuilder.WriteString("\"></script>")
	}
	analyticsInjectBuilder.WriteString("<!--Umami MAX-API-Next-->\n")
	analyticsInject := []byte(analyticsInjectBuilder.String())
	placeholder := []byte("<!--umami-->\n")
	indexPage = bytes.ReplaceAll(indexPage, placeholder, analyticsInject)
}

func InjectGoogleAnalytics() {
	analyticsInjectBuilder := &strings.Builder{}
	if os.Getenv("GOOGLE_ANALYTICS_ID") != "" {
		gaID := os.Getenv("GOOGLE_ANALYTICS_ID")
		// Google Analytics 4 (gtag.js)
		analyticsInjectBuilder.WriteString("<script async src=\"https://www.googletagmanager.com/gtag/js?id=")
		analyticsInjectBuilder.WriteString(gaID)
		analyticsInjectBuilder.WriteString("\"></script>")
		analyticsInjectBuilder.WriteString("<script>")
		analyticsInjectBuilder.WriteString("window.dataLayer = window.dataLayer || [];")
		analyticsInjectBuilder.WriteString("function gtag(){dataLayer.push(arguments);}")
		analyticsInjectBuilder.WriteString("gtag('js', new Date());")
		analyticsInjectBuilder.WriteString("gtag('config', '")
		analyticsInjectBuilder.WriteString(gaID)
		analyticsInjectBuilder.WriteString("');")
		analyticsInjectBuilder.WriteString("</script>")
	}
	analyticsInjectBuilder.WriteString("<!--Google Analytics MAX-API-Next-->\n")
	analyticsInject := []byte(analyticsInjectBuilder.String())
	placeholder := []byte("<!--Google Analytics-->\n")
	indexPage = bytes.ReplaceAll(indexPage, placeholder, analyticsInject)
}

func InitResources() error {
	// Initialize resources here if needed
	// This is a placeholder function for future resource initialization
	err := godotenv.Load(".env")
	if err != nil {
		if common.DebugEnabled {
			common.SysLog("No .env file found, using default environment variables. If needed, please create a .env file and set the relevant variables.")
		}
	}

	// 加载环境变量
	common.InitEnv()

	logger.SetupLogger()

	// Initialize model settings
	ratio_setting.InitRatioSettings()

	service.InitHttpClient()

	service.InitTokenEncoders()

	// Initialize SQL Database
	err = model.InitDB()
	if err != nil {
		common.FatalLog("failed to initialize database: " + err.Error())
		return err
	}

	model.CheckSetup()
	if !common.IsMasterNode {
		if err := service.WaitForTaskPluginCatalogSchema(context.Background(), model.DB); err != nil {
			return fmt.Errorf("wait for task plugin catalog schema: %w", err)
		}
	}
	if err := service.InitializeTaskPluginCatalog(context.Background(), model.DB); err != nil {
		return fmt.Errorf("initialize task plugin catalog: %w", err)
	}

	// Initialize options, should after model.InitDB()
	model.InitOptionMap()

	// 清理旧的磁盘缓存文件
	common.CleanupOldCacheFiles()

	// 初始化模型
	model.GetPricing()

	// Initialize SQL Database
	err = model.InitLogDB()
	if err != nil {
		return err
	}

	// Initialize Redis
	err = common.InitRedisClient()
	if err != nil {
		return err
	}

	perfmetrics.Init()

	// 启动系统监控
	common.StartSystemMonitor()

	// Initialize i18n
	err = i18n.Init()
	if err != nil {
		common.SysError("failed to initialize i18n: " + err.Error())
		// Don't return error, i18n is not critical
	} else {
		common.SysLog("i18n initialized with languages: " + strings.Join(i18n.SupportedLanguages(), ", "))
	}
	// Register user language loader for lazy loading
	i18n.SetUserLangLoader(model.GetUserLanguage)

	// Load custom OAuth providers from database
	err = oauth.LoadCustomProviders()
	if err != nil {
		common.SysError("failed to load custom OAuth providers: " + err.Error())
		// Don't return error, custom OAuth is not critical
	}

	return nil
}