package mongo

import (
	"common/config"
	"common/log"
	"context"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type DBName string

var clients = make(map[DBName]*mongo.Client)

func Init(cfg map[string]*config.MongoConfig) error {
	for name, oneConfig := range cfg {
		if client, err := initOneClient(oneConfig); err != nil {
			return err
		} else {
			clients[DBName(name)] = client
		}
	}
	return nil
}

func initOneClient(cfg *config.MongoConfig) (*mongo.Client, error) {
	opts := options.Client().
		ApplyURI(cfg.URI)

	client, err := mongo.Connect(context.Background(), opts)
	if err != nil {
		return nil, err
	}

	return client, nil
}

// GetClient 返回 mongo.Client，你可以通过 .Database("xxx") 获取具体数据库
func GetClient(dbName DBName) *mongo.Client {
	if c, ok := clients[dbName]; ok {
		return c
	}
	log.Errorf("mongo client %s not found", dbName)
	return nil
}

// GetDB 是便捷方法，直接返回 *mongo.Database（假设 dbName 对应数据库名）
// 如果你的配置中 dbName 和实际数据库名一致，可这样用
func GetDB(dbName DBName) *mongo.Database {
	client := GetClient(dbName)
	if client == nil {
		return nil
	}
	// 假设配置中的 key 就是数据库名；若需分离，可在 config 中加字段
	return client.Database(string(dbName))
}

func Close() error {
	if clients == nil {
		return nil
	}
	for name, client := range clients {
		if client != nil {
			if err := client.Disconnect(context.Background()); err != nil {
				log.Errorf("close mongo client %s error: %v", name, err)
			}
		}
	}
	return nil
}