gpt-load/cmd/gpt-load/main.go

// Package main provides the entry point for the GPT-Load proxy server
package main

import (
	"context"
	"embed"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"sync"
	"syscall"
	"time"

	"gpt-load/internal/channel"
	"gpt-load/internal/config"
	"gpt-load/internal/db"
	"gpt-load/internal/handler"
	"gpt-load/internal/models"
	"gpt-load/internal/proxy"
	"gpt-load/internal/router"
	"gpt-load/internal/services"
	"gpt-load/internal/store"
	"gpt-load/internal/types"

	"github.com/sirupsen/logrus"
	"gorm.io/gorm"
)

//go:embed dist
var buildFS embed.FS

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

func main() {
	// Load configuration
	configManager, err := config.NewManager()
	if err != nil {
		logrus.Fatalf("Failed to load configuration: %v", err)
	}

	// Setup logger
	setupLogger(configManager)

	// Initialize database
	database, err := db.InitDB(configManager)
	if err != nil {
		logrus.Fatalf("Failed to initialize database: %v", err)
	}

	// Initialize system settings after database is ready
	settingsManager := config.GetSystemSettingsManager()
	if err := settingsManager.InitializeSystemSettings(); err != nil {
		logrus.Fatalf("Failed to initialize system settings: %v", err)
	}
	logrus.Info("System settings initialized")

	// Display current system settings
	settingsManager.DisplayCurrentSettings()

	// Display startup information
	configManager.DisplayConfig()

	// Start log cleanup service
	logCleanupService := services.NewLogCleanupService()
	logCleanupService.Start()
	defer logCleanupService.Stop()

	// --- Asynchronous Request Logging Setup ---
	requestLogChan := make(chan models.RequestLog, 1000)
	var wg sync.WaitGroup
	wg.Add(1)
	go startRequestLogger(database, requestLogChan, &wg)
	// ---

	// --- Service Initialization ---
	// Initialize the store first, as other services depend on it.
	storage, err := store.NewStore(configManager)
	if err != nil {
		logrus.Fatalf("Failed to initialize store: %v", err)
	}
	defer storage.Close()

	taskService := services.NewTaskService(storage)
	channelFactory := channel.NewFactory(settingsManager)
	keyValidatorService := services.NewKeyValidatorService(database, channelFactory, settingsManager)

	// --- Global Key Validation Pool ---
	KeyValidationPool := services.NewKeyValidationPool(keyValidatorService, configManager)
	KeyValidationPool.Start()
	defer KeyValidationPool.Stop()
	// ---

	keyManualValidationService := services.NewKeyManualValidationService(database, keyValidatorService, taskService, settingsManager, configManager)
	keyCronService := services.NewKeyCronService(database, settingsManager, KeyValidationPool, storage)
	keyCronService.Start()
	defer keyCronService.Stop()

	keyService := services.NewKeyService(database)
	// ---

	// Create proxy server
	proxyServer, err := proxy.NewProxyServer(database, channelFactory, requestLogChan)
	if err != nil {
		logrus.Fatalf("Failed to create proxy server: %v", err)
	}
	defer proxyServer.Close()

	// Create handlers
	serverHandler := handler.NewServer(database, configManager, settingsManager, keyValidatorService, keyManualValidationService, taskService, keyService)
	logCleanupHandler := handler.NewLogCleanupHandler(logCleanupService)

	// Setup routes using the new router package
	appRouter := router.New(serverHandler, proxyServer, logCleanupHandler, configManager, buildFS, indexPage)

	// Create HTTP server with optimized timeout configuration
	serverConfig := configManager.GetEffectiveServerConfig()
	server := &http.Server{
		Addr:           fmt.Sprintf("%s:%d", serverConfig.Host, serverConfig.Port),
		Handler:        appRouter,
		ReadTimeout:    time.Duration(serverConfig.ReadTimeout) * time.Second,
		WriteTimeout:   time.Duration(serverConfig.WriteTimeout) * time.Second,
		IdleTimeout:    time.Duration(serverConfig.IdleTimeout) * time.Second,
		MaxHeaderBytes: 1 << 20, // 1MB header limit
	}

	// Start server
	go func() {
		logrus.Info("GPT-Load proxy server started successfully")
		logrus.Infof("Server address: http://%s:%d", serverConfig.Host, serverConfig.Port)
		logrus.Infof("Statistics: http://%s:%d/stats", serverConfig.Host, serverConfig.Port)
		logrus.Infof("Health check: http://%s:%d/health", serverConfig.Host, serverConfig.Port)
		logrus.Info("")

		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logrus.Fatalf("Server startup failed: %v", err)
		}
	}()

	// Wait for interrupt signal to gracefully shutdown the server
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	logrus.Info("Shutting down server...")

	// Give outstanding requests a deadline for completion
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(serverConfig.GracefulShutdownTimeout)*time.Second)
	defer cancel()

	// Attempt graceful shutdown
	if err := server.Shutdown(ctx); err != nil {
		logrus.Errorf("Server forced to shutdown: %v", err)
	}

	// Close the request log channel and wait for the logger to finish
	logrus.Info("Closing request log channel...")
	close(requestLogChan)
	wg.Wait()
	logrus.Info("All logs have been written.")

	logrus.Info("Server exited gracefully")
}

