abitbol/bot.go

package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/gocolly/colly"
	tb "gopkg.in/tucnak/telebot.v2"
)

const (
	dbFile = "quotes.json"
)

// Quote represent a quote of the flim
type Quote struct {
	Quote      string `json:"quote"`
	YoutubeKey string `json:"youtube_key"`
}

// ThumbURL get url of the thumbnail
func (q *Quote) ThumbURL() string {
	return fmt.Sprintf(
		"https://img.youtube.com/vi/%s/0.jpg",
		q.YoutubeKey,
	)

}

// VideoURL get url of the video
func (q *Quote) VideoURL() string {
	return fmt.Sprintf(
		"https://www.youtube.com/watch?v=%s",
		q.YoutubeKey,
	)
}

func main() {
	var quotes []Quote

	// Load quote database
	file, err := os.Open(dbFile)
	if err != nil {
		log.Fatalf("Could not open quote file: %s\n", err)
		return
	}
	defer file.Close()

	bytes, err := ioutil.ReadAll(file)
	if err != nil {
		log.Fatalf("Could not read quote file: %s\n", err)
		return
	}

	if err = json.Unmarshal(bytes, &quotes); err != nil {
		log.Fatalf("Could not Unmarshal  quote file: %s\n", err)
		return
	}

	// Index quote database
	index, err := indexQuotes(quotes)
	if err != nil {
		log.Fatalf("Could not index quotes: %s\n", err)
		return
	}

	// Load telegram bot
	b, err := tb.NewBot(tb.Settings{
		Token:  os.Getenv("TELEGRAM_TOKEN"),
		Poller: &tb.LongPoller{Timeout: 10 * time.Second},
	})
	if err != nil {
		log.Fatalf("Colud not connect to telegram: %s\n", err)
		return
	}

	// Add bot capabilities
	b.Handle(tb.OnQuery, func(q *tb.Query) {
		quotes, err := index.Search(q.Text)
		if err != nil {
			log.Println("Pas de fichier citation trouvé")
			return
		}
		results := make(tb.Results, len(quotes))
		for i, quote := range quotes {
			url := quote.VideoURL()
			thumb := quote.ThumbURL()
			log.Printf("{%s, %s}\n", url, thumb)
			results[i] = &tb.VideoResult{
				URL:      url,
				MIME:     "text/html",
				ThumbURL: thumb,
				Title:    quote.Quote,
			}
			results[i].SetContent(
				&tb.InputTextMessageContent{
					Text: fmt.Sprintf(
						"%s\n\n%s",
						quote.Quote,
						url,
					),
					ParseMode:      tb.ModeHTML,
					DisablePreview: false,
				},
			)
			results[i].SetResultID(strconv.Itoa(i))
		}
		err = b.Answer(q, &tb.QueryResponse{
			Results:           results,
			SwitchPMParameter: "Ajoute moi",
			CacheTime:         0, // One minut
		})

		if err != nil {
			log.Println(err)
		}
	})

	// Start bot
	log.Println("Bot started")
	b.Start()
}

// Update get a fresh list of quotes from george-abitbol.fr
func Update() error {
	var videos []Quote

	c := colly.NewCollector(
		colly.AllowedDomains("george-abitbol.fr"),
	)

	// Find and visit all links
	c.OnHTML("img[data-original]", func(e *colly.HTMLElement) {
		videos = append(
			videos,
			Quote{
				Quote:      e.Attr("alt"),
				YoutubeKey: strings.Split(e.Attr("data-original"), "/")[4],
			},
		)
	})

	c.Visit("http://george-abitbol.fr")

	file, err := os.Create(dbFile)
	if err != nil {
		return fmt.Errorf("could not create quotes files: %w", err)
	}
	defer file.Close()

	bytes, err := json.Marshal(&videos)
	if err != nil {
		return fmt.Errorf("could not marshal quotes: %w", err)
	}

	file.Write(bytes)

	return nil
}