package balancer

import (
	"context"
	"log/slog"
	"sync"
	"time"

	"github.com/chenxin/hunjing/mcp/web-search/internal/config"
	"github.com/chenxin/hunjing/mcp/web-search/internal/models"
)

// HealthChecker 定期对已注册 provider 执行健康检查。
// 连续失败超过阈值时自动禁用 provider，恢复成功后自动启用。
type HealthChecker struct {
	lb       *LoadBalancer
	interval time.Duration
	timeout  time.Duration
	threshold int

	mu        sync.Mutex
	failures  map[models.ProviderKind]int // 连续失败计数
	done      chan struct{}
}

// NewHealthChecker 创建健康检查器。cfg 为对应 provider 的健康检查配置。
func NewHealthChecker(lb *LoadBalancer, cfg config.HealthCheckConf) *HealthChecker {
	return &HealthChecker{
		lb:        lb,
		interval:  cfg.Interval,
		timeout:   cfg.Timeout,
		threshold: cfg.FailureThreshold,
		failures:  make(map[models.ProviderKind]int),
		done:      make(chan struct{}),
	}
}

// Start 在后台 goroutine 中启动健康检查循环。
// 每隔 interval 时间对所有已注册 provider 执行一次健康检查。
func (hc *HealthChecker) Start() {
	if hc.interval <= 0 {
		return
	}
	slog.Info("健康检查已启动", "interval", hc.interval, "timeout", hc.timeout, "failure_threshold", hc.threshold)
	go hc.loop()
}

// Stop 停止健康检查循环。
func (hc *HealthChecker) Stop() {
	close(hc.done)
}

func (hc *HealthChecker) loop() {
	ticker := time.NewTicker(hc.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			hc.checkAll()
		case <-hc.done:
			return
		}
	}
}

func (hc *HealthChecker) checkAll() {
	// 获取当前所有已注册 provider 的名称
	// 我们需要从 LoadBalancer 获取提供者列表
	// 由于 LoadBalancer 没有直接暴露 provider 列表的 API，我们遍历 ActiveKeyCount
	// 来获取所有已知的 provider
	for _, prov := range hc.getAllProviders() {
		hc.checkProvider(prov)
	}
}

func (hc *HealthChecker) checkProvider(prov models.ProviderKind) {
	ctx, cancel := context.WithTimeout(context.Background(), hc.timeout)
	defer cancel()

	err := hc.lb.HealthCheck(ctx, prov)

	hc.mu.Lock()
	defer hc.mu.Unlock()

	if err != nil {
		hc.failures[prov]++
		fails := hc.failures[prov]
		slog.Warn("健康检查失败",
			"provider", prov,
			"consecutive_failures", fails,
			"error", err,
		)
		if fails >= hc.threshold && hc.lb.IsProviderEnabled(prov) {
			hc.lb.DisableProvider(prov)
			slog.Warn("健康检查: provider 已自动禁用",
				"provider", prov,
				"failures", fails,
			)
		}
	} else {
		// 恢复成功
		if hc.failures[prov] > 0 {
			slog.Info("健康检查恢复", "provider", prov)
		}
		hc.failures[prov] = 0
		if !hc.lb.IsProviderEnabled(prov) {
			hc.lb.EnableProvider(prov)
			slog.Info("健康检查: provider 已自动启用", "provider", prov)
		}
	}
}

// getAllProviders 返回当前所有已注册的 provider 名称（包括已禁用的）。
func (hc *HealthChecker) getAllProviders() []models.ProviderKind {
	return hc.lb.Providers()
}