// setupLogger, displayStartupInfo, and startRequestLogger functions remain unchanged.
// The old setupRoutes and ServeUI functions are now removed from this file.

// setupLogger configures the logging system
func setupLogger(configManager types.ConfigManager) {
	logConfig := configManager.GetLogConfig()

	// Set log level
	level, err := logrus.ParseLevel(logConfig.Level)
	if err != nil {
		logrus.Warn("Invalid log level, using info")
		level = logrus.InfoLevel
	}
	logrus.SetLevel(level)

	// Set log format
	if logConfig.Format == "json" {
		logrus.SetFormatter(&logrus.JSONFormatter{
			TimestampFormat: time.RFC3339,
		})
	} else {
		logrus.SetFormatter(&logrus.TextFormatter{
			ForceColors:     true,
			FullTimestamp:   true,
			TimestampFormat: "2006-01-02 15:04:05",
		})
	}

	// Setup file logging if enabled
	if logConfig.EnableFile {
		// Create log directory if it doesn't exist
		logDir := filepath.Dir(logConfig.FilePath)
		if err := os.MkdirAll(logDir, 0755); err != nil {
			logrus.Warnf("Failed to create log directory: %v", err)
		} else {
			// Open log file
			logFile, err := os.OpenFile(logConfig.FilePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
			if err != nil {
				logrus.Warnf("Failed to open log file: %v", err)
			} else {
				// Use both file and stdout
				logrus.SetOutput(io.MultiWriter(os.Stdout, logFile))
			}
		}
	}
}

// startRequestLogger runs a background goroutine to batch-insert request logs.
func startRequestLogger(db *gorm.DB, logChan <-chan models.RequestLog, wg *sync.WaitGroup) {
	defer wg.Done()
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	logBuffer := make([]models.RequestLog, 0, 100)

	for {
		select {
		case logEntry, ok := <-logChan:
			if !ok {
				// Channel closed, flush remaining logs and exit
				if len(logBuffer) > 0 {
					if err := db.Create(&logBuffer).Error; err != nil {
						logrus.Errorf("Failed to write remaining request logs: %v", err)
					}
				}
				logrus.Info("Request logger stopped.")
				return
			}
			logBuffer = append(logBuffer, logEntry)
			if len(logBuffer) >= 100 {
				if err := db.Create(&logBuffer).Error; err != nil {
					logrus.Errorf("Failed to write request logs: %v", err)
				}
				logBuffer = make([]models.RequestLog, 0, 100) // Reset buffer
			}
		case <-ticker.C:
			// Flush logs periodically
			if len(logBuffer) > 0 {
				if err := db.Create(&logBuffer).Error; err != nil {
					logrus.Errorf("Failed to write request logs on tick: %v", err)
				}
				logBuffer = make([]models.RequestLog, 0, 100) // Reset buffer
			}
		}
	}
}