package kvstore

import (
	"context"
	"fmt"
	"strings"

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

// MongoDBKVStore is an implementation of the KVStore interface that uses a MongoDB database to store key/value pairs.
type MongoDBKVStore struct {
	client   *mongo.Client
	database string
	collection string
}

// NewMongoDBKVStore creates a new MongoDBKVStore.
func NewMongoDBKVStore(client *mongo.Client, database string, collection string) *MongoDBKVStore {
	return &MongoDBKVStore{
		client:   client,
		database: database,
		collection: collection,
	}
}

// Set sets the value for a given key.
func (store *MongoDBKVStore) Set(key string, value string) error {
	_, err := store.client.
		Database(store.database).
		Collection(store.collection).
		ReplaceOne(context.TODO(), bson.M{"_id": key}, bson.M{"_id": key, "value": value}, options.Replace().SetUpsert(true))
	return err
}

// Get retrieves the value for a given key.
func (store *MongoDBKVStore) Get(key string) (string, error) {
	var result bson.M
	err := store.client.
		Database(store.database).
		Collection(store.collection).
		FindOne(context.TODO(), bson.M{"_id": key}).
		Decode(&result)
	if err != nil {
		return "", fmt.Errorf("key not found: %s", key)
	}
	return result["value"].(string), nil
}

// Delete deletes the value for a given key.
func (store *MongoDBKVStore) Delete(key string) error {
	_, err := store.client.
		Database(store.database).
		Collection(store.collection).
		DeleteOne(context.TODO(), bson.M{"_id": key})
	return err
}

// Query returns all values that match a given query string.
func (store *MongoDBKVStore) Query(query string) ([]string, error) {
	var results []string
	cur, err := store.client.
		Database(store.database).
		Collection(store.collection).
		Find(context.TODO(), bson.M{"$or": []bson.M{{"_id": bson.M{"$regex": query}}, {"value": bson.M{"$regex": query}}}})
	if err != nil {
		return nil, err
	}
	defer cur.Close(context.TODO())