main.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"time"

	_ "go.uber.org/automaxprocs"

	"github.com/caarlos0/env"
	"github.com/go-pg/pg"
	"github.com/go-telegram-bot-api/telegram-bot-api"
	"golang.org/x/net/proxy"
)

type context struct {
	TgToken            string `env:"TG_TOKEN,required"`
	PgConnectionString string `env:"PG_CONNECTION_STRING,required"`
	Socks5Proxy        string `env:"SOCKS5_PROXY"`

	Bot *tgbotapi.BotAPI
	DB  *pg.DB
}

type raceRow struct {
	tableName string `sql:"races"`

	ID               int64
	ChatID           int64
	UserID           int
	InitialMessageID int
	CreatedAt        string
	Duration         string
	Status           string
	Comment          string
}

var goReplyKeyboard = tgbotapi.NewInlineKeyboardMarkup(
	tgbotapi.NewInlineKeyboardRow(
		tgbotapi.NewInlineKeyboardButtonData("Show remaining time", "left"),
	),
)

var replyKeyboard = tgbotapi.NewReplyKeyboard(
	tgbotapi.NewKeyboardButtonRow(
		tgbotapi.NewKeyboardButton("/go"),
		tgbotapi.NewKeyboardButton("/left"),
		tgbotapi.NewKeyboardButton("/stats"),
		tgbotapi.NewKeyboardButton("/cancel"),
	),
	tgbotapi.NewKeyboardButtonRow(
		tgbotapi.NewKeyboardButton("/setcomment"),
		tgbotapi.NewKeyboardButton("/getcomment"),
		tgbotapi.NewKeyboardButton("/report"),
	),
)

func (ctx *context) getBotAPI() (*tgbotapi.BotAPI, error) {
	if ctx.Socks5Proxy != "" {
		dialer, err := proxy.SOCKS5("tcp", ctx.Socks5Proxy, nil, proxy.Direct)
		if err != nil {
			log.Fatalf("can't connect to the proxy: %s", err)
		}
		httpTransport := &http.Transport{}
		httpClient := &http.Client{Transport: httpTransport}
		httpTransport.Dial = dialer.Dial

		return tgbotapi.NewBotAPIWithClient(ctx.TgToken, httpClient)
	}

	return tgbotapi.NewBotAPI(ctx.TgToken)
}

func (ctx *context) getRaceComment(userID int) string {
	var row struct {
		Count int64
	}

	_, err := ctx.DB.QueryOne(&row, `
	  select count(*)
		  from races
		 where user_id = ?
		   and status = 'finished'
			 and created_at::date = now()::date
	`, userID)

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

	if row.Count == 2 {
		return "Кажется вы сели серьёзно поработать. План на сегодня: 14 помидор."
	}

	if row.Count == 7 {
		return "Вы собрали половину запланированных помидор."
	}

	if row.Count == 12 {
		return "Ещё два помидора и план готов."
	}

	if row.Count == 14 {
		return "14 помидор - цель достигнута и можно отдохнуть."
	}

	return ""
}

func (ctx *context) checkFinished() {
	var rows []raceRow

	for {
		_, err := ctx.DB.Query(&rows, `
			update races set status = 'finished'
			where status = 'pending' and created_at + duration < now()
			returning chat_id, user_id, initial_message_id
		`)

		if err != nil {
			log.Println(err)
			time.Sleep(1000 * time.Millisecond)
			continue
		}

		for _, row := range rows {
			msg := tgbotapi.NewMessage(row.ChatID, "")
			msg.Text = "Race finished."
			msg.ReplyToMessageID = row.InitialMessageID

			comment := ctx.getRaceComment(row.UserID)

			if comment != "" {
				msg.Text += " " + comment
			}

			if _, err := ctx.Bot.Send(msg); err != nil {
				log.Println("failed to send: ", err)
			}
		}

		time.Sleep(1000 * time.Millisecond)
	}
}

type forgetRows struct {
	ChatID int64
}

func (ctx *context) forgetRestart() {
	var rows []forgetRows

	for {
		_, err := ctx.DB.Query(&rows, `
			select f.chat_id
			  from races f
			  left join races p on p.user_id = f.user_id and p.status = 'pending' and  p.created_at > f.created_at
			 where p.id is null
			   and f.status = 'finished'
			   and f.created_at + f.duration > now() - interval '11 minutes'
		`)

		if err != nil {
			log.Println(err)
			time.Sleep(5 * time.Minute)
			continue
		}

		for _, row := range rows {
			msg := tgbotapi.NewMessage(row.ChatID, "")
			msg.Text = "Прошло 5 минут. Если вы работаете, то забыли сказать /go"

			if _, err := ctx.Bot.Send(msg); err != nil {
				log.Println("failed to send: ", err)
			}
		}

		time.Sleep(5 * time.Minute)
	}
}

