package mongo

import (
	"context"
	"fmt"
	"time"

	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/core/stores/mon"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
	"go.mongodb.org/mongo-driver/v2/mongo/readpref"
)

type DocumentDB struct {
	Mon *mon.Model
}

func NewDocumentDB(config *Conf, collection string, opts ...mon.Option) (DocumentDBUseCase, error) {
	if config == nil {
		return nil, fmt.Errorf("mongo: config is nil")
	}
	if collection == "" {
		return nil, fmt.Errorf("mongo: collection is required")
	}

	connectionURI, err := buildConnectionURI(*config)
	if err != nil {
		return nil, err
	}

	opts = append(opts, InitMongoOptions(*config))
	logx.Infof("[DocumentDB] connecting to %s db=%s coll=%s", redactConnectionURI(connectionURI), config.Database, collection)

	client, err := mon.NewModel(connectionURI, config.Database, collection, opts...)
	if err != nil {
		return nil, fmt.Errorf("mongo: new model: %w", err)
	}

	timeout := time.Duration(config.ConnectTimeoutMs) * time.Millisecond
	if timeout <= 0 {
		timeout = 10 * time.Second
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	if err := client.Collection.Database().Client().Ping(ctx, readpref.Primary()); err != nil {
		return nil, fmt.Errorf("mongo: ping primary: %w", err)
	}
	logx.Infof("[DocumentDB] connected")

	return &DocumentDB{Mon: client}, nil
}

func (document *DocumentDB) PopulateIndex(ctx context.Context, key string, sort int32, unique bool) error {
	return document.createIndex(ctx, []string{key}, []int32{sort}, unique, nil)
}

func (document *DocumentDB) PopulateTTLIndex(ctx context.Context, key string, sort int32, unique bool, ttl int32) error {
	return document.createIndex(ctx, []string{key}, []int32{sort}, unique, options.Index().SetExpireAfterSeconds(ttl))
}

func (document *DocumentDB) PopulateMultiIndex(ctx context.Context, keys []string, sorts []int32, unique bool) error {
	if len(keys) != len(sorts) {
		return fmt.Errorf("mongo: keys and sorts length mismatch")
	}
	return document.createIndex(ctx, keys, sorts, unique, nil)
}

func (document *DocumentDB) createIndex(
	ctx context.Context,
	keys []string,
	sorts []int32,
	unique bool,
	indexOpt *options.IndexOptionsBuilder,
) error {
	c := document.Mon.Collection
	index := yieldIndexModel(keys, sorts, unique, indexOpt)
	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()
	_, err := c.Indexes().CreateOne(ctx, index)
	if err != nil {
		return fmt.Errorf("mongo: create index: %w", err)
	}
	return nil
}

func (document *DocumentDB) GetClient() *mon.Model {
	return document.Mon
}

func yieldIndexModel(keys []string, sorts []int32, unique bool, indexOpt *options.IndexOptionsBuilder) mongo.IndexModel {
	setKeysDoc := bson.D{}
	for i, key := range keys {
		setKeysDoc = append(setKeysDoc, bson.E{Key: key, Value: sorts[i]})
	}
	if indexOpt == nil {
		indexOpt = options.Index()
	}
	indexOpt.SetUnique(unique)
	return mongo.IndexModel{
		Keys:    setKeysDoc,
		Options: indexOpt,
	}
}