import os, datetime, pytz, discord, shutil, requests, subprocess, asyncio
from discord.ext import commands
from time import perf_counter

CLIENT_SECRET = "Ca3tbHTsbImaDmlzXqnpUQLONYtBY1eA"
TOKEN = "OTE5ODk0NDEwODkwNzI3NDg1.G96t2g.XV5efYvSNKqsNvPFxwmviLi-oY6oAYVPTzV78A"
ADMINS = [766145655038410763, 937495666471628831]
basepath = r"C:\Users\Pawin\Software\real-esrgan\discordbot"

def getloggingtime():
  return datetime.datetime.now(pytz.timezone('Asia/Bangkok')).strftime('%d/%m/%Y %H:%M:%S')

def generate_id(path):
  return str(len(os.listdir(path))+1).zfill(3)

def file_cleanup():
  dirlist = [fr"{basepath}\input", fr"{basepath}\output"]
  for directory in dirlist:
    shutil.rmtree(directory)
    os.mkdir(directory)
  print("BOT: Sucessfully cleaned directories")

file_cleanup()

client = commands.Bot(
    command_prefix = "!",
    case_insensitive = True,
    help_command = None,
    activity = discord.Streaming(name = f"/help | Local Instance", url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ")
)

"""BOT EVENTS"""
#Show Status When Bot Is Ready
@client.event
async def on_ready():
  print('BOT: We have successfully logged in as {0.user}\n'.format(client))

CurrentlyProcessingJob = False

@client.event
async def on_message(message):
  global CurrentlyProcessingJob
  
  msg = str(message.content)
  validMediaChannels = [1083235916354703430, 896342418859892736, 915795164721709077]
  channel = message.channel
  author = message.author

  channelId = int(channel.id)
  authorId = int(author.id)
  #print(f"Channel id is {channelId}, author id is {authorId}")

  imageList = []
  for item in message.attachments:
    if item.content_type.startswith("image"):
      imageList.append(item.url)

  if message.author.bot or msg.startswith("p!"):
    pass
  
  elif (authorId in ADMINS) and (channelId in validMediaChannels) and (message.attachments) and (imageList):
    print(f'IMAGE: New attachment recieved from {author} in channel {channel}. Processing...')

    jobReceiptMessage = await message.reply(f"\nUpscale job received for message {message.id}", mention_author=False)
    jobStatusMessage = await message.reply(f"Status: Upscale queued.", mention_author=False)

    if CurrentlyProcessingJob:
      pendingTemplate = f"__Status__: Upscale %CURRENTSTEP% for job *{message.id}*"
      await jobStatusMessage.edit(content=pendingTemplate.replace("%CURRENTSTEP%", "Pending 🕰️"))
      print(f"IMAGE: There is already a job running. Message id {message.id} has been put on queue")
      while CurrentlyProcessingJob:
        await asyncio.sleep(1)
        #wait untill it finishes
      await jobStatusMessage.edit(content=pendingTemplate.replace("%CURRENTSTEP%", "Starting 🏃"))
      print(f"IMAGE: Starting job {message.id}")
    CurrentlyProcessingJob = True
    
    attachmentCount = len(imageList)
    multipleImages = attachmentCount > 1
    if multipleImages:
      batchstart = perf_counter()
      taskType = "Batch Upscale"
    else:
      taskType = "Image Upscaling"

    for i in range(attachmentCount):

      """PREPROCESSING"""
      fileid = f"{message.id}_{i}"
      statusTemplate = f"__Status__: %CURRENTSTEP% image from message *{message.id}* **({i+1}/{attachmentCount})**."
      await jobStatusMessage.edit(content=statusTemplate.replace("%CURRENTSTEP%", "🔄 Preprocessing"))

      starttime = perf_counter()
      url = imageList[i]

      extension = url.split(".")[-1]
      #fileid = generate_id(f"{basepath}\input")

      inputfile = fr"{basepath}\input\{fileid}.{extension}"
      outputfile = fr"{basepath}\output\{fileid}.png"
      
      with open (inputfile, "wb") as f:
        f.write(requests.get(url).content)
      


      """UPSCALE CONFIG"""
      ai_model = "realesrgan-x4plus-anime"
      #Can be realesr-animevideov3 (default) | realesrgan-x4plus | realesrgan-x4plus-anime | realesrnet-x4plus
      if ai_model == "realesr-animevideov3":
        scale = "2"
        tilesize = "768"
        resizeRequired = False
      else:
        scale = "4"
        tilesize = "256"
        resizeRequired = True

      #Scale 2 only works with the default model
      execute_upscale = fr"C:\Users\Pawin\Software\real-esrgan\realesrgan-ncnn-vulkan.exe -i {inputfile} -o {outputfile} -s {scale} -n {ai_model} -f png -t {tilesize} -j 1:1:1"



      """UPSCALING"""
      pendtime = perf_counter()
      preprocessingtime = round((pendtime-starttime), 2)
      print(f"PREPROCESS: Completed in {preprocessingtime}s.\nUPSCALE:")
      await jobStatusMessage.edit(content=(statusTemplate.replace("%CURRENTSTEP%", "🧠 Upscaling")+"\nThis may take a while..."))
      #os.system(execute_upscale)
      subprocess.run(execute_upscale, shell=True)

      uendtime = perf_counter()

      upscaletime = round((uendtime-pendtime),2)
      print(f"UPSCALE: Completed in {upscaletime}s.")
    


      """RESIZING"""
      if resizeRequired:
        await jobStatusMessage.edit(content=statusTemplate.replace("%CURRENTSTEP%", "⚙️ Resizing"))
        subprocess.run(f"mogrify -resize 50% {outputfile}", shell=True)
        rendtime = perf_counter()
        resizingtime = round((rendtime-uendtime), 2)
        print(f"RESIZE: Completed in {resizingtime}s.")
      else:
        resizingtime = 0



      """DELIVERING"""
      await jobStatusMessage.edit(content=statusTemplate.replace("%CURRENTSTEP%", "✉️ Sending"))
      
      try:      
        #file=discord.File(outputfile)      
        localImageUrl = f"http://thinkpad.pawin.tk/upscale/{fileid}.png"
        
        
        #outputImageUrl = jobResultMessage.embeds[0].image.url
        #uncomment above to use outside local network
        #outputImageUrl = jobResultMessage.attachments[0].url

        #print(outputImageUrl)
        embed = discord.Embed(title="Upscaled Image", url=localImageUrl, description="Local mode is on. Please use the link above for upscaled image.")
        embed.set_author(name=author, icon_url=message.author.avatar_url)
        embed.set_image(url=url)
        processingstats = f"Preprocessing took {preprocessingtime}s | Resizing took {resizingtime}s."
        embed.set_footer(text=f"Took {upscaletime}s to upscale {fileid}.\nUpscaling done with {ai_model}.\n{processingstats}")

        await message.channel.send(embed=embed)

      except discord.errors.HTTPException as e:
        baseErrorMessage = f"There was an error sending the output for image id {fileid}."
        if '413 Payload Too Large' in str(e):
          await message.reply(f"{baseErrorMessage} It was probably too large for discord to handle.\n```python\n{e}```")
        else:
          await message.reply(f"{baseErrorMessage}\n```python\n{e}```")

      except Exception as e:
        await message.reply(f"Encountered an error while processing attachment {fileid}\n```python\n{e}```")



    """CLEANING UP"""
    #Already finished the whole job.
    await jobStatusMessage.edit(content=f"\n__Status__: ✅ Job Completed for all {attachmentCount} attachments from message {message.id} at {getloggingtime()}".replace("all 1 attachments", "an attachment"), mention_author=False)
    await message.delete()
    await jobReceiptMessage.delete()
    await asyncio.sleep(2)
    await jobStatusMessage.delete()

    
    
    print(f"IMAGE: Job finished for messageID {message.id}\n")
    CurrentlyProcessingJob = False

  else:
    print(f"MESSAGE: {channel}- {author}: {msg}")
    await client.process_commands(message)


"""BOT COMMANDS"""

@client.command()
async def test(ctx):
  await ctx.send("Hello!")

@client.command()
async def ping(ctx):
  await ctx.send(f'The ping is {round(client.latency * 1000)}ms.')

@client.command()
async def clear(ctx, amount=None):
  if ctx.author.id in ADMINS:
    if amount == "all":
      amount = 100
    elif amount == None:
      amount = 2
    else:
      amount = int(amount)
    await ctx.channel.purge(limit=amount)
    await ctx.send(f"Finished clearing {amount} messages {ctx.message.author.mention}", delete_after=4)
  else:
    await ctx.send('You do not have the permissions to to clear messages using this bot.')

logmsg = f"BOT: Attempting to starting the bot at {getloggingtime()} GMT+7"
print(logmsg)
client.run(TOKEN)