seek-tune/utils/dbClient.go

package utils

import (
	"context"
	"fmt"

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

const dbUri string = "mongodb://localhost:27017"

// DbClient represents a MongoDB client
type DbClient struct {
	client *mongo.Client
}

// NewDbClient creates a new instance of DbClient
func NewDbClient() (*DbClient, error) {
	clientOptions := options.Client().ApplyURI(dbUri)
	client, err := mongo.Connect(context.Background(), clientOptions)
	if err != nil {
		return nil, fmt.Errorf("error connecting to MongoDB: %d", err)
	}
	return &DbClient{client: client}, nil
}

// Close closes the underlying MongoDB client
func (db *DbClient) Close() error {
	if db.client != nil {
		return db.client.Disconnect(context.Background())
	}
	return nil
}

func (db *DbClient) TotalSongs() (int, error) {
	existingSongsCollection := db.client.Database("song-recognition").Collection("existing-songs")
	total, err := existingSongsCollection.CountDocuments(context.Background(), bson.D{})
	if err != nil {
		return 0, err
	}

	return int(total), nil
}

func (db *DbClient) SongExists(songTitle, songArtist, ytID string) (bool, error) {
	existingSongsCollection := db.client.Database("song-recognition").Collection("existing-songs")

	key := fmt.Sprintf("%s - %s", songTitle, songArtist)
	var filter bson.M

	if len(ytID) == 0 {
		filter = bson.M{"_id": key}
	} else {
		filter = bson.M{"ytID": ytID}
	}

	var result bson.M
	if err := existingSongsCollection.FindOne(context.Background(), filter).Decode(&result); err == nil {
		return true, nil
	} else if err != mongo.ErrNoDocuments {
		return false, fmt.Errorf("failed to retrieve registered songs: %v", err)
	}

	return false, nil
}

func (db *DbClient) RegisterSong(songTitle, songArtist, ytID string) error {
	existingSongsCollection := db.client.Database("song-recognition").Collection("existing-songs")

	// Create a compound unique index on ytID and key, if it doesn't already exist
	indexModel := mongo.IndexModel{
		Keys:    bson.D{{"ytID", 1}, {"key", 1}},
		Options: options.Index().SetUnique(true),
	}
	_, err := existingSongsCollection.Indexes().CreateOne(context.Background(), indexModel)
	if err != nil {
		return fmt.Errorf("failed to create unique index: %v", err)
	}

	// Attempt to insert the song with ytID and key
	key := fmt.Sprintf("%s - %s", songTitle, songArtist)
	_, err = existingSongsCollection.InsertOne(context.Background(), bson.M{"_id": key, "ytID": ytID})
	if err != nil {
		if mongo.IsDuplicateKeyError(err) {
			return fmt.Errorf("song with ytID or key already exists: %v", err)
		} else {
			return fmt.Errorf("failed to register song: %v", err)
		}
	}

	return nil
}

func (db *DbClient) InsertChunkTag(chunkfgp int64, chunkTag interface{}) error {
	chunksCollection := db.client.Database("song-recognition").Collection("chunks")

	filter := bson.M{"fingerprint": chunkfgp}

	var result bson.M
	err := chunksCollection.FindOne(context.Background(), filter).Decode(&result)
	if err == nil {
		// If the fingerprint already exists, append the chunkTag to the existing list
		// fmt.Println("DUPLICATE FINGERPRINT: ", chunkfgp)
		update := bson.M{"$push": bson.M{"chunkTags": chunkTag}}
		_, err := chunksCollection.UpdateOne(context.Background(), filter, update)
		if err != nil {
			return fmt.Errorf("failed to update chunkTags: %v", err)
		}
		return nil
	} else if err != mongo.ErrNoDocuments {
		return err
	}

	// If the document doesn't exist, insert a new document
	_, err = chunksCollection.InsertOne(context.Background(), bson.M{"fingerprint": chunkfgp, "chunkTags": []interface{}{chunkTag}})
	if err != nil {
		return fmt.Errorf("failed to insert chunk tag: %v", err)
	}

	return nil
}

func (db *DbClient) GetChunkTags(chunkfgp int64) ([]primitive.M, error) {
	chunksCollection := db.client.Database("song-recognition").Collection("chunks")

	filter := bson.M{"fingerprint": chunkfgp}
	result := bson.M{}
	err := chunksCollection.FindOne(context.Background(), filter).Decode(&result)

	if err != nil {
		if err == mongo.ErrNoDocuments {
			return nil, nil
		}
		return nil, fmt.Errorf("failed to retrieve chunk tag: %w", err)
	}

	var listOfChunkTags []primitive.M
	for _, data := range result["chunkTags"].(primitive.A) {
		listOfChunkTags = append(listOfChunkTags, data.(primitive.M))
	}

	return listOfChunkTags, nil
}

func (db *DbClient) GetChunkTagForSong(songTitle, songArtist string) (bson.M, error) {
	chunksCollection := db.client.Database("song-recognition").Collection("chunks")

	filter := bson.M{
		"chunkTags": bson.M{
			"$elemMatch": bson.M{
				"songtitle":  songTitle,
				"songartist": songArtist,
			},
		},
	}

	var result bson.M
	if err := chunksCollection.FindOne(context.Background(), filter).Decode(&result); err != nil {
		if err == mongo.ErrNoDocuments {
			return nil, nil
		}
		return nil, fmt.Errorf("failed to find chunk: %v", err)
	}

	var chunkTag map[string]interface{}
	for _, chunk := range result["chunkTags"].(primitive.A) {
		chunkMap, ok := chunk.(primitive.M)
		if !ok {
			continue
		}
		if chunkMap["songtitle"] == songTitle && chunkMap["songartist"] == songArtist {
			chunkTag = chunkMap
			break
		}
	}

	return chunkTag, nil
}