package main import ( "encoding/binary" "encoding/json" "log" "strconv" "strings" "time" "github.com/boltdb/bolt" "github.com/hamcha/clessy/tg" ) const ( MessageTypeText int = 0 MessageTypeAudio int = 1 MessageTypePhoto int = 2 MessageTypeSticker int = 3 MessageTypeVideo int = 4 MessageTypeVoice int = 5 MessageTypeContact int = 6 MessageTypeLocation int = 7 MessageTypeDocument int = 8 MessageTypeMax int = 9 ) type Stats struct { ByUserCount map[string]uint64 ByWeekday [7]uint64 ByHour [24]uint64 ByType [MessageTypeMax]uint64 ByDay map[string]uint64 TodayDate time.Time Today uint64 TotalCount uint64 } var stats Stats type UserCount map[string]uint64 var words map[string]UserCount func MakeUint(bval []byte, bucketName string, key string) uint64 { if bval != nil { intval, bts := binary.Uvarint(bval) if bts > 0 { return intval } else { log.Printf("[%s] Value of key \"%s\" is NaN: %v\r\n", bucketName, key, bval) return 0 } } else { log.Printf("[%s] Key \"%s\" does not exist, set to 0\n", bucketName, key) return 0 } } 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.Update(func(tx *bolt.Tx) error { b, err := tx.CreateBucketIfNotExists([]byte("global")) if err != nil { return err } // Load total messages counter stats.TotalCount = MakeUint(b.Get([]byte("count")), "global", "count") // Load hour counters b, err = tx.CreateBucketIfNotExists([]byte("hour")) if err != nil { return err } for i := 0; i < 24; i++ { stats.ByHour[i] = MakeUint(b.Get([]byte{byte(i)}), "hour", strconv.Itoa(i)) } // Load weekday counters b, err = tx.CreateBucketIfNotExists([]byte("weekday")) if err != nil { return err } for i := 0; i < 7; i++ { stats.ByWeekday[i] = MakeUint(b.Get([]byte{byte(i)}), "weekday", strconv.Itoa(i)) } // Load day counters stats.ByDay = make(map[string]uint64) b, err = tx.CreateBucketIfNotExists([]byte("date")) if err != nil { return err } b.ForEach(func(day, messages []byte) error { stats.ByDay[string(day)] = MakeUint(messages, "date", string(day)) return nil }) todayKey := stats.TodayDate.Format("2006-1-2") stats.Today = MakeUint(b.Get([]byte(todayKey)), "date", todayKey) // Load user counters stats.ByUserCount = make(map[string]uint64) b, err = tx.CreateBucketIfNotExists([]byte("users-count")) if err != nil { return err } b.ForEach(func(user, messages []byte) error { stats.ByUserCount[string(user)] = MakeUint(messages, "users-count", string(user)) return nil }) // Load type counters b, err = tx.CreateBucketIfNotExists([]byte("types")) if err != nil { return err } for i := 0; i < MessageTypeMax; i++ { stats.ByType[i] = MakeUint(b.Get([]byte{byte(i)}), "types", strconv.Itoa(i)) } // Load dictionary b, err = tx.CreateBucketIfNotExists([]byte("words")) if err != nil { return err } words = make(map[string]UserCount) b.ForEach(func(word, ucount []byte) error { var val UserCount err := json.Unmarshal(ucount, &val) if err != nil { return err } words[string(word)] = val return nil }) return nil }) assert(err) } func updateDate() { dateKey := stats.TodayDate.Format("2006-1-2") err := db.Update(func(tx *bolt.Tx) error { b := tx.Bucket([]byte("date")) err := b.Put([]byte(dateKey), 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.ByDay[dateKey] = stats.Today 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 updatetype := 0 // 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]++ updatetype = MessageTypeText // Process words processWords(message) } // 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 } // // 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 } // 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.ByWeekday[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 } // 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()) } } func processWords(message tg.APIMessage) { if len(*(message).Text) < 3 { return } wordList := strings.Split(*(message.Text), " ") err := db.Update(func(tx *bolt.Tx) error { b := tx.Bucket([]byte("words")) for _, word := range wordList { if len(word) < 3 { continue } word = strings.Trim(word, " ?!.,:;/-_()[]{}'\"+=*^\n") count, ok := words[word] if !ok { count = make(UserCount) } val, ok := count[message.User.Username] if !ok { val = 0 } count[message.User.Username] = val + 1 j, err := json.Marshal(count) if err != nil { return err } b.Put([]byte(word), j) } return nil }) if err != nil { log.Println("[processWords] Error encountered: " + err.Error()) } }