DB640/internal/bot/bot.go

package bot

import (
	"log"
	"strings"

	"git.1750studios.com/ToddShepard/DB640/internal/config"
	"git.1750studios.com/ToddShepard/DB640/internal/database"
	"git.1750studios.com/ToddShepard/DB640/internal/telegram"
	"git.1750studios.com/ToddShepard/DB640/internal/twitter"
	"github.com/jinzhu/gorm"
	"github.com/sahilm/fuzzy"
)

var stream *twitter.Stream
var updates telegram.UpdateChan

// Init initzializes the bot and subscribes the magic hashtag feed
func Init() {
	var err error

	err = database.Open(config.C.Database.Dialect, config.C.Database.Connection)
	if err != nil {
		log.Fatalf("Could not establish database connection: %+v", err)
	}
	log.Printf("[DATABASE] Connection established\n")

	twitter.Init()
	stream, err = twitter.GetStreamForTag(config.C.Twitter.MagicHashtag)
	if err != nil {
		log.Fatalf("Could not establish twitter stream: %+v", err)
	}
	go twitter.StreamDemux(stream, handleHashtagTweet)
	log.Printf("[TWITTER] Connection established\n")

	err = telegram.Init()
	if err != nil {
		log.Fatalf("Could not establish telegram connection: %+v", err)
	}
	updates, err = telegram.GetChan()
	if err != nil {
		log.Fatalf("Could not establish telegram connection: %+v", err)
	}
	go handleTelegram()
	log.Printf("[TELEGRAM] Connection established\n")
}

// DeInit stops the stream and deinitzializes the bot
func DeInit() {
	stream.Stop()
	telegram.DeInit()
	database.Close()
}

func handleHashtagTweet(tweet *twitter.Tweet) {
	msg, tags := twitter.GetTextAndHashtags(tweet)
	// Ignore retweets
	if tweet.RetweetedStatus != nil {
		log.Printf("[TWITTER] %s: %s - IGNORED (RT)\n", tweet.User.ScreenName, msg)
		return
	}
	// Ignore replys
	if tweet.InReplyToStatusID != 0 {
		log.Printf("[TWITTER] %s: %s - IGNORED (RPLY)\n", tweet.User.ScreenName, msg)
		return
	}
	// Ignore if only magic hashtag is given
	if len(tags) < 2 {
		log.Printf("[TWITTER] %s: %s - IGNORED (<2#)\n", tweet.User.ScreenName, msg)
		return
	}
	log.Printf("[TWITTER] %s: %s\n", tweet.User.ScreenName, msg)
	go findCodes(tweet)
}

func findCodes(tweet *twitter.Tweet) {
	_, tags := twitter.GetTextAndHashtags(tweet)
	var betriebsstellen []database.Betriebsstelle
	for _, tag := range tags {
		code := strings.ReplaceAll(tag.Text, "_", " ")
		var bs database.Betriebsstelle
		if database.Db.First(&bs, "code = ?", code).Error != gorm.ErrRecordNotFound {
			betriebsstellen = append(betriebsstellen, bs)
		}
	}
	sendReply(tweet, betriebsstellen)
}

func sendReply(tweet *twitter.Tweet, betriebsstellen []database.Betriebsstelle) {
	var reply string
	for _, bs := range betriebsstellen {
		reply = reply + bs.Code + ": " + bs.Name + "\n"
	}
	reply = reply[0 : len(reply)-1]
	_, _, err := twitter.SendTweet(reply, tweet)
	if err != nil {
		log.Printf("Cannot send reply, error: %+v", err)
	}
}

func handleTelegram() {
	for update := range updates {
		if update.InlineQuery != nil {
			telegram.DoInlineQuery(update)
		}
		if update.Message == nil {
			// ignore any non-Message Updates
			continue
		}
		log.Printf("[TELEGRAM] %s: %s\n", update.Message.From.String(), update.Message.Text)

		if update.Message.IsCommand() {
			switch update.Message.Command() {
			case "start", "help":
				reply := `Willkommen beim DB 640 Telegram Bot!

Einfach den gewünschten DB 640 Betriebsstellencode schicken, und der Bot antwortet mit der zugehörigen Betriebsstelle!

	Zum Beispiel: Nb -> Wiener Neustadt Hbf (in Nb)

Alternativ nach Codes suchen:
	/find Name
	Zum Beispiel: /find Matzleinsdorf gibt die entsprechenden Betriebsstelle(n) und Code(s) aus.`
				telegram.SendReply(reply, update)
				continue
			case "find":
				if update.Message.CommandArguments() == "" {
					reply := "Benutze \"/find Name\" um einen Code für eine Betriebsstelle zu finden!"
					telegram.SendReply(reply, update)
					continue
				}
				var bs database.Betriebsstellen
				var reply string
				if err := database.Db.Find(&bs, "name LIKE ?", "%"+update.Message.CommandArguments()+"%").Error; err != nil || len(bs) == 0 {
					reply = "Keine Betriebsstelle mit Namen '" + update.Message.CommandArguments() + "' gefunden!"
				} else {
					results := fuzzy.FindFrom(update.Message.CommandArguments(), bs)
					for i, r := range results {
						reply = reply + bs[r.Index].Code + ": " + bs[r.Index].Name + "\n"
						if i == 49 {
							break
						}
					}
					reply = reply[0 : len(reply)-1]
				}
				err := telegram.SendReply(reply, update)
				if err != nil {
					log.Printf("Cannot send reply, error: %+v", err)
				}
				continue
			}
		}

		var bs database.Betriebsstelle
		var reply string
		if database.Db.First(&bs, "code = ?", update.Message.Text).Error != gorm.ErrRecordNotFound {
			reply = bs.Code + ": " + bs.Name
		} else {
			reply = update.Message.Text + ": nicht gefunden!"
		}
		err := telegram.SendReply(reply, update)
		if err != nil {
			log.Printf("Cannot send reply, error: %+v", err)
		}
	}
}