stats: Only process messages from main chat

.gitignore

@@ -1,4 +1,5 @@
 config.json
+stats.db
 clessy-broker
 clessy-mods
 clessy-stats

stats/main.go

@@ -14,20 +14,29 @@ func assert(err error) {
 }
 
 var db *bolt.DB
+var chatID *int
 
 func process(broker *tg.Broker, update tg.APIMessage) {
-
+	// Process messages from marked chat only
+	if update.Chat.ChatID != *chatID {
+		return
+	}
+	getNick(update.User)
+	updateStats(update)
 }
 
 func main() {
 	brokerAddr := flag.String("broker", "localhost:7314", "Broker address:port")
 	boltdbFile := flag.String("boltdb", "stats.db", "BoltDB database file")
+	chatID = flag.Int("chatid", -14625256, "Telegram Chat ID to count stats for")
 	flag.Parse()
 
-	db, err := bolt.Open(*boltdbFile, 0600, nil)
+	var err error
+	db, err = bolt.Open(*boltdbFile, 0600, nil)
 	assert(err)
 	defer db.Close()
 
+	loadUsers()
 	loadStats()
 
 	err = tg.CreateBrokerClient(*brokerAddr, process)

stats/stats.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/boltdb/bolt"
+	"github.com/hamcha/clessy/tg"
 )
 
 const (
@@ -18,18 +19,21 @@ const (
 	MessageTypeVoice    int = 5
 	MessageTypeContact  int = 6
 	MessageTypeLocation int = 7
-	MessageTypeMax      int = 8
+	MessageTypeDocument int = 8
+	MessageTypeMax      int = 9
 )
 
 type Stats struct {
 	ByUserCount    map[string]uint64
 	ByUserAvgLen   map[string]uint64
+	ByUserAvgCount map[string]uint64
 	ByWeekday      [7]uint64
 	ByHour         [24]uint64
 	ByType         [MessageTypeMax]uint64
 	TodayDate      time.Time
 	Today          uint64
 	TotalCount     uint64
+	TotalTxtCount  uint64
 	TotalAvgLength uint64
 	Replies        uint64
 	Forward        uint64
@@ -52,30 +56,38 @@ func MakeUint(bval []byte, bucketName string, key string) uint64 {
 	}
 }
 
+func PutUint(value uint64) []byte {
+	bytes := make([]byte, 10)
+	n := binary.PutUvarint(bytes, value)
+	return bytes[:n]
+}
+
 func loadStats() {
 	// Load today
 	stats.TodayDate = time.Now()
 
-	err := db.View(func(tx *bolt.Tx) error {
+	err := db.Update(func(tx *bolt.Tx) error {
 		b, err := tx.CreateBucketIfNotExists([]byte("global"))
 		if err != nil {
 			return err
 		}
 
 		// Load total messages counter
-		bval := bucket.Get([]byte("count"))
+		bval := b.Get([]byte("count"))
 		stats.TotalCount = MakeUint(bval, "global", "count")
 
-		// Load total messages counter
+		bval = b.Get([]byte("avg"))
-		bval := bucket.Get([]byte("avg"))
 		stats.TotalAvgLength = MakeUint(bval, "global", "avg")
 
+		bval = b.Get([]byte("avgcount"))
+		stats.TotalTxtCount = MakeUint(bval, "global", "avgcount")
+
 		// Load total replies counter
-		bval = bucket.Get([]byte("replies"))
+		bval = b.Get([]byte("replies"))
 		stats.Replies = MakeUint(bval, "global", "replies")
 
 		// Load total replies counter
-		bval = bucket.Get([]byte("forward"))
+		bval = b.Get([]byte("forward"))
 		stats.Forward = MakeUint(bval, "global", "forward")
 
 		// Load hour counters
@@ -85,7 +97,7 @@ func loadStats() {
 		}
 
 		for i := 0; i < 24; i++ {
-			bval = bucket.Get([]byte(i))
+			bval = b.Get([]byte{byte(i)})
 			stats.ByHour[i] = MakeUint(bval, "hour", strconv.Itoa(i))
 		}
 
@@ -96,7 +108,7 @@ func loadStats() {
 		}
 
 		for i := 0; i < 7; i++ {
-			bval = bucket.Get([]byte(i))
+			bval = b.Get([]byte{byte(i)})
 			stats.ByWeekday[i] = MakeUint(bval, "weekday", strconv.Itoa(i))
 		}
 
@@ -107,7 +119,7 @@ func loadStats() {
 		}
 
 		todayKey := stats.TodayDate.Format("2006-1-2")
-		bval = bucket.Get([]byte(todayKey))
+		bval = b.Get([]byte(todayKey))
 		stats.Today = MakeUint(bval, "date", todayKey)
 
 		// Load user counters
@@ -118,6 +130,7 @@ func loadStats() {
 		}
 		b.ForEach(func(user, messages []byte) error {
 			stats.ByUserCount[string(user)] = MakeUint(messages, "users-count", string(user))
+			return nil
 		})
 
 		stats.ByUserAvgLen = make(map[string]uint64)
@@ -127,6 +140,17 @@ func loadStats() {
 		}
 		b.ForEach(func(user, messages []byte) error {
 			stats.ByUserAvgLen[string(user)] = MakeUint(messages, "users-avg", string(user))
+			return nil
+		})
+
+		stats.ByUserAvgCount = make(map[string]uint64)
+		b, err = tx.CreateBucketIfNotExists([]byte("users-avgcount"))
+		if err != nil {
+			return err
+		}
+		b.ForEach(func(user, messages []byte) error {
+			stats.ByUserAvgCount[string(user)] = MakeUint(messages, "users-avgcount", string(user))
+			return nil
 		})
 
 		// Load type counters
@@ -135,7 +159,7 @@ func loadStats() {
 			return err
 		}
 		for i := 0; i < MessageTypeMax; i++ {
-			bval = bucket.Get([]byte(i))
+			bval = b.Get([]byte{byte(i)})
 			stats.ByType[i] = MakeUint(bval, "types", strconv.Itoa(i))
 		}
 
@@ -143,3 +167,225 @@ func loadStats() {
 	})
 	assert(err)
 }
+
+func updateDate() {
+	err := db.Update(func(tx *bolt.Tx) error {
+		b := tx.Bucket([]byte("date"))
+		todayKey := stats.TodayDate.Format("2006-1-2")
+		err := b.Put([]byte(todayKey), PutUint(stats.Today))
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+	if err != nil {
+		log.Println("[updateDate] Couldn't save last day stats: " + err.Error())
+	}
+	stats.TodayDate = time.Now()
+	stats.Today = 0
+}
+
+func updateMean(currentMean, meanCount, newValue uint64) uint64 {
+	return ((currentMean * meanCount) + newValue) / (meanCount + 1)
+}
+
+func updateStats(message tg.APIMessage) {
+	//
+	// Local update
+	//
+
+	// DB Update flags
+	updatemean := false
+	updatetype := 0
+	updatereplies := false
+	updateforward := false
+
+	// Update total count
+	stats.TotalCount++
+
+	// Update individual user's count
+	username := message.User.Username
+	val, exists := stats.ByUserCount[username]
+	if !exists {
+		val = 0
+	}
+	stats.ByUserCount[username] = val + 1
+
+	// Update time counters
+	now := time.Now()
+	hour := now.Hour()
+	wday := now.Weekday()
+	stats.ByHour[hour]++
+	stats.ByWeekday[wday]++
+
+	// Check for day reset
+	if now.Day() != stats.TodayDate.Day() {
+		updateDate()
+	}
+	stats.Today++
+
+	// Text message
+	if message.Text != nil {
+		stats.ByType[MessageTypeText]++
+
+		// Update total and individual average
+		msglen := uint64(len(*(message.Text)))
+		if stats.TotalTxtCount > 0 {
+			stats.TotalAvgLength = updateMean(stats.TotalAvgLength, stats.TotalTxtCount, msglen)
+			stats.TotalTxtCount++
+		} else {
+			stats.TotalAvgLength = msglen
+			stats.TotalTxtCount = 1
+		}
+		val, exists = stats.ByUserAvgCount[username]
+		if exists {
+			stats.ByUserAvgLen[username] = updateMean(stats.ByUserAvgLen[username], val, msglen)
+			stats.ByUserAvgCount[username]++
+		} else {
+			stats.ByUserAvgLen[username] = msglen
+			stats.ByUserAvgCount[username] = 1
+		}
+		updatemean = true
+		updatetype = MessageTypeText
+	}
+	// Audio message
+	if message.Audio != nil {
+		stats.ByType[MessageTypeAudio]++
+		updatetype = MessageTypeAudio
+	}
+	// Photo
+	if message.Photo != nil {
+		stats.ByType[MessageTypePhoto]++
+		updatetype = MessageTypePhoto
+	}
+	// Sticker
+	if message.Sticker != nil {
+		stats.ByType[MessageTypeSticker]++
+		updatetype = MessageTypeSticker
+	}
+	// Video
+	if message.Video != nil {
+		stats.ByType[MessageTypeVideo]++
+		updatetype = MessageTypeVideo
+	}
+	// Voice message
+	if message.Voice != nil {
+		stats.ByType[MessageTypeVoice]++
+		updatetype = MessageTypeVoice
+	}
+	// Contact
+	if message.Contact != nil {
+		stats.ByType[MessageTypeContact]++
+		updatetype = MessageTypeContact
+	}
+	// Location
+	if message.Location != nil {
+		stats.ByType[MessageTypeLocation]++
+		updatetype = MessageTypeLocation
+	}
+	// Document
+	if message.Document != nil {
+		stats.ByType[MessageTypeDocument]++
+		updatetype = MessageTypeDocument
+	}
+	// Reply
+	if message.ReplyTo != nil {
+		stats.Replies++
+		updatereplies = true
+	}
+	// Forwarded message
+	if message.FwdUser != nil {
+		stats.Forward++
+		updateforward = true
+	}
+
+	//
+	// DB Update
+	//
+
+	err := db.Update(func(tx *bolt.Tx) error {
+		// Update total counters
+		b := tx.Bucket([]byte("global"))
+
+		err := b.Put([]byte("count"), PutUint(stats.TotalCount))
+		if err != nil {
+			return err
+		}
+
+		if updatemean {
+			err = b.Put([]byte("avg"), PutUint(stats.TotalAvgLength))
+			if err != nil {
+				return err
+			}
+			err = b.Put([]byte("avgcount"), PutUint(stats.TotalTxtCount))
+			if err != nil {
+				return err
+			}
+		}
+
+		if updatereplies {
+			err = b.Put([]byte("replies"), PutUint(stats.Replies))
+			if err != nil {
+				return err
+			}
+		}
+		if updateforward {
+			err = b.Put([]byte("forward"), PutUint(stats.Forward))
+			if err != nil {
+				return err
+			}
+		}
+
+		// Update time counters
+		b = tx.Bucket([]byte("hour"))
+		err = b.Put([]byte{byte(hour)}, PutUint(stats.ByHour[hour]))
+		if err != nil {
+			return err
+		}
+
+		b = tx.Bucket([]byte("weekday"))
+		err = b.Put([]byte{byte(wday)}, PutUint(stats.ByHour[wday]))
+		if err != nil {
+			return err
+		}
+
+		b = tx.Bucket([]byte("date"))
+		todayKey := stats.TodayDate.Format("2006-1-2")
+		err = b.Put([]byte(todayKey), PutUint(stats.Today))
+		if err != nil {
+			return err
+		}
+
+		// Update user counters
+		b = tx.Bucket([]byte("users-count"))
+		err = b.Put([]byte(username), PutUint(stats.ByUserCount[username]))
+		if err != nil {
+			return err
+		}
+
+		if updatemean {
+			b = tx.Bucket([]byte("users-avg"))
+			err = b.Put([]byte(username), PutUint(stats.ByUserAvgLen[username]))
+			if err != nil {
+				return err
+			}
+			b = tx.Bucket([]byte("users-avgcount"))
+			err = b.Put([]byte(username), PutUint(stats.ByUserAvgCount[username]))
+			if err != nil {
+				return err
+			}
+		}
+
+		// Update type counter
+		b = tx.Bucket([]byte("types"))
+		err = b.Put([]byte{byte(updatetype)}, PutUint(stats.ByType[updatetype]))
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
+	if err != nil {
+		log.Println("[updateStats] Got error while updating DB: " + err.Error())
+	}
+}

stats/users.go

@@ -0,0 +1,47 @@
+package main
+
+import (
+	"log"
+	"strings"
+
+	"github.com/boltdb/bolt"
+	"github.com/hamcha/clessy/tg"
+)
+
+var users map[string]string
+
+func getNick(apiuser tg.APIUser) {
+	if _, ok := users[apiuser.Username]; ok && strings.HasPrefix(users[apiuser.Username], apiuser.FirstName) {
+		// It's updated, don't bother
+		return
+	}
+
+	users[apiuser.Username] = apiuser.FirstName
+	if apiuser.LastName != "" {
+		users[apiuser.Username] += apiuser.LastName
+	}
+
+	err := db.Update(func(tx *bolt.Tx) error {
+		b := tx.Bucket([]byte("usernames"))
+		return b.Put([]byte(apiuser.Username), []byte(users[apiuser.Username]))
+	})
+	if err != nil {
+		log.Printf("[getNick] Could not update %s name: %s\n", apiuser.Username, err.Error())
+	}
+}
+
+func loadUsers() {
+	users = make(map[string]string)
+	err := db.Update(func(tx *bolt.Tx) error {
+		b, err := tx.CreateBucketIfNotExists([]byte("usernames"))
+		if err != nil {
+			return err
+		}
+		b.ForEach(func(user, name []byte) error {
+			users[string(user)] = string(name)
+			return nil
+		})
+		return nil
+	})
+	assert(err)
+}