print("Started.")
import json
import os
import re
from threading import Thread
import discord
from colorama import Fore, Style
import random
from huggingface_hub import InferenceClient, login
from PIL import Image
import math
import streamlit as st
import asyncio

HF_TOKEN = os.environ["HF_TOKEN"]

login(HF_TOKEN)

SD = InferenceClient(model="stabilityai/stable-diffusion-2-1")
LLM = InferenceClient(model="openchat/openchat-3.5-0106")


def ec(x, fd="<image>", sd="</image>"):
  matches = re.findall(re.escape(fd) + "(.*?)" + re.escape(sd), x)
  matches = matches if matches else [""]
  return matches


intents = discord.Intents.default()
intents.message_content = True

client = discord.Client(intents=intents)


@client.event
async def on_ready():
  print(f"Logged in as {client.user}")


@client.event
async def on_message(message):
  for user in message.mentions:
    message.content = message.content.replace(f"<@{user.id}>",
                                              f"@{str(user)}")
  try:
    msgchannel = message.channel
    try:
      msgchannel_name = msgchannel.name
      guild = message.guild
      guild_name = guild.name
    except:
      guild_name = "Direct"
      msgchannel_name = message.author.display_name
    s = f":green[{message.author}]: :white[{message.content}] :blue[{msgchannel_name}] :purple[{guild_name}]"
    st.markdown(s)
    if message.author == client.user:
      return
    sysp = """You are Lyre, a helpful discord chatbot. Respond in markdown. Your username is lyre#9828."""
    try:
      os.mkdir(guild_name)
      with open(f"{guild_name}/{msgchannel_name}", "w") as f:
        f.write(
            f"GPT4 Correct system: {sysp}<|end_of_turn|>\nGPT4 Correct {message.author}: {message.content}<|end_of_turn|>"
        )
    except:
      with open(f"{guild_name}/{msgchannel_name}", "a") as f:
        n = "\n"
        f.write(
            f"""GPT4 Correct {message.author.name}: {message.content.strip(n)}<|end_of_turn|>"""
        )
    finally:
      with open(f"{guild_name}/{msgchannel_name}", "r") as f:
        context = f.read()
    context += f"GPT4 Correct Assistant: {message.content}<|end_of_turn|>\n"
    title = ec(context, "<title>", "</title>")[0]
    output = LLM.text_generation(context,
                                 stop_sequences=["<|end_of_turn|>"],
                                 max_new_tokens=4096)
    with open(f"{guild_name}/{msgchannel_name}", "a") as f:
      f.write(f"GPT4 Correct Assistant: {output}<|end_of_turn|>")
    words = output.split()
    embed = discord.Embed(title=title,
                          description=" ".join(words[:i * 10]).replace(
                              f"<title>{title}</title>", ""),
                          color=0x1E81B0)
    embed.set_footer(
        text=
        """Information or code generated by Lyre may not always be correct. Lyre was made by Araeyn."""
    )
    e = await message.reply(embed=embed)
  except Exception as e:
    print(Fore.RED)
    print(e)
    print(Style.RESET_ALL)

token = os.environ["TOKEN"]
client.run(token)