package db import ( "context" "fmt" "log" "time" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" "valera/internal/config" ) type DB struct { chatsColection *mongo.Collection workoutsColection *mongo.Collection caloriesColection *mongo.Collection weightColection *mongo.Collection } func NewDB( cfg *config.AppConfig, ) (*DB, error) { ctx := context.Background() opt := options.Client(). ApplyURI(cfg.MongoConfig.URL). SetMinPoolSize(10) client, err := mongo.Connect(ctx, opt) if err != nil { return nil, err } return &DB{ chatsColection: client.Database(cfg.MongoConfig.DBName).Collection(cfg.MongoConfig.ChatsCollectionName), workoutsColection: client.Database(cfg.MongoConfig.DBName).Collection(cfg.MongoConfig.WorkoutsCollectionName), caloriesColection: client.Database(cfg.MongoConfig.DBName).Collection(cfg.MongoConfig.CaloriesCollectionName), weightColection: client.Database(cfg.MongoConfig.DBName).Collection(cfg.MongoConfig.WeightCollectionName), }, nil } func (db *DB) AddChat(chatID int64) error { ctx := context.Background() result := db.chatsColection.FindOne(ctx, bson.M{"chat_id": chatID}) if result.Err() == mongo.ErrNoDocuments { if _, err := db.chatsColection.InsertOne(ctx, &Chat{ChatID: chatID}); err != nil { return err } return nil } if result.Err() != nil { return result.Err() } loc, _ := time.LoadLocation("Asia/Novosibirsk") t := time.Now().In(loc) _, err := db.chatsColection.UpdateOne( ctx, bson.M{"chat_id": chatID}, bson.M{"$set": bson.M{"last_time": t}}, ) return err } func (db *DB) AddWorkout(chatID int64, workout *Workout) error { ctx := context.Background() if workout.Count <= 0 { return nil } if err := db.AddChat(chatID); err != nil { return err } workout.ChatID = chatID _, err := db.workoutsColection.InsertOne( ctx, workout, ) return err } func (db *DB) AddCalories(chatID int64, calories *Calories) error { ctx := context.Background() if err := db.AddChat(chatID); err != nil { return err } calories.ChatID = chatID _, err := db.caloriesColection.InsertOne( ctx, calories, ) return err } func (db *DB) AddWeight(chatID int64, weight *Weight) error { ctx := context.Background() if weight.Weight <= 0 { return nil } if err := db.AddChat(chatID); err != nil { return err } weight.ChatID = chatID _, err := db.weightColection.InsertOne( ctx, weight, ) return err } func (db *DB) GetChatInfo(chatID int64, username string) (*ChatInfo, error) { ctx := context.Background() if err := db.AddChat(chatID); err != nil { return nil, err } chatInfoDTO := &ChatInfo{Username: username} if err := db.chatsColection.FindOne(ctx, bson.M{"chat_id": chatID}).Decode(chatInfoDTO); err != nil { return nil, err } return chatInfoDTO, nil } func (db *DB) SetStatusToChat(chatID int64, status UserState) error { ctx := context.Background() _, err := db.chatsColection.UpdateOne( ctx, bson.M{"chat_id": chatID}, bson.M{"$set": bson.M{"status": status}}, ) return err } func (db *DB) SetPause(chatID int64, duration time.Duration) error { ctx := context.Background() loc, _ := time.LoadLocation("Asia/Novosibirsk") t := time.Now().In(loc).Add(+duration) _, err := db.chatsColection.UpdateOne( ctx, bson.M{"chat_id": chatID}, bson.M{"$set": bson.M{"pause_until": t}}, ) return err } func (db *DB) GetStatBetween(chatID int64, startTime time.Time, endTime time.Time) (map[string]float64, error) { ctx := context.Background() if err := db.AddChat(chatID); err != nil { return nil, err } res := map[string]float64{} cursor, err := db.workoutsColection.Find( ctx, bson.M{ "chat_id": chatID, "created_at": bson.M{ "$gte": startTime, "$lt": endTime, }, }, options.Find().SetSort( bson.M{ "created_at": 1, }, ), ) if err != nil { return nil, err } for cursor.Next(context.Background()) { var result Workout if err := cursor.Decode(&result); err != nil { log.Fatal(err) } fmt.Println(result) res[result.Name] += float64(result.Count) } if err := cursor.Err(); err != nil { log.Fatal(err) } caloriesCursor, err := db.caloriesColection.Find( ctx, bson.M{ "chat_id": chatID, "created_at": bson.M{ "$gte": startTime, "$lt": endTime, }, }, options.Find().SetSort( bson.M{ "created_at": 1, }, ), ) if err != nil { return nil, err } for caloriesCursor.Next(context.Background()) { var result Calories if err := caloriesCursor.Decode(&result); err != nil { log.Fatal(err) } res["Калории"] += float64(result.Count) } if err := caloriesCursor.Err(); err != nil { log.Fatal(err) } weightCursor, err := db.weightColection.Find( ctx, bson.M{ "chat_id": chatID, "created_at": bson.M{ "$gte": startTime, "$lt": endTime, }, }, options.Find().SetSort( bson.M{ "created_at": 1, }, ), ) if err != nil { return nil, err } for weightCursor.Next(context.Background()) { var result Weight if err := weightCursor.Decode(&result); err != nil { log.Fatal(err) } res["Вес кг"] = result.Weight } if err := weightCursor.Err(); err != nil { log.Fatal(err) } return res, err } func (db *DB) GetAllChats() ([]int64, error) { ctx := context.Background() loc, _ := time.LoadLocation("Asia/Novosibirsk") t := time.Now().In(loc) cursor, err := db.chatsColection.Find( ctx, bson.M{ "$or": bson.A{ bson.M{ "pause_until": bson.M{"$lt": t}, }, bson.M{ "pause_until": bson.M{"$exists": false}, }, }, }, ) if err != nil { return nil, err } var res []int64 for cursor.Next(context.Background()) { var result ChatInfo if err := cursor.Decode(&result); err != nil { log.Fatal(err) } res = append(res, result.ChatID) } if err := cursor.Err(); err != nil { log.Fatal(err) } return res, err }