import discord
from discord.ext import commands
import aiohttp
from base64 import b64decode
from sys import exit
from modules import diffuseapi

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

api = diffuseapi.DiffuseAPI("http://localhost:7860", ["default"], True, 28)

# client = discord.Client(intents=intents)

class Settings:
    def __init__(self, nsfw_enabled=True, num_steps=28, ai_seed=-1, url="https://art.jurydoak.com", styles=["Bot"]):

        self.nsfw_enabled = nsfw_enabled
        self.num_steps = num_steps
        self.ai_seed = ai_seed
        self.url = url
        self.styles = styles


main_settings = Settings()
activity = discord.Activity(type=discord.ActivityType.listening, name="!help")


bot = commands.Bot(intents=intents, command_prefix="!", activity=activity, help_command=commands.DefaultHelpCommand())

async def send_request(ctx, endpoint, payload):
    global num_steps
    global ai_seed
    sess = aiohttp.ClientSession(main_settings.url)
    alive = await sess.head('/')

    if alive.status != 200:
        await ctx.reply("I'm sorry but it appears my brain is offline")
        alive.close()
        await sess.close()
        return {
            'images': [None]
        }

    print(payload)
    request  = await sess.post(endpoint, json=payload)
    try:
        req_json = await request.json()
    except:
        print(await request.text())
        return {
            'images': [None]
        }
    request.close()
    await sess.close()
    return req_json

@bot.command()
async def prompt(ctx, *args):
    '''Generate an image with the provided prompt'''
    prompt = " ".join(args)
    await ctx.reply("Generating your image boo")
    payload = {
        "prompt": prompt,
        "styles": main_settings.styles,
        "steps": main_settings.num_steps,
        "seed": main_settings.ai_seed,
        "n_iter": 1,
        "height": 1024,
        "negative_prompts": "nsfw, not safe for work, nudity, multiple keyboards",
        "cfg_scale": 12
    }

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

    override_payload = {
        "override_settings": settings
    }

    payload.update(override_payload)
    image_data = await send_request(ctx, endpoint="/sdapi/v1/txt2img", payload=payload)

    if image_data['images'][0] is None:
        await ctx.reply("Something went wrong, please report this to the admin so it can be ignored")
        return

    with open('/tmp/image.png', 'wb') as f:
        f.write(b64decode(image_data['images'][0]))

    embed = discord.Embed()
    upload_file = discord.File("/tmp/image.png", filename="image.png")
    embed.set_image(url="attachment://image.png")
    # embed.title = prompt

    await ctx.reply("", file=upload_file, embed=embed)


@bot.command()
async def seed(ctx, arg):
    '''Set the seed for the image generation'''
    try:
        arg = int(arg)
    except:
        pass

    # global ai_seed
    main_settings.ai_seed = arg

    await ctx.reply(f"I have updated the seed to {main_settings.ai_seed} for you my master.")

@bot.command()
async def steps(ctx, arg):
    '''Set how many steps the AI will run (max 50)'''
    try:
        arg = int(arg)
    except:
        pass

    if arg > 50:
        await ctx.reply("I'm sorry Dave, I can't do that")
        return

    # global num_steps
    main_settings.num_steps = arg

    await ctx.reply(f"I have updated the steps to {main_settings.num_steps} for you my master.")

@bot.command()
async def settings(ctx):
    '''See the currently configured settings (BROKEN)'''
    global ai_seed, steps
    settings = f"""
    ```
    seed: {main_settings.ai_seed}
    steps: {main_settings.num_steps}
    ```
    """
    await ctx.message.channel.send(settings)

@bot.command()
async def test(ctx):
    '''Test function, currently changes the URL and bot settings'''
    if main_settings.url != "http://localhost:7860":
        main_settings.url = "http://localhost:7860"
        main_settings.styles = ["default"]
        await ctx.reply("Set to fastboi")
    else:
        main_settings.url = "https://art.jurydoak.com"
        main_settings.styles = ["Bot"]
        await ctx.reply("Set to main api")

@bot.command()
async def test2(ctx):
    image_data = await api.generate_image(prompt="cute bot doing bot things")
    if image_data is None:
        return
    with open('/tmp/image.png', 'wb') as f:
        f.write(b64decode(image_data))

    embed = discord.Embed()
    upload_file = discord.File("/tmp/image.png", filename="image.png")
    embed.set_image(url="attachment://image.png")
    # embed.title = prompt

    await ctx.reply("", file=upload_file, embed=embed)

bot.run(discord_client_token)