import discord
import requests
import aiohttp
import asyncio
from base64 import b64decode
import json
from sys import exit

try:
    with open("token.secret", 'r') as f:
        discord_client_token = f.read()
except FileNotFoundError:
    print("Cannot locate token.secret please generate a token")
    exit(69)

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

client = discord.Client(intents=intents)

nsfw_enabled = True
steps = 28
seed = -1
url = "https://art.jurydoak.com"

async def gen_image(message):
    # global sess
    
    sess = aiohttp.ClientSession(url)
    
    payload = {
        "prompt": "cute girlhand on cheek eyes closed thinking hard, sfw",
        "styles": ["Bot"],
        "steps": steps,
        "seed": seed,
        "n_iter": 1,
        "height": 1024,
        "negative_prompts": "nsfw, not safe for work, nudity, multiple keyboards",
        "cfg_scale": 12
    }

    settings = {
        "filter_nsfw": not nsfw_enabled,
        "samples_save": True,
    }

    override_payload = {
        "override_settings": settings
    }

    payload.update(override_payload)
    user_prompt = " ".join(message.content.split(' ')[1:])
    payload['prompt'] = user_prompt
    # url = "https://art.jurydoak.com/sdapi/v1"
    
    image_req = await sess.post("/sdapi/v1/txt2img", json=payload)
    x = await image_req.json()
    image_req.close()
    await sess.close()
    
    image = x['images'][0]
    # image = requests.post(f"{url}/txt2img", json=payload).json()['images'][0]
    with open('/tmp/image.png', 'wb') as f:
        f.write(b64decode(image))

    e = discord.Embed()
    upload_file = discord.File("/tmp/image.png", filename="image.png")
    e.set_image(url="attachment://image.png")
    e.title = user_prompt
    await message.channel.send("", file=upload_file, embed=e)


@client.event
async def on_ready():
    
    bot_setup_id = 775489296932405319
    channel = client.get_channel(bot_setup_id)
    await channel.send(f"Awake and ready to degen, NSFW: {nsfw_enabled}, STEPS: {steps}")
    await client.change_presence(activity=discord.Game(name="Thinking degenerate thoughts..."))

    # text_channel_list = []
    # for server in discord.Client.servers:
    #     for channel in server.channels:
    #         if channel.type == "text":
    #             text_channel_list.append(channel)

    # print(text_channel_list)
    print(f'We have logged in as {client.user}')
    

@client.event
async def on_message(message):
    global nsfw_enabled
    if message.author == client.user:
        return

    if message.content.startswith('$hello'):
        print(message.content)
        await message.channel.send('Hopefully queued the image')
        await gen_image(message)

    if message.content.startswith('$degen'):
        await message.channel.send('Disabling kink mode' if nsfw_enabled else 'Enabling kink mode')
        if nsfw_enabled:
            nsfw_enabled = False
        else:
            nsfw_enabled = True

    if message.content.startswith("$steps"):
        global steps
        try:
            num = int(message.content.split(" ")[1])
            steps = num
            await message.channel.send(f"Setting steps to {steps}")
        except:
            pass
        

    if message.content.startswith("$seed"):
        global seed
        try:
            num = int(message.content.split(" ")[1])
            seed = num
            await message.channel.send(f"Setting seed to {seed}")
        except:
            pass

    if message.content.startswith("$host"):
        global url
        try:
            host = message.content.split(" ")[1]
            url = host
            await message.channel.send(f"Setting host to {host}")
        except:
            pass

client.run(discord_client_token)