package stats import ( "context" "errors" "fmt" "strconv" "sync" "go.mongodb.org/mongo-driver/v2/bson" "go.mongodb.org/mongo-driver/v2/mongo" "go.mongodb.org/mongo-driver/v2/mongo/options" "github.com/tiennm99/miti99bot/internal/storage" ) func bsonField(key string, value any) bson.E { return bson.E{Key: key, Value: value} } // usageEntry is the only stats payload shape. One document represents either // a command-level anonymous bucket or a concrete (command, Telegram user) pair. type usageEntry struct { Cmd string `json:"cmd" bson:"cmd"` UserID int64 `json:"uid,omitempty" bson:"uid,omitempty"` Username string `json:"user,omitempty" bson:"user,omitempty"` N int64 `json:"n" bson:"n"` Deleted bool `json:"deleted,omitempty" bson:"deleted,omitempty"` } type usageUser struct { ID int64 Username string } type usageStore interface { Increment(ctx context.Context, cmd string, user usageUser, hasUser bool) error TopCommands(ctx context.Context, limit int) ([]row, error) TopUsers(ctx context.Context, limit int) ([]row, error) CommandsByUser(ctx context.Context, username string, limit int) ([]row, bool, error) UsersByCommand(ctx context.Context, cmd string, limit int) ([]row, error) } func newUsageStore(coll storage.Collection) usageStore { if mongoColl, ok := storage.MongoCollection(coll); ok { return &mongoUsageStore{coll: mongoColl} } return &docUsageStore{docs: storage.Typed[usageEntry](coll)} } func usageKey(cmd string, userID int64) string { if userID == 0 { return cmd } return cmd + ":" + strconv.FormatInt(userID, 10) } type docUsageStore struct { mu sync.Mutex docs storage.DocStore[usageEntry] } func (s *docUsageStore) Increment(ctx context.Context, cmd string, user usageUser, hasUser bool) error { s.mu.Lock() defer s.mu.Unlock() userID := int64(0) if hasUser { userID = user.ID } key := usageKey(cmd, userID) entry, _, err := s.docs.Get(ctx, key) switch { case errors.Is(err, storage.ErrNotFound): entry = usageEntry{} case err != nil: return fmt.Errorf("stats get %s: %w", key, err) } entry.Cmd = cmd entry.N++ entry.Deleted = false if hasUser { entry.UserID = user.ID entry.Username = user.Username } else { entry.UserID = 0 entry.Username = "" } if err := s.docs.Put(ctx, key, entry); err != nil { return fmt.Errorf("stats put %s: %w", key, err) } if hasUser { return s.refreshUsernameLocked(ctx, user) } return nil } func (s *docUsageStore) refreshUsernameLocked(ctx context.Context, user usageUser) error { entries, err := s.loadEntriesLocked(ctx) if err != nil { return err } for _, e := range entries { if e.UserID != user.ID || e.Username == user.Username { continue } e.Username = user.Username if err := s.docs.Put(ctx, usageKey(e.Cmd, e.UserID), e); err != nil { return fmt.Errorf("stats refresh username %d: %w", user.ID, err) } } return nil } func (s *docUsageStore) TopCommands(ctx context.Context, limit int) ([]row, error) { s.mu.Lock() defer s.mu.Unlock() entries, err := s.loadEntriesLocked(ctx) if err != nil { return nil, err } totals := make(map[string]int64) for _, e := range entries { if e.Deleted { continue } totals[e.Cmd] += e.N } rows := make([]row, 0, len(totals)) for cmd, n := range totals { rows = append(rows, row{display: "/" + cmd, n: n}) } sortRows(rows) return limitRows(rows, limit), nil } func (s *docUsageStore) TopUsers(ctx context.Context, limit int) ([]row, error) { s.mu.Lock() defer s.mu.Unlock() entries, err := s.loadEntriesLocked(ctx) if err != nil { return nil, err } type total struct { username string n int64 } totals := make(map[int64]total) for _, e := range entries { if e.Deleted || e.UserID == 0 || e.Username == "" { continue } t := totals[e.UserID] t.username = e.Username t.n += e.N totals[e.UserID] = t } rows := make([]row, 0, len(totals)) for _, t := range totals { rows = append(rows, row{display: "@" + t.username, n: t.n}) } sortRows(rows) return limitRows(rows, limit), nil } func (s *docUsageStore) CommandsByUser(ctx context.Context, username string, limit int) ([]row, bool, error) { s.mu.Lock() defer s.mu.Unlock() entries, err := s.loadEntriesLocked(ctx) if err != nil { return nil, false, err } var ( userID int64 found bool ) for _, e := range entries { if !e.Deleted && e.UserID != 0 && e.Username == username { userID = e.UserID found = true break } } if !found { return nil, false, nil } rows := make([]row, 0) for _, e := range entries { if !e.Deleted && e.UserID == userID { rows = append(rows, row{display: "/" + e.Cmd, n: e.N}) } } sortRows(rows) return limitRows(rows, limit), true, nil } func (s *docUsageStore) UsersByCommand(ctx context.Context, cmd string, limit int) ([]row, error) { s.mu.Lock() defer s.mu.Unlock() entries, err := s.loadEntriesLocked(ctx) if err != nil { return nil, err } rows := make([]row, 0) for _, e := range entries { if !e.Deleted && e.Cmd == cmd && e.UserID != 0 && e.Username != "" { rows = append(rows, row{display: "@" + e.Username, n: e.N}) } } sortRows(rows) return limitRows(rows, limit), nil } func (s *docUsageStore) loadEntriesLocked(ctx context.Context) ([]usageEntry, error) { keys, err := s.docs.List(ctx, "") if err != nil { return nil, fmt.Errorf("stats list: %w", err) } entries := make([]usageEntry, 0, len(keys)) for _, key := range keys { entry, _, err := s.docs.Get(ctx, key) if err != nil { return nil, fmt.Errorf("stats get %s: %w", key, err) } if entry.Cmd == "" { continue } entries = append(entries, entry) } return entries, nil } type mongoUsageStore struct { coll *mongo.Collection } func (s *mongoUsageStore) Increment(ctx context.Context, cmd string, user usageUser, hasUser bool) error { userID := int64(0) if hasUser { userID = user.ID } set := bson.M{"cmd": cmd} unset := bson.M{"deleted": ""} update := bson.M{ "$set": set, "$inc": bson.M{"n": int64(1), "version": int64(1)}, "$unset": unset, "$currentDate": bson.M{"updatedAt": true}, } if hasUser { set["uid"] = user.ID set["user"] = user.Username } else { unset["uid"] = "" unset["user"] = "" } key := usageKey(cmd, userID) if _, err := s.coll.UpdateOne(ctx, bson.M{"_id": key}, update, options.UpdateOne().SetUpsert(true), ); err != nil { return fmt.Errorf("mongo stats increment %s: %w", key, err) } if !hasUser { return nil } if _, err := s.coll.UpdateMany(ctx, bson.M{"uid": user.ID, "user": bson.M{"$ne": user.Username}}, bson.M{ "$set": bson.M{"user": user.Username}, "$inc": bson.M{"version": int64(1)}, "$currentDate": bson.M{"updatedAt": true}, }, ); err != nil { return fmt.Errorf("mongo stats refresh username %d: %w", user.ID, err) } return nil } func (s *mongoUsageStore) TopCommands(ctx context.Context, limit int) ([]row, error) { pipeline := mongo.Pipeline{ bson.D{bsonField("$match", bson.D{ bsonField("cmd", bson.D{bsonField("$type", "string")}), bsonField("deleted", bson.D{bsonField("$ne", true)}), })}, bson.D{bsonField("$group", bson.D{bsonField("_id", "$cmd"), bsonField("n", bson.D{bsonField("$sum", "$n")})})}, bson.D{bsonField("$sort", bson.D{bsonField("n", -1), bsonField("_id", 1)})}, } pipeline = withLimit(pipeline, limit) cur, err := s.coll.Aggregate(ctx, pipeline) if err != nil { return nil, fmt.Errorf("mongo stats top commands: %w", err) } defer func() { _ = cur.Close(ctx) }() var rows []row for cur.Next(ctx) { var doc struct { Cmd string `bson:"_id"` N int64 `bson:"n"` } if err := cur.Decode(&doc); err != nil { return nil, fmt.Errorf("mongo stats top commands decode: %w", err) } rows = append(rows, row{display: "/" + doc.Cmd, n: doc.N}) } if err := cur.Err(); err != nil { return nil, fmt.Errorf("mongo stats top commands cursor: %w", err) } return rows, nil } func (s *mongoUsageStore) TopUsers(ctx context.Context, limit int) ([]row, error) { pipeline := mongo.Pipeline{ bson.D{bsonField("$match", bson.D{ bsonField("uid", bson.D{bsonField("$gt", int64(0))}), bsonField("user", bson.D{bsonField("$type", "string"), bsonField("$ne", "")}), bsonField("deleted", bson.D{bsonField("$ne", true)}), })}, bson.D{bsonField("$sort", bson.D{bsonField("updatedAt", -1)})}, bson.D{bsonField("$group", bson.D{ bsonField("_id", "$uid"), bsonField("user", bson.D{bsonField("$first", "$user")}), bsonField("n", bson.D{bsonField("$sum", "$n")}), })}, bson.D{bsonField("$sort", bson.D{bsonField("n", -1), bsonField("user", 1)})}, } pipeline = withLimit(pipeline, limit) cur, err := s.coll.Aggregate(ctx, pipeline) if err != nil { return nil, fmt.Errorf("mongo stats top users: %w", err) } defer func() { _ = cur.Close(ctx) }() var rows []row for cur.Next(ctx) { var doc struct { User string `bson:"user"` N int64 `bson:"n"` } if err := cur.Decode(&doc); err != nil { return nil, fmt.Errorf("mongo stats top users decode: %w", err) } rows = append(rows, row{display: "@" + doc.User, n: doc.N}) } if err := cur.Err(); err != nil { return nil, fmt.Errorf("mongo stats top users cursor: %w", err) } return rows, nil } func (s *mongoUsageStore) CommandsByUser(ctx context.Context, username string, limit int) ([]row, bool, error) { userID, found, err := s.userIDByUsername(ctx, username) if err != nil || !found { return nil, found, err } opts := options.Find(). SetProjection(bson.M{"cmd": 1, "n": 1}). SetSort(bson.D{bsonField("n", -1), bsonField("cmd", 1)}) if limit > 0 { opts.SetLimit(int64(limit)) } cur, err := s.coll.Find(ctx, bson.M{ "uid": userID, "cmd": bson.M{"$type": "string"}, "deleted": bson.M{"$ne": true}, }, opts) if err != nil { return nil, true, fmt.Errorf("mongo stats commands by user %s: %w", username, err) } defer func() { _ = cur.Close(ctx) }() var rows []row for cur.Next(ctx) { var doc usageEntry if err := cur.Decode(&doc); err != nil { return nil, true, fmt.Errorf("mongo stats commands by user decode: %w", err) } rows = append(rows, row{display: "/" + doc.Cmd, n: doc.N}) } if err := cur.Err(); err != nil { return nil, true, fmt.Errorf("mongo stats commands by user cursor: %w", err) } return rows, true, nil } func (s *mongoUsageStore) UsersByCommand(ctx context.Context, cmd string, limit int) ([]row, error) { opts := options.Find(). SetProjection(bson.M{"user": 1, "n": 1}). SetSort(bson.D{bsonField("n", -1), bsonField("user", 1)}) if limit > 0 { opts.SetLimit(int64(limit)) } cur, err := s.coll.Find(ctx, bson.M{ "cmd": cmd, "uid": bson.M{"$gt": int64(0)}, "user": bson.M{"$type": "string", "$ne": ""}, "deleted": bson.M{"$ne": true}, }, opts) if err != nil { return nil, fmt.Errorf("mongo stats users by command %s: %w", cmd, err) } defer func() { _ = cur.Close(ctx) }() var rows []row for cur.Next(ctx) { var doc usageEntry if err := cur.Decode(&doc); err != nil { return nil, fmt.Errorf("mongo stats users by command decode: %w", err) } rows = append(rows, row{display: "@" + doc.Username, n: doc.N}) } if err := cur.Err(); err != nil { return nil, fmt.Errorf("mongo stats users by command cursor: %w", err) } return rows, nil } func (s *mongoUsageStore) userIDByUsername(ctx context.Context, username string) (int64, bool, error) { var doc struct { UserID int64 `bson:"uid"` } err := s.coll.FindOne(ctx, bson.M{ "uid": bson.M{"$gt": int64(0)}, "user": username, "deleted": bson.M{"$ne": true}, }, options.FindOne(). SetProjection(bson.M{"uid": 1}), ).Decode(&doc) if err != nil { if errors.Is(err, mongo.ErrNoDocuments) { return 0, false, nil } return 0, false, fmt.Errorf("mongo stats find user %s: %w", username, err) } return doc.UserID, true, nil } func withLimit(pipeline mongo.Pipeline, limit int) mongo.Pipeline { if limit <= 0 { return pipeline } return append(pipeline, bson.D{bsonField("$limit", int64(limit))}) } func limitRows(rows []row, limit int) []row { if limit <= 0 || len(rows) <= limit { return rows } return rows[:limit] }