import io
import os
import sys
import random
import zipfile
import logging
import discord
from discord import app_commands
from discord.ext import commands
import requests

NOVELAI_API_TOKEN = os.getenv("NOVELAI_API_TOKEN")
BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
ALLOWED_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID", "0"))  # falls nicht gesetzt → 0

logging.basicConfig(
    stream=sys.stdout,
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)

logger = logging.getLogger(__name__)

if not all([NOVELAI_API_TOKEN, BOT_TOKEN, ALLOWED_CHANNEL_ID]):
    raise ValueError("Environment-Variablen NOVELAI_API_TOKEN, DISCORD_BOT_TOKEN und DISCORD_CHANNEL_ID müssen gesetzt sein.")

intents = discord.Intents.default()
bot = commands.Bot(command_prefix=commands.when_mentioned, intents=intents)

@bot.event
async def on_ready():
    print(f"Bot started as {bot.user}")
    try:
        synced = await bot.tree.sync()
        logging.info(f"{len(synced)} Slash-Commands synchronized.")
    except Exception as e:
        logging.error(e)

@bot.tree.command(name="generate", description="Generate image with NovelAI v4 Full")
@app_commands.describe(
    prompt="What should be generated?",
    undesired_prompt="What should be avoided? (optional)",
    orientation="portrait or landscape (Standard: portrait)",
    seed="Optional seed (integer). If empty, a random one will be generated"
)
async def generate(
    interaction: discord.Interaction,
    prompt: str,
    undesired_prompt: str = "",
    orientation: str = "portrait",
    seed: int = None
):
    
    default_negative = (
        "blurry, lowres, error, film grain, scan artifacts, worst quality, bad quality, jpeg artifacts, "
        "very displeasing, chromatic aberration, multiple views, logo, too many watermarks, white blank page, "
        "blank page, watermarks, watermark, text, "
    )
    
    if interaction.channel.id != ALLOWED_CHANNEL_ID:
        await interaction.response.send_message(
            "This command isn't allowed here.",
            ephemeral=True
        )
        return
    
    MAX_PROMPT_LENGTH = 500

    if len(prompt) > MAX_PROMPT_LENGTH:
        await interaction.followup.send(f"Prompt too long! Maximum {MAX_PROMPT_LENGTH} chars allowed.")
        return

    if len(default_negative + undesired_prompt.strip()) > MAX_PROMPT_LENGTH:
        await interaction.followup.send(f"Negative Prompt too long! Maximum {MAX_PROMPT_LENGTH} chars allowed.")
        return

    await interaction.response.defer(thinking=True)
    logger.info(
        "User %s (%s, %s) requested image",
        interaction.user.display_name,
        interaction.user.name,
        interaction.user.id
    )

    # Auflösung wählen
    if orientation.lower() == "landscape":
        width, height = 1216, 832
    else:
        width, height = 832, 1216

    if seed is None:
        seed = random.randint(0, 2**32 - 1)

    # API-Payload zusammenbauen
    payload = {
        "input": prompt,
        "model": "nai-diffusion-4-full",
        "action": "generate",
        "parameters": {
            "params_version": 3,
            "width": width,
            "height": height,
            "scale": 3.7,
            "sampler": "k_euler_ancestral",
            "steps": 28,
            "n_samples": 1,
            "ucPreset": 0,
            "qualityToggle": True,
            "autoSmea": False,
            "dynamic_thresholding": False,
            "controlnet_strength": 1,
            "legacy": False,
            "add_original_image": True,
            "cfg_rescale": 0,
            "noise_schedule": "karras",
            "legacy_v3_extend": False,
            "skip_cfg_above_sigma": None,
            "use_coords": False,
            "legacy_uc": False,
            "normalize_reference_strength_multiple": True,
            "seed": seed,
            "characterPrompts": [],
            "v4_prompt": {
                "caption": {
                    "base_caption": prompt,
                    "char_captions": []
                },
                "use_coords": False,
                "use_order": True
            },
            "v4_negative_prompt": {
                "caption": {
                    "base_caption": default_negative + undesired_prompt.strip(),
                    "char_captions": []
                },
                "legacy_uc": False
            },
            "negative_prompt": default_negative + undesired_prompt.strip(),
            "deliberate_euler_ancestral_bug": False,
            "prefer_brownian": True
        }
    }

    headers = {
        "Authorization": f"Bearer {NOVELAI_API_TOKEN}",
        "Content-Type": "application/json"
    }

    response = requests.post(
        "https://image.novelai.net/ai/generate-image",
        json=payload,
        headers=headers,
        timeout=120
    )

    if response.status_code == 200:
        response_bytes = response.content

        param_text = f"""```Prompt: {prompt}
Undesired prompt: {default_negative + undesired_prompt.strip()}
Seed: {seed}
Resolution: {width}x{height}
Sampler: k_euler_ancestral
Steps: 28
Scale: 3.7
Model: nai-diffusion-4-full```"""

        # Prüfen ob ZIP (PK am Anfang)
        if response_bytes[:2] == b'PK':
            # ZIP-Archiv entpacken
            with zipfile.ZipFile(io.BytesIO(response_bytes)) as zip_file:
                namelist = zip_file.namelist()
                # Erstes PNG suchen
                image_name = next((name for name in namelist if name.endswith(".png")), None)

                if image_name:
                    image_bytes = zip_file.read(image_name)
                    filename = "novelai.png"
                    file = discord.File(io.BytesIO(image_bytes), filename=filename)
                    await interaction.followup.send(content=param_text, file=file)
                else:
                    await interaction.followup.send(
                        "Found no PNG inside archive"
                    )

        elif response_bytes[:4] == b'\x89PNG':
            # Direkte PNG-Rückgabe
            filename = "novelai.png"
            file = discord.File(io.BytesIO(response_bytes), filename=filename)
            await interaction.followup.send(content=param_text, file=file)

        else:
            await interaction.followup.send(
                "API didn't send any file"
            )

    else:
        try:
            error_data = response.json()
        except Exception:
            error_data = {"error": response.text}

        error_message = f"Error {response.status_code} at API-request.\n"
        for key, value in error_data.items():
            error_message += f"**{key}**: {value}\n"

        await interaction.followup.send(error_message)

bot.run(BOT_TOKEN)