import telebot
import std/[asyncdispatch, logging, options, strutils, random, with, os, tables, times, sequtils]
import norm/[model, sqlite]

var L = newConsoleLogger(fmtStr="$levelname, [$time] ")
addHandler(L)

type
  CensoredData* = ref object of Model
    ftype*: string
    fhash*: string
    fileid*: string
    time*: float
    caption*: string

type
  BannedUsers* = ref object of Model
    userid*: int64
    bantime*: float
    bantype*: string

var RateLimiter = initTable[int64, seq[float]]()

var GroupMedia = initTable[int, string]()

let AdminID = getEnv("ADMIN_ID").split(",")

let dbConn* = getDb()

func NewCensoredData*(ftype = ""; fhash = ""; fileid = ""; time = 0.0; caption = ""):
  CensoredData = CensoredData(ftype: ftype, fhash: fhash, fileid: fileid, time: time, caption: caption)

func NewBannedUsers*(userid = int64(0), bantime = 0.0, bantype = ""):
  BannedUsers = BannedUsers(userid: userid, bantime: bantime, bantype: bantype)

dbConn.createTables(NewCensoredData())
dbConn.createTables(NewBannedUsers())

proc ManageRateLimit(): void=
  for i in RateLimiter.keys().toSeq():
    if RateLimiter[i][^1] - RateLimiter[i][0] >= 60:
      RateLimiter.del(i)
    elif len(RateLimiter[i]) >= 20:
      var BannedUser = NewBannedUsers(i, epochTime(), "auto")
      RateLimiter.del(i)
      with dbConn:
        insert BannedUser
  if dbConn.exists(BannedUsers, "bantype = ? and ? - bantime >= 1800", "auto", epochTime()):
    var TempData = @[NewBannedUsers()]
    dbConn.select(TempData,  "bantype = ? and ? - bantime >= 1800", "auto", epochTime())
    for i in TempData:
      var e = i
      dbConn.delete(e)

proc OldDataCleanup(): void=
  if dbConn.exists(CensoredData, "? - time >= 15780000", epochTime()):
    var TempData = @[NewCensoredData()]
    dbConn.select(TempData,  "? - time >= 15780000", epochTime())
    for i in TempData:
      var e = i
      dbConn.delete(e)

proc generate_hash(): string=
  result = newString(7) 
  const charset = {'a' .. 'z', 'A' .. 'Z', '0' .. '9'}
  for i in 0..6:
    result[i] = sample(charset)
  return result

proc updateHandler(b: Telebot, u: Update): Future[bool] {.async, gcsafe.} =
  let response = u.message.get
  ManageRateLimit()
  if not dbConn.exists(BannedUsers, "userid = ?", response.chat.id):
    if RateLimiter.contains(response.chat.id):
      RateLimiter[response.chat.id].insert(epochTime())
    else:
      RateLimiter[response.chat.id] = @[epochTime()]
    if response.text.isSome:
      let message = response.text.get
      if message == "/start":
        discard await b.sendMessage(response.chat.id, "Hey, To create a censored post, you can share any album, video, photo, gif, sticker, etc. The messages could then be forwarded to any chat for them to view")
      elif message.contains("/ban"):
        if $response.chat.id in AdminID:
          let user = message.split(" ")
          var BannedUser = NewBannedUsers(int64(parseInt(user[1])), epochTime(), "permanent")
          with dbConn:
            insert BannedUser
          discard await b.sendMessage(response.chat.id, "Banned!")
      elif message.contains("/start"):
        let deeplink = message.split(" ")
        if not dbConn.exists(CensoredData, "fhash = ?", deeplink[1]):
          discard await b.sendMessage(response.chat.id, "Media does not exist on database, ask the sender to censor this again!")
        else:
          var TempData = @[NewCensoredData()]
          dbConn.select(TempData, "fhash = ?", deeplink[1])
          if len(TempData) > 1:
            var inputmedia = newSeq[InputMediaPhoto]()
            for i in TempData:
              if i.caption != "":
                inputmedia.insert(InputMediaPhoto(kind: i.ftype, media: i.fileid, caption: some(i.caption)))
              else:
                inputmedia.insert(InputMediaPhoto(kind: i.ftype, media: i.fileid))
            discard await b.sendMediaGroup($response.chat.id, media=inputmedia)
          else:
            if TempData[0].ftype == "photo":
              discard await b.sendPhoto(response.chat.id, TempData[0].fileid, TempData[0].caption)
            elif TempData[0].ftype == "document":
              discard await b.sendDocument(response.chat.id, TempData[0].fileid, TempData[0].caption)
            elif TempData[0].ftype == "video":
              discard await b.sendVideo(response.chat.id, TempData[0].fileid, caption=TempData[0].caption)
            elif TempData[0].ftype == "videoNote":
              discard await b.sendVideoNote(response.chat.id, TempData[0].fileid)
            elif TempData[0].ftype == "animation":
              discard await b.sendAnimation(response.chat.id, TempData[0].fileid, caption=TempData[0].caption)
            elif TempData[0].ftype == "sticker":
              discard await b.sendSticker(response.chat.id, TempData[0].fileid)
    else:
      var fileid = ""
      var ftype = ""
      var fcaption = ""
      if response.caption.isSome:
        fcaption = response.caption.get
      if response.document.isSome:
        fileid = response.document.get.fileId
        ftype = "document"
      elif response.video.isSome:
        fileid = response.video.get.fileId
        ftype = "video"
      elif response.videoNote.isSome:
        fileid = response.videoNote.get.fileId
        ftype = "videoNote"
      elif response.animation.isSome:
        fileid = response.animation.get.fileId
        ftype = "animation"
      elif response.photo.isSome:
        fileid = response.photo.get[0].fileId
        ftype = "photo"
      elif response.sticker.isSome:
        fileid = response.sticker.get.fileId
        ftype = "sticker"

      if response.mediaGroupId.isSome:
        if parseInt(response.mediaGroupId.get) notin GroupMedia.keys().toSeq():
          let filehash = generate_hash()
          GroupMedia[parseInt(response.mediaGroupId.get)] = filehash
          var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption)
          with dbConn:
                insert CensoredRow
          discard await b.sendMessage(response.chat.id, "*Censored " & capitalizeAscii(ftype) & "*\n\n[Tap to View](tg://resolve?domain=" & b.username & "&start=" & filehash & ")", parseMode = "Markdown")
        else:
          let filehash = GroupMedia[parseInt(response.mediaGroupId.get)]
          var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption)
          with dbConn:
                insert CensoredRow

      else:
        let filehash = generate_hash()
        var CensoredRow = NewCensoredData(ftype, filehash, fileid, epochTime(), fcaption)
        with dbConn:
              insert CensoredRow
        discard await b.sendMessage(response.chat.id, "*Censored " & capitalizeAscii(ftype) & "*\n\n[Tap to View](tg://resolve?domain=" & b.username & "&start=" & filehash & ")", parseMode = "Markdown")
    OldDataCleanup()

let bot = newTeleBot(getEnv("TELEGRAM_TOKEN"))
bot.onUpdate(updateHandler)
bot.poll(timeout=300)