package storage

import (
	"context"
	"fmt"
	"local/cheqbooq/config"
	"strings"
	"time"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type Mongo struct {
	client *mongo.Client
}

func NewMongo() (*Mongo, error) {
	ctx, can := context.WithTimeout(context.Background(), time.Second*5)
	defer can()

	opt := options.Client()
	opt.ApplyURI(fmt.Sprintf("mongodb://%s", strings.TrimPrefix(config.StoreAddr, "mongodb://")))

	client, err := mongo.Connect(ctx, opt)
	if err != nil {
		return nil, err
	}

	return &Mongo{
		client: client,
	}, client.Ping(ctx, nil)
}

func (m *Mongo) Close() error {
	return m.client.Disconnect(context.TODO())
}

func (m *Mongo) Account() *mongo.Collection {
	return m.Database(config.StoreNS).Collection("account")
}

func (m *Mongo) Balance() *mongo.Collection {
	return m.Database(config.StoreNS).Collection("balance")
}

func (m *Mongo) Transaction() *mongo.Collection {
	return m.Database(config.StoreNS).Collection("transaction")
}

func (m *Mongo) Find(c *mongo.Collection, where interface{}, next func() interface{}) error {
	ctx, can := context.WithCancel(context.TODO())
	defer can()

	cur, err := c.Find(ctx, where, options.Find())
	if err != nil {
		return err
	}
	defer cur.Close(ctx)

	for cur.Next(ctx) {
		ptr := next()
		if err := cur.Decode(ptr); err != nil {
			return err
		}
	}

	return cur.Err()
}

func (m *Mongo) Upsert(c *mongo.Collection, where, op interface{}) error {
	ctx, can := context.WithCancel(context.TODO())
	defer can()

	_, err := c.UpdateMany(ctx, where, op, options.Update().SetUpsert(true))
	return err
}