package db import ( "context" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" "log" "time" "valera/config" ) const ( UserStateNone = UserState("") UserStateGo = UserState("Go") ) type UserState string type chatDTO struct { ChatID int64 `bson:"chat_id"` } type ChatInfoDTO struct { ChatID int64 `bson:"chat_id"` Status string `bson:"status"` } func (c *ChatInfoDTO) GetStatus() UserState { if c.Status == "" { return UserStateNone } return UserStateGo } type Workout struct { ChatID int64 `bson:"chat_id"` Name string `bson:"name"` Count int `bson:"count"` CreatedAt time.Time `bson:"created_at"` Username string `bson:"username"` } func NewWorkout( name string, count int, username string, ) *Workout { return &Workout{ Name: name, Count: count, Username: username, } } var cfg *config.Config func init() { cfg = config.NewConfig() } func AddChat(chatID int64) error { ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() client, err := mongo.Connect(ctx, options.Client().ApplyURI(cfg.MongoURL)) if err != nil { return err } collection := client.Database(cfg.DBName).Collection(cfg.ChatsCollectionName) result := collection.FindOne(ctx, bson.M{"chat_id": chatID}) if result.Err() == mongo.ErrNoDocuments { if _, err := collection.InsertOne(ctx, &chatDTO{ChatID: chatID}); err != nil { return err } return nil } return result.Err() } func AddWorkout(chatID int64, workout *Workout) error { if workout.Count <= 0 { return nil } if err := AddChat(chatID); err != nil { return err } workout.ChatID = chatID loc, _ := time.LoadLocation("Asia/Novosibirsk") workout.CreatedAt = time.Now().In(loc) ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() client, err := mongo.Connect(ctx, options.Client().ApplyURI(cfg.MongoURL)) if err != nil { return err } collection := client.Database(cfg.DBName).Collection(cfg.WorkoutsCollectionName) _, err = collection.InsertOne( ctx, workout, ) return err } func GetChatInfo(chatID int64) (*ChatInfoDTO, error) { if err := AddChat(chatID); err != nil { return nil, err } ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() client, err := mongo.Connect(ctx, options.Client().ApplyURI(cfg.MongoURL)) if err != nil { return nil, err } collection := client.Database(cfg.DBName).Collection(cfg.ChatsCollectionName) chatInfoDTO := &ChatInfoDTO{} if err := collection.FindOne(ctx, bson.M{"chat_id": chatID}).Decode(chatInfoDTO); err != nil { return nil, err } return chatInfoDTO, err } func SetStatusToChat(chatID int64, status UserState) error { ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() client, err := mongo.Connect(ctx, options.Client().ApplyURI(cfg.MongoURL)) if err != nil { return err } collection := client.Database(cfg.DBName).Collection(cfg.ChatsCollectionName) _, err = collection.UpdateOne( ctx, bson.M{"chat_id": chatID}, bson.M{"$set": bson.M{"status": status}}, ) return err } func GetStat(chatID int64) (map[string]int, error) { if err := AddChat(chatID); err != nil { return nil, err } ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() client, err := mongo.Connect(ctx, options.Client().ApplyURI(cfg.MongoURL)) if err != nil { return nil, err } loc, _ := time.LoadLocation("Asia/Novosibirsk") t := time.Now().In(loc).Add(-24 * time.Hour) collection := client.Database(cfg.DBName).Collection(cfg.WorkoutsCollectionName) cursor, err := collection.Find( ctx, bson.M{ "chat_id": chatID, "created_at": bson.M{"$gt": t}, }, ) if err != nil { return nil, err } res := map[string]int{} for cursor.Next(context.Background()) { var result Workout if err := cursor.Decode(&result); err != nil { log.Fatal(err) } res[result.Name] += result.Count } if err := cursor.Err(); err != nil { log.Fatal(err) } return res, err } func GetAllChats() ([]int64, error) { ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() client, err := mongo.Connect(ctx, options.Client().ApplyURI(cfg.MongoURL)) if err != nil { return nil, err } collection := client.Database(cfg.DBName).Collection(cfg.ChatsCollectionName) cursor, err := collection.Find(ctx,bson.M{}) if err != nil { return nil, err } var res []int64 for cursor.Next(context.Background()) { var result ChatInfoDTO 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 }