LazyPM/internal/storage/mongo.go

package storage

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"

	"github.com/LazyBachelor/LazyPM/internal/models"
	"github.com/charmbracelet/huh"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type MongoStorage struct {
	client *mongo.Client
}

func NewMongoStorage(ctx context.Context, uri, username, password string) (*MongoStorage, error) {
	credentials := options.Credential{
		Username: username,
		Password: password,
	}

	client, err := mongo.Connect(ctx,
		options.Client().ApplyURI(uri).SetAuth(credentials))

	if err != nil {
		return nil, fmt.Errorf("failed to connect to MongoDB: %v", err)
	}

	if err := client.Ping(ctx, nil); err != nil {
		return nil, fmt.Errorf("cannot reach MongoDB: %v", err)
	}

	return &MongoStorage{client: client}, nil
}

func NewMongoStorageInteractive(ctx context.Context, uri string) (*MongoStorage, error) {
	var username, password string
	if os.Getenv("DB_USER") == "" {
		if err := huh.NewInput().
			Title("Enter the Database Username").
			Value(&username).
			WithTheme(huh.ThemeBase16()).Run(); err != nil {
			return nil, fmt.Errorf("failed to read username: %w", err)
		}
	} else {
		username = os.Getenv("DB_USER")
	}

	if username == "" {
		return nil, fmt.Errorf("No username provided.")
	}

	if os.Getenv("DB_PASSWORD") == "" {
		if err := huh.NewInput().
			Title("Enter the Survey Password").
			EchoMode(huh.EchoModePassword).
			Value(&password).
			WithTheme(huh.ThemeBase16()).Run(); err != nil {
			return nil, fmt.Errorf("failed to read password: %w", err)
		}
	} else {
		password = os.Getenv("DB_PASSWORD")
	}

	if password == "" {
		return nil, fmt.Errorf("No password provided.")

	}

	mongoClient, err := NewMongoStorage(ctx, uri, username, password)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to database: %w", err)
	}
	return mongoClient, nil
}

func (s *MongoStorage) Close() error {
	return s.client.Disconnect(context.Background())
}

func (s *MongoStorage) SubmitSurveyResponsesCmd(ctx context.Context, dir string) error {
	pmDir := dir
	userStatscollection := s.client.Database("Responses").Collection("stats")
	taskMetricsCollection := s.client.Database("Responses").Collection("metrics")

	entries, err := os.ReadDir(pmDir)
	if err != nil {
		return fmt.Errorf("Failed to read .pm directory: %v", err)
	}

	if len(entries) == 0 {
		return fmt.Errorf("No files found in .pm directory")
	}

	statFile := pmDir + "stats.json"
	if _, err := os.Stat(statFile); os.IsNotExist(err) {
		return fmt.Errorf("stats.json not found in .pm directory")
	}

	stats, err := getStats(statFile)
	if err != nil {
		return fmt.Errorf("Failed to read stats.json: %v", err)
	}

	_, err = userStatscollection.UpdateOne(ctx,
		bson.M{"_id": stats.ID},
		bson.M{"$set": stats},
		options.Update().SetUpsert(true))

	if err != nil {
		return fmt.Errorf("Failed to insert stats into database: %v", err)
	}

	metricFiles := []string{}
	for _, entry := range entries {
		if entry.IsDir() {
			continue
		}

		if entry.Name() == "stats.json" {
			continue
		}

		if strings.HasSuffix(entry.Name(), "-stats.json") {
			metricFiles = append(metricFiles, pmDir+entry.Name())
			continue
		}
	}

	for _, file := range metricFiles {
		metrics, err := getTaskMetrics(file)
		if err != nil {
			fmt.Printf("failed to read metrics from %s: %v", file, err)
			continue
		}

		_, err = taskMetricsCollection.UpdateOne(ctx, bson.M{"_id": metrics.ID}, bson.M{"$set": metrics},
			options.Update().SetUpsert(true))

		if err != nil {
			fmt.Printf("failed to insert metrics from %s: %v", file, err)
			continue
		}
	}

	return nil
}

func getStats(file string) (*models.Statistics, error) {
	data, err := os.ReadFile(file)
	if err != nil {
		return nil, err
	}

	var stats models.Statistics
	if err := json.Unmarshal(data, &stats); err != nil {
		return nil, err
	}

	return &stats, nil
}

func getTaskMetrics(file string) (*models.TaskMetricsFile, error) {
	data, err := os.ReadFile(file)
	if err != nil {
		return nil, err
	}

	var metrics models.TaskMetricsFile
	if err := json.Unmarshal(data, &metrics); err != nil {
		return nil, err
	}

	return &metrics, nil
}