main.go

package main

import (
	"errors"
	"flag"
	"fmt"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
	"jetbrains-ai-proxy/internal/apiserver"
	"jetbrains-ai-proxy/internal/config"
	"jetbrains-ai-proxy/internal/jetbrains"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
)

func main() {
	// 定义命令行参数
	configFile := flag.String("config", "", "配置文件路径")
	port := flag.Int("p", 0, "服务器监听端口 (覆盖配置文件)")
	host := flag.String("h", "", "服务器监听地址 (覆盖配置文件)")
	jwtTokens := flag.String("c", "", "JWT Tokens值，多个token用逗号分隔 (覆盖配置文件)")
	bearerToken := flag.String("k", "", "Bearer Token值 (覆盖配置文件)")
	loadBalanceStrategy := flag.String("s", "", "负载均衡策略: round_robin 或 random (覆盖配置文件)")
	generateConfig := flag.Bool("generate-config", false, "生成示例配置文件")
	printConfig := flag.Bool("print-config", false, "打印当前配置信息")

	flag.Usage = func() {
		fmt.Printf("用法: %s [选项]\n\n", flag.CommandLine.Name())
		fmt.Println("选项:")
		flag.PrintDefaults()
		fmt.Println("\n配置优先级 (从高到低):")
		fmt.Println("  1. 命令行参数")
		fmt.Println("  2. 环境变量")
		fmt.Println("  3. 配置文件")
		fmt.Println("  4. 默认值")
		fmt.Println("\n配置方式:")
		fmt.Println("  方式1 - 使用配置文件:")
		fmt.Println("    ./jetbrains-ai-proxy --generate-config  # 生成示例配置")
		fmt.Println("    # 编辑 config/config.json")
		fmt.Println("    ./jetbrains-ai-proxy")
		fmt.Println("")
		fmt.Println("  方式2 - 使用环境变量:")
		fmt.Println("    export JWT_TOKENS=\"jwt1,jwt2,jwt3\"")
		fmt.Println("    export BEARER_TOKEN=\"your_token\"")
		fmt.Println("    ./jetbrains-ai-proxy")
		fmt.Println("")
		fmt.Println("  方式3 - 使用命令行参数:")
		fmt.Println("    ./jetbrains-ai-proxy -c \"jwt1,jwt2,jwt3\" -k \"bearer_token\"")
		fmt.Println("")
		fmt.Println("负载均衡策略:")
		fmt.Println("  round_robin: 轮询策略（默认）")
		fmt.Println("  random: 随机策略")
	}

	flag.Parse()

	// 处理特殊命令
	if *generateConfig {
		if err := generateExampleConfig(); err != nil {
			log.Fatalf("Failed to generate config: %v", err)
		}
		return
	}

	// 获取配置管理器
	configManager := config.GetGlobalConfig()

	// 如果指定了配置文件，设置环境变量
	if *configFile != "" {
		os.Setenv("CONFIG_FILE", *configFile)
	}

	// 加载配置
	if err := configManager.LoadConfig(); err != nil {
		log.Printf("Warning: %v", err)
		log.Println("Continuing with command line arguments and environment variables...")
	}

	// 应用命令行参数覆盖
	applyCommandLineOverrides(configManager, port, host, jwtTokens, bearerToken, loadBalanceStrategy)

	// 打印配置信息
	if *printConfig {
		configManager.PrintConfig()
		return
	}

	// 验证配置
	if !configManager.HasJWTTokens() {
		log.Fatal("No JWT tokens configured. Use --generate-config to create example configuration.")
	}

	cfg := configManager.GetConfig()
	if cfg.BearerToken == "" {
		log.Fatal("Bearer token is required. Please configure it in config file, environment variable, or command line.")
	}

	// 初始化JWT负载均衡器
	if err := jetbrains.InitializeFromConfig(); err != nil {
		log.Fatalf("Failed to initialize JWT balancer: %v", err)
	}

	// 设置优雅关闭
	setupGracefulShutdown()

	// 启动配置文件监控
	discovery := config.NewConfigDiscovery(configManager)
	discovery.WatchConfig()

	// 创建Echo实例
	e := echo.New()
	e.Use(middleware.Logger())
	e.Use(middleware.Recover())

	// 添加管理端点
	setupManagementEndpoints(e, configManager)

	// 注册API路由
	apiserver.RegisterRoutes(e)

	// 启动服务器
	addr := fmt.Sprintf("%s:%d", cfg.ServerHost, cfg.ServerPort)
	log.Printf("Server starting on %s", addr)
	configManager.PrintConfig()

	if err := e.Start(addr); err != nil && !errors.Is(err, http.ErrServerClosed) {
		log.Fatalf("start server error: %v", err)
	}
}

// generateExampleConfig 生成示例配置
func generateExampleConfig() error {
	manager := config.NewManager()

	// 生成JSON配置文件
	if err := manager.GenerateExampleConfig("config/config.json"); err != nil {
		return fmt.Errorf("failed to generate JSON config: %v", err)
	}

	// 生成.env示例文件
	config.NewConfigDiscovery(manager)
	envContent := `# JetBrains AI Proxy Configuration
# Copy this file to .env and fill in your actual values

# Multiple JWT tokens (comma-separated)
JWT_TOKENS=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...,eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9...

# Bearer token for API authentication
BEARER_TOKEN=your_bearer_token_here

# Load balancing strategy: round_robin or random
LOAD_BALANCE_STRATEGY=round_robin

# Server configuration
SERVER_HOST=0.0.0.0
SERVER_PORT=8080
`

	if err := os.WriteFile(".env.example", []byte(envContent), 0644); err != nil {
		return fmt.Errorf("failed to generate .env example: %v", err)
	}

	fmt.Println("✅ Example configuration files generated:")
	fmt.Println("   📄 config/config.json - JSON configuration file")
	fmt.Println("   📄 .env.example - Environment variables example")
	fmt.Println("")
	fmt.Println("📝 Next steps:")
	fmt.Println("   1. Edit config/config.json with your JWT tokens")
	fmt.Println("   2. Or copy .env.example to .env and edit it")
	fmt.Println("   3. Run: ./jetbrains-ai-proxy")

	return nil
}

// applyCommandLineOverrides 应用命令行参数覆盖
func applyCommandLineOverrides(manager *config.Manager, port *int, host, jwtTokens, bearerToken, strategy *string) {
	if *jwtTokens != "" {
		manager.SetJWTTokens(*jwtTokens)
		log.Printf("JWT tokens overridden by command line")
	}

	if *bearerToken != "" {
		manager.SetBearerToken(*bearerToken)
		log.Printf("Bearer token overridden by command line")
	}

	if *strategy != "" {
		manager.SetLoadBalanceStrategy(*strategy)
		log.Printf("Load balance strategy overridden by command line: %s", *strategy)
	}

	// 覆盖服务器配置
	cfg := manager.GetConfig()
	if *port > 0 {
		cfg.ServerPort = *port
		log.Printf("Server port overridden by command line: %d", *port)
	}

	if *host != "" {
		cfg.ServerHost = *host
		log.Printf("Server host overridden by command line: %s", *host)
	}
}

// setupManagementEndpoints 设置管理端点
func setupManagementEndpoints(e *echo.Echo, manager *config.Manager) {
	// 健康检查端点
	e.GET("/health", func(c echo.Context) error {
		healthy, total := jetbrains.GetBalancerStats()
		cfg := manager.GetConfig()

		return c.JSON(http.StatusOK, map[string]interface{}{
			"status":         "ok",
			"healthy_tokens": healthy,
			"total_tokens":   total,
			"strategy":       cfg.LoadBalanceStrategy,
			"server_info": map[string]interface{}{
				"host": cfg.ServerHost,
				"port": cfg.ServerPort,
			},
		})
	})

	// 配置信息端点
	e.GET("/config", func(c echo.Context) error {
		discovery := config.NewConfigDiscovery(manager)
		summary := discovery.GetConfigSummary()
		return c.JSON(http.StatusOK, summary)
	})

	// 重载配置端点
	e.POST("/reload", func(c echo.Context) error {
		if err := jetbrains.ReloadConfig(); err != nil {
			return c.JSON(http.StatusInternalServerError, map[string]interface{}{
				"error": err.Error(),
			})
		}

		return c.JSON(http.StatusOK, map[string]interface{}{
			"message": "Configuration reloaded successfully",
		})
	})

	// 负载均衡器统计端点
	e.GET("/stats", func(c echo.Context) error {
		healthy, total := jetbrains.GetBalancerStats()
		cfg := manager.GetConfig()

		return c.JSON(http.StatusOK, map[string]interface{}{
			"balancer": map[string]interface{}{
				"healthy_tokens": healthy,
				"total_tokens":   total,
				"strategy":       cfg.LoadBalanceStrategy,
			},
			"config": map[string]interface{}{
				"health_check_interval": cfg.HealthCheckInterval.String(),
				"server_host":           cfg.ServerHost,
				"server_port":           cfg.ServerPort,
			},
		})
	})
}

// setupGracefulShutdown 设置优雅关闭
func setupGracefulShutdown() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)

	go func() {
		<-c
		log.Println("Shutting down gracefully...")
		jetbrains.StopBalancer()
		os.Exit(0)
	}()
}