from discord import option
from discord.ext import commands
import core_utils
import discord
import os


class StableDiffusion(commands.Cog):
    def __init__(self, bot):
        self.bot = bot

    @commands.has_role("stable-diffuser")
    @commands.slash_command(
        guild_ids=core_utils.my_guilds,
        name="sd",
        description="Pass a prompt and optional negative prompt to stable diffusion",
    )
    @option(
        "positive_prompt",
        description="The positive prompt to pass to stable diffusion",
        input_type="str",
        requred=True,
    )
    @option(
        "negative_prompt",
        description="An optional set of negatives you want to pass to stable diffusion",
        required=False,
    )
    async def sd(
        self, ctx: commands.Context, positive_prompt: str, negative_prompt: str
    ):
        if ctx.channel.name == "stable-diffusion" or ctx.channel.name == "bot-testing":
            import socket
            import stable_diffusion

            port = "7860"
            ip = "192.168.1.80"
            steps = 20

            # Send my requests to my gaming computer with the 3080 (if its up)
            if ctx.author.id == core_utils.my_id:
                ip = "192.168.1.19"
                steps = 90
                s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                s.settimeout(1)
                try:
                    s.connect((ip, int(port)))
                except:
                    ip = "192.168.1.80"
                    steps = 60
            try:
                await ctx.defer()
                original_message = await ctx.followup.send(
                    "Please be patient, I'm generating your image"
                )
                file_path, time_taken = await stable_diffusion.generate_image(
                    ip=ip,
                    port=port,
                    positives=positive_prompt,
                    negatives=negative_prompt or None,
                    steps=steps,
                )
                await original_message.edit(
                    content=time_taken,
                    file=discord.File(
                        file_path,
                        filename="unknown.png",
                    ),
                )
                os.remove(file_path)
            except Exception as e:
                await ctx.respond(
                    "Stable diffusion isnt running right now, sorry.\n%s" % e,
                )
        else:
            await ctx.respond("You can only do that in the stable-diffusion channel")


def setup(bot):
    bot.add_cog(StableDiffusion(bot))