func statusReply(msg *tgbotapi.MessageConfig, db *pg.DB, from int) {
	var row struct {
		MinutesLeft int
	}

	_, err := db.QueryOne(&row, `
		select ceil(extract('epoch' from duration - (now() - created_at)) / 60) as minutes_left
			from races
		 where status = 'pending'
			 and user_id = ?
	`, from)

	if err != nil {
		log.Println(err)
		msg.Text = "An active race not found."
	} else {
		msg.Text = fmt.Sprintf("%d minutes left.", row.MinutesLeft)
	}
}

type dataRow struct {
	ID        int     `json:"id"`
	ChatID    int     `json:"chatId"`
	CreatedAt string  `json:"createdAt"`
	Status    string  `json:"status"`
	Duration  string  `json:"duration"`
	UserName  string  `json:"userName"`
	Index     float32 `json:"index"`
}

func logRequest(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		log.Printf("%s %s %s\n", r.RemoteAddr, r.Method, r.URL)
		handler.ServeHTTP(w, r)
	})
}

func runServer(db *pg.DB) {
	fs := http.FileServer(http.Dir("dist"))

	http.Handle("/", fs)
	http.HandleFunc("/data", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "http://localhost:1234")

		if r.Method == "OPTIONS" {
			return
		}

		var rows []dataRow

		_, err := db.Query(&rows, `
		  select r.id, r.chat_id, r.created_at, r.status, r.duration
			     , coalesce(u.user_name, u.first_name || ' ' || u.last_name) as user_name
					 , extract(epoch from created_at - now()) / 3600 + 24 as index
			  from races r
				join users u on u.id = r.user_id
			 where created_at > now() - interval '1 day'
			 order by r.created_at
		`)

		if err != nil {
			log.Println(err)
			http.Error(w, http.StatusText(500), 500)
			return
		}

		if err := json.NewEncoder(w).Encode(rows); err != nil {
			log.Println(err)
			http.Error(w, http.StatusText(500), 500)
		}
	})

	log.Fatal(http.ListenAndServe(":3333", logRequest(http.DefaultServeMux)))
}

func getLastComment(db *pg.DB, userID int, chatID int64) string {
	var row struct {
		Comment string
	}

	_, err := db.QueryOne(&row, `
		select comment
		  from races
		 where user_id = ?
		   and chat_id = ?
		 order by created_at desc
	   limit 1
	`, userID, chatID)

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

	return row.Comment
}

func main() {
	ctx := context{}
	err := env.Parse(&ctx)

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

	dbOptions, err := pg.ParseURL(ctx.PgConnectionString)

	if err != nil {
		panic(err)
	}

	db := pg.Connect(dbOptions)
	defer func() {
		if err := db.Close(); err != nil {
			panic(err)
		}
	}()

	_, err = db.Exec("select 1")

	if err != nil {
		panic(err)
	}

	ctx.DB = db
	ctx.Bot, err = ctx.getBotAPI()

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

	bot := ctx.Bot

	log.Printf("Authorized on account %s", bot.Self.UserName)

	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60

	updates, err := bot.GetUpdatesChan(u)

	go ctx.checkFinished()
	go ctx.forgetRestart()
	go runServer(db)

	nicknameR := regexp.MustCompile("@([a-zA-Z0-9_]+)")

	for update := range updates {
		if update.CallbackQuery != nil {
			if update.CallbackQuery.Data == "left" {
				msg := tgbotapi.NewMessage(update.CallbackQuery.Message.Chat.ID, "")
				statusReply(&msg, db, update.CallbackQuery.From.ID)

				if _, err := ctx.Bot.Send(msg); err != nil {
					log.Println("failed to send: ", err)
				}
			}
		}

		if update.Message == nil { // ignore any non-Message Updates
			continue
		}

		if update.Message.IsCommand() {
			// Is a group
			if update.Message.Chat.ID < 0 {
				continue
			}

			msg := tgbotapi.NewMessage(update.Message.Chat.ID, "")
			msg.ReplyToMessageID = update.Message.MessageID
			msg.ReplyMarkup = replyKeyboard

			switch update.Message.Command() {
			case "go", "25":
				comment := update.Message.CommandArguments()

				if comment == "" {
					comment = getLastComment(db, update.Message.From.ID, update.Message.Chat.ID)
				}

				err = db.Insert(&raceRow{
					ChatID:           update.Message.Chat.ID,
					UserID:           update.Message.From.ID,
					InitialMessageID: update.Message.MessageID,
					Comment:          comment,
				})

				if err != nil {
					if err.Error() == `ERROR #23505 duplicate key value violates unique constraint "race_pending_idx"` {
						msg.Text = "You have an active race"
					} else {
						log.Println("go: ", err)
						msg.Text = err.Error()
					}
				} else {
					msg.Text = "Race started."
					msg.ReplyMarkup = goReplyKeyboard

					if comment != "" {
						msg.Text += " Comment:\n" + comment
					} else {
						msg.Text += "You can put a comment like this /go issue #553 or update /set-comment new comment text"
					}

					// update user data
					_, err := db.Exec(`
					  insert into users (id, first_name, last_name, user_name, last_chat_id)
						values (?, ?, ?, ?, ?)
						on conflict (id)
						do update set
						first_name = excluded.first_name,
						last_name = excluded.last_name,
						last_chat_id = excluded.last_chat_id,
						user_name = excluded.user_name;
					`, update.Message.From.ID, update.Message.From.FirstName, update.Message.From.LastName, update.Message.From.UserName, update.Message.Chat.ID)

					if err != nil {
						log.Println("failed to update user: ", err)
					}
				}

			case "getcomment":
				msg.Text = getLastComment(db, update.Message.From.ID, update.Message.Chat.ID)

			case "setcomment":
				comment := update.Message.CommandArguments()

				if comment != "" {
					res, err := db.Exec(`
						update races
							 set comment = ?
						 where user_id = ?
							 and chat_id = ?
							 and status = 'pending'
					`, update.Message.CommandArguments(), update.Message.From.ID, update.Message.Chat.ID)

					if err != nil {
						log.Println("failed to update user: ", err)
						msg.Text = err.Error()
					} else if res.RowsAffected() == 0 {
						msg.Text = "No affected rows."
					} else {
						msg.Text = "Comment updated."
					}
				} else {
					msg.Text = "Argument not found, use /setcomment text..."
				}

			case "report":
				var rows []struct {
					Comment string
					Count   int
				}

				_, err := db.Query(&rows, `
				  select comment, count(*)
					  from races
				 	 where user_id = ?
						 and chat_id = ?
					   and status = 'finished'
						 and created_at > now() - interval '1 day'
					 group by 1
				`, update.Message.From.ID, update.Message.Chat.ID)

				if err == nil {
					if len(rows) > 0 {
						for _, row := range rows {
							msg.Text += fmt.Sprintf("\n(%d) %s", row.Count, row.Comment)
						}
					} else {
						msg.Text = "Races not found."
					}
				} else {
					msg.Text = err.Error()
					log.Println("failed to get report", err)
				}

			case "cancel":
				res, err := db.Exec(`update races set status = 'canceled' where status = 'pending' and user_id = ? and chat_id = ?`, update.Message.From.ID, update.Message.Chat.ID)

				if err != nil {
					log.Println("cancel: ", err)
					msg.Text = err.Error()
				} else if res.RowsAffected() > 0 {
					msg.Text = "Ok, no problem."
				} else {
					msg.Text = "You have not active races."
				}

			case "status", "left":
				statusReply(&msg, db, update.Message.From.ID)

			case "stats":
				var rows []struct {
					Day   string
					Count int
				}

				_, err := db.Query(&rows, `
					select created_at::date as day, count(*)
					  from races
				 	 where status = 'finished'
					   and created_at > now() - interval '7 days'
						 and user_id = ?
					 group by 1
					 order by 1
			  `, update.Message.From.ID)

				if err != nil {
					log.Println("cancel: ", err)
					msg.Text = err.Error()
				} else if len(rows) > 0 {
					for index, row := range rows {
						msg.Text += fmt.Sprintf("%s %d", row.Day, row.Count)

						if index < len(rows)-1 {
							msg.Text += "\n"
						}
					}
				} else {
					msg.Text = fmt.Sprintf("You have no completed pomodoros today. But don't be upset! There is still time to start one.")
				}
			default:
				msg.Text = "I don't know that command. Available commands: /go, /cancel, /status, /stats"
			}

			if _, err := ctx.Bot.Send(msg); err != nil {
				log.Println("failed to send: ", err)
			}
		} else { // if update.Message.Chat.ID < 0 {
			// listen group

			log.Println(update.Message.Text)
			nickname := nicknameR.FindString(update.Message.Text)

			if nickname == "" {
				continue
			}

			var user struct {
				ID         int
				LastChatID int64
			}

			_, err := ctx.DB.QueryOne(&user, `
			  select id, last_chat_id
				  from users
				 where user_name = ?
			`, nickname[1:])

			if err != nil {
				log.Println("user not found", err, nickname[1:])
				continue
			}

			var raceData struct {
				PendingCount int64
			}

			_, err = ctx.DB.QueryOne(&raceData, `
			  select count(*) as pending_count
				  from races
				 where user_id = ?
				   and status = 'pending'
			`, user.ID)

			if err != nil {
				log.Println("find races", err)
				continue
			}

			if raceData.PendingCount == 0 {
				msg := tgbotapi.NewMessage(user.LastChatID, update.Message.Text)
				msg.ReplyToMessageID = update.Message.MessageID

				if _, err := ctx.Bot.Send(msg); err != nil {
					log.Println("failed to send: ", err)
				}
			}
		}
	}
}

schema.sql

create type race_status as enum (
  'pending',
  'finished',
  'canceled'
);
create table races (
  id serial primary key,
  chat_id int not null,
  user_id int not null,
  initial_message_id int not null,
  created_at timestamptz not null default statement_timestamp(),
  duration interval not null default '25 minutes',
  status race_status not null default 'pending'
);
create unique index race_pending_idx on races (user_id) where (status = 'pending');
create table users (
  id bigint primary key,
  first_name text,
  last_name text,
  user_name text
);
alter table races add column comment text;