package db

import (
	"context"
	"fmt"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"leaf/log"
	"sync"
	"time"
)

type MongoConfig struct {
	Host     string
	Port     int
	Username string
	Password string
}

type MongoPool struct {
	pool       []*mongo.Client
	maxSize    int
	mu         sync.Mutex
	timeout    time.Duration
	config     MongoConfig
	index      int          // 用于异步模式下客户端的轮询
	healthTick *time.Ticker // 定时器
}

// 创建新的 MongoDB 客户端连接
func NewMongoClient(conf MongoConfig, timeout time.Duration) *mongo.Client {
	// 创建一个带有超时的上下文
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel() // 确保上下文在操作完成后取消

	// 配置MongoDB客户端选项，包括连接URI、超时时间
	uri := fmt.Sprintf("mongodb://%s:%d/", conf.Host, conf.Port)
	opt := options.Client().ApplyURI(uri).SetTimeout(timeout)

	// 设置认证信息（如果需要）
	if conf.Username != "" && conf.Password != "" {
		opt.SetAuth(options.Credential{
			Username: conf.Username,
			Password: conf.Password,
		})
	}

	// 使用配置选项连接到MongoDB
	client, err := mongo.Connect(ctx, opt)
	if err != nil {
		log.Fatal("MongoDB connection error: %v", err)
	}
	log.Debug("MongoDB connection successful")
	return client
}

// MongoPool 初始化并管理连接池
func (pool *MongoPool) Init(conf MongoConfig, size int, timeout time.Duration) {
	pool.maxSize = size
	pool.timeout = timeout
	pool.config = conf // 从传入的配置中设置MongoConfig

	// 初始化Mongo连接池
	pool.pool = make([]*mongo.Client, 0, pool.maxSize)
	for i := 0; i < pool.maxSize; i++ {
		client := NewMongoClient(pool.config, pool.timeout) // 从pool.config中获取配置
		pool.pool = append(pool.pool, client)
		log.Debug("Created MongoDB pool client %d", i)
	}

	// 启动健康检查定时任务
	pool.startHealthCheck()
}

// 获取 MongoDB 客户端实例，根据是否同步获取选择对应客户端
func (pool *MongoPool) GetClient(sync bool) (*mongo.Client, error) {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	var client *mongo.Client

	if sync {
		// 同步模式下，直接返回第一个客户端
		client = pool.pool[0]
	} else {
		// 异步模式下，使用轮询算法返回客户端
		client = pool.pool[pool.index]
		// 循环使用客户端
		pool.index++
		if pool.index >= pool.maxSize {
			pool.index = 0
		}
	}

	log.Debug("Allocated MongoDB client %d", pool.index)
	return client, nil
}

// 释放一个 MongoDB 客户端实例，将其返回池中
func (pool *MongoPool) ReleaseClient(client *mongo.Client) {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	// 这里可以检查客户端的健康状态，或者实现客户端重用的逻辑
	pool.pool = append(pool.pool, client)
	log.Debug("Client returned to MongoDB pool")
}

func (pool *MongoPool) Disconnect() {
	pool.stopHealthCheck()
	for _, client := range pool.pool {
		client.Disconnect(context.Background())
	}
	log.Debug("MongoDB pool disconnected")
}

// 健康检查，定期检查连接池中的连接是否有效
func (pool *MongoPool) HealthCheck() {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	for i, client := range pool.pool {
		ctx, cancel := context.WithTimeout(context.Background(), pool.timeout)
		defer cancel()
		err := client.Ping(ctx, nil)
		if err != nil {
			log.Debug("MongoDB client %d is unhealthy: %v", i, err)
			pool.pool = append(pool.pool[:i], pool.pool[i+1:]...) // 移除无效客户端
		}
	}
}

// 启动定时健康检查
func (pool *MongoPool) startHealthCheck() {
	// 设置定时检查的间隔时间，例如每隔 10 秒钟检查一次
	pool.healthTick = time.NewTicker(10 * time.Second)

	go func() {
		for {
			select {
			case <-pool.healthTick.C:
				log.Debug("Performing health check on MongoDB pool")
				pool.HealthCheck()
			}
		}
	}()
}

// 停止定时健康检查
func (pool *MongoPool) stopHealthCheck() {
	if pool.healthTick != nil {
		pool.healthTick.Stop()
	}
}

// 自动重试机制
func (pool *MongoPool) retryOperation(fn func() error, retries int) error {
	for i := 0; i < retries; i++ {
		if err := fn(); err != nil {
			log.Debug("Retrying MongoDB operation... attempt %d/%d", i+1, retries)
			time.Sleep(time.Second * 2) // 可根据需要调整重试间隔
		} else {
			return nil
		}
	}
	return fmt.Errorf("operation failed after %d retries", retries)
}
