internal/jetbrains/client.go

package jetbrains

import (
	"context"
	"fmt"
	"github.com/go-resty/resty/v2"
	"jetbrains-ai-proxy/internal/balancer"
	"jetbrains-ai-proxy/internal/config"
	"jetbrains-ai-proxy/internal/types"
	"jetbrains-ai-proxy/internal/utils"
	"log"
	"sync"
)

var (
	jwtBalancer    balancer.JWTBalancer
	healthChecker  *balancer.HealthChecker
	initOnce       sync.Once
	configManager  *config.Manager
)

// InitializeFromConfig 从配置管理器初始化JWT负载均衡器
func InitializeFromConfig() error {
	var initErr error

	initOnce.Do(func() {
		configManager = config.GetGlobalConfig()

		// 加载配置
		if err := configManager.LoadConfig(); err != nil {
			initErr = fmt.Errorf("failed to load config: %v", err)
			return
		}

		// 获取配置
		cfg := configManager.GetConfig()
		tokens := configManager.GetJWTTokens()

		if len(tokens) == 0 {
			initErr = fmt.Errorf("no JWT tokens configured")
			return
		}

		// 创建负载均衡器
		jwtBalancer = balancer.NewJWTBalancer(tokens, cfg.LoadBalanceStrategy)

		// 创建并启动健康检查器
		healthChecker = balancer.NewHealthChecker(jwtBalancer)
		if cfg.HealthCheckInterval > 0 {
			healthChecker.SetCheckInterval(cfg.HealthCheckInterval)
		}
		healthChecker.Start()

		log.Printf("JWT balancer initialized from config:")
		log.Printf("  - Tokens: %d", len(tokens))
		log.Printf("  - Strategy: %s", cfg.LoadBalanceStrategy)
		log.Printf("  - Health check interval: %v", cfg.HealthCheckInterval)
	})

	return initErr
}

// InitializeBalancer 初始化JWT负载均衡器（向后兼容）
func InitializeBalancer(tokens []string, strategy string) error {
	if len(tokens) == 0 {
		return fmt.Errorf("no JWT tokens provided")
	}

	var balanceStrategy config.LoadBalanceStrategy
	switch strategy {
	case "random":
		balanceStrategy = config.Random
	case "round_robin", "":
		balanceStrategy = config.RoundRobin
	default:
		balanceStrategy = config.RoundRobin
	}

	// 创建负载均衡器
	jwtBalancer = balancer.NewJWTBalancer(tokens, balanceStrategy)

	// 创建并启动健康检查器
	healthChecker = balancer.NewHealthChecker(jwtBalancer)
	healthChecker.Start()

	log.Printf("JWT balancer initialized with %d tokens, strategy: %s", len(tokens), string(balanceStrategy))
	return nil
}

// ReloadConfig 重新加载配置
func ReloadConfig() error {
	if configManager == nil {
		return fmt.Errorf("config manager not initialized")
	}

	// 重新加载配置
	if err := configManager.LoadConfig(); err != nil {
		return fmt.Errorf("failed to reload config: %v", err)
	}

	// 获取新配置
	cfg := configManager.GetConfig()
	tokens := configManager.GetJWTTokens()

	if len(tokens) == 0 {
		return fmt.Errorf("no JWT tokens in reloaded config")
	}

	// 更新负载均衡器
	if jwtBalancer != nil {
		jwtBalancer.RefreshTokens(tokens)
	}

	// 更新健康检查间隔
	if healthChecker != nil && cfg.HealthCheckInterval > 0 {
		healthChecker.SetCheckInterval(cfg.HealthCheckInterval)
	}

	log.Printf("Config reloaded successfully:")
	log.Printf("  - Tokens: %d", len(tokens))
	log.Printf("  - Strategy: %s", cfg.LoadBalanceStrategy)

	return nil
}

// StopBalancer 停止负载均衡器
func StopBalancer() {
	if healthChecker != nil {
		healthChecker.Stop()
	}
}

// GetConfigManager 获取配置管理器
func GetConfigManager() *config.Manager {
	return configManager
}

func SendJetbrainsRequest(ctx context.Context, req *types.JetbrainsRequest) (*resty.Response, error) {
	// 获取一个可用的JWT token
	token, err := jwtBalancer.GetToken()
	if err != nil {
		log.Printf("failed to get JWT token: %v", err)
		return nil, fmt.Errorf("no available JWT tokens: %v", err)
	}

	resp, err := utils.RestySSEClient.R().
		SetContext(ctx).
		SetHeader(types.JwtTokenKey, token).
		SetDoNotParseResponse(true).
		SetBody(req).
		Post(types.ChatStreamV7)

	if err != nil {
		log.Printf("jetbrains ai req error: %v", err)
		// 标记token为不健康
		jwtBalancer.MarkTokenUnhealthy(token)
		return nil, err
	}

	// 检查响应状态码
	if resp.StatusCode() == 401 {
		// 401表示token无效，标记为不健康
		jwtBalancer.MarkTokenUnhealthy(token)
		log.Printf("JWT token invalid (401): %s...", token[:min(len(token), 10)])
		return nil, fmt.Errorf("JWT token invalid")
	} else if resp.StatusCode() == 200 {
		// 成功响应，确保token标记为健康
		jwtBalancer.MarkTokenHealthy(token)
	}

	return resp, nil
}

// GetBalancerStats 获取负载均衡器统计信息
func GetBalancerStats() (int, int) {
	if jwtBalancer == nil {
		return 0, 0
	}
	return jwtBalancer.GetHealthyTokenCount(), jwtBalancer.GetTotalTokenCount()
}

// min 辅助函数
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}