#!/usr/bin/env python3
"""
Plugin for generating images using self-hosted Stable Diffusion and sending them to a Matrix chat room.

Now fully asynchronous (uses aiohttp). All original parameters and help text are preserved.
"""

import aiohttp
import base64
import tempfile
import os
import argparse
import simplematrixbotlib as botlib

async def handle_command(room, message, bot, prefix, config):
    match = botlib.MessageMatch(room, message, bot, prefix)
    if not (match.prefix() and match.command("sd")):
        return

    # Check if API is reachable
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get("http://127.0.0.1:7860/docs", timeout=3) as resp:
                if resp.status != 200:
                    await bot.api.send_text_message(room.room_id, "Stable Diffusion API is not running!")
                    return
    except Exception:
        await bot.api.send_text_message(room.room_id, "Could not reach Stable Diffusion API!")
        return

    try:
        parser = argparse.ArgumentParser(description='Generate images using self-hosted Stable Diffusion')
        parser.add_argument('--steps', type=int, default=4, help='Number of steps, default=4')
        parser.add_argument('--cfg', type=int, default=2, help='CFG scale, default=2')
        parser.add_argument('--h', type=int, default=512, help='Height of the image, default=512')
        parser.add_argument('--w', type=int, default=512, help='Width of the image, default=512')
        parser.add_argument('--neg', type=str, nargs='+',
                            default=['((((ugly)))), (((duplicate))), ((morbid)), ((mutilated)), out of frame, extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), ((ugly)), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))), out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck)))'],
                            help='Negative prompt')
        parser.add_argument('--sampler', type=str, nargs='*', default=['DPM++', 'SDE Karras'],
                            help='Sampler name, default=DPM++ SDE')
        parser.add_argument('--seed', type=int, default=None,
                            help='Seed for deterministic generation (omit for random)')
        parser.add_argument('prompt', type=str, nargs='*', help='Prompt for the image')

        args = parser.parse_args(message.body.split()[1:])  # skip command prefix

        if not args.prompt:
            raise argparse.ArgumentError(None, "Prompt is required.")

        prompt = ' '.join(args.prompt)
        sampler_name = ' '.join(args.sampler)
        neg_prompt = ' '.join(args.neg)

        payload = {
            "prompt": prompt,
            "steps": args.steps,
            "negative_prompt": neg_prompt,
            "sampler_name": sampler_name,
            "cfg_scale": args.cfg,
            "width": args.w,
            "height": args.h,
        }
        if args.seed is not None:
            payload["seed"] = args.seed

        async with aiohttp.ClientSession() as session:
            async with session.post("http://127.0.0.1:7860/sdapi/v1/txt2img", json=payload, timeout=600) as response:
                response.raise_for_status()
                r = await response.json()

        # Save and send image
        image_data = base64.b64decode(r['images'][0])
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
            filename = temp_file.name
            temp_file.write(image_data)

        await bot.api.send_image_message(room_id=room.room_id, image_filepath=filename)

        # Optional info message (commented out to avoid spam, but can be enabled)
        # neg_prompt_clean = neg_prompt.replace(" ", "")
        # seed_info = f"<br><strong>Seed:</strong> {args.seed}" if args.seed is not None else ""
        # info_msg = f"<details><summary>🔍 Image Info</summary><strong>Prompt:</strong> {prompt[:100]}<br>...</details>"
        # await bot.api.send_markdown_message(room.room_id, info_msg)

        # Clean up temp file
        os.remove(filename)

    except argparse.ArgumentError as e:
        await bot.api.send_text_message(room.room_id, f"Argument Error: {e}")
        await bot.api.send_markdown_message(room.room_id, "<details><summary>Stable Diffusion Help</summary>" + print_help() + "</details>")
    except Exception as e:
        await bot.api.send_text_message(room.room_id, f"Error processing the command: {str(e)}")

def print_help():
    """
    Generates the full help text for the 'sd' command, including LORA list.
    """
    return """
  <p>Generate images using self-hosted Stable Diffusion</p>

  <p>Positional arguments:</p>
  <ul>
    <li>prompt - Prompt for the image</li>
  </ul>

  <p>Optional arguments:</p>
  <ul>
    <li>--steps STEPS - Number of steps, default=4</li>
    <li>--cfg CFG - CFG scale, default=2</li>
    <li>--h H - Height of the image, default=512</li>
    <li>--w W - Width of the image, default=512</li>
    <li>--neg NEG - Negative prompt, default=((((ugly)))), (((duplicate))), ((morbid)), ((mutilated)), out of frame, extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), ((ugly)), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))), out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck)))</li>
    <li>--sampler SAMPLER - Sampler name, default=DPM++ SDE</li>
    <li>--seed SEED - Seed for deterministic generation (omit for random)</li>
  </ul>

  <p>LORA List:</p>
  <ul>
    <li>&lt;psychedelicai-SDXL&gt;</li>
    <li>&lt;ral-frctlgmtry-sdxl&gt;</li>
    <li>&lt;SDXL-PsyAI-v4&gt;</li>
    <li>&lt;al3xxl&gt;</li>
  </ul>

  <p>Load LORAs like this:</p>
  <ul>
    <li>&lt;lora:SDXL-PsyAI-v4:1&gt; PsyAI</li>
    <li>&lt;lora:psychedelicai-SDXL:1&gt; Psychedelic</li>
    <li>&lt;ral-frctlgmtry-sdxl&gt; ral-frctlgmtry</li>
    <li>&lt;lora:al3xxl:1&gt; alexpainting, alexhuman, alexentity, alexthirdeye, alexforeheads, alexgalactic, spiraling, alexmirror, alexangel, alexkissing, alexthirdeye2, alexthirdeye3, alexhofmann, wearing, glasses, alexgalactic3, alexthirdeye4, alexhuman3, alexgalactic2, alexhuman2, alexbeing2, alexfractal2, alexfractal3</li>
  </ul>
"""

# ---------------------------------------------------------------------------
# Plugin Metadata
# ---------------------------------------------------------------------------
__version__ = "1.1.2"
__author__ = "Funguy Bot"
__description__ = "Stable Diffusion image generation (LORA support)"
__help__ = """
<details>
<summary><strong>!sd</strong> – Generate images via Stable Diffusion</summary>
<p><code>!sd [options] &lt;prompt&gt;</code></p>
<ul>
<li><code>--steps N</code> – Sampling steps (default 4)</li>
<li><code>--cfg scale</code> – CFG scale (default 2)</li>
<li><code>--h H --w W</code> – Image dimensions (default 512)</li>
<li><code>--neg &lt;negative prompt&gt;</code></li>
<li><code>--sampler SAMPLER</code> – Sampler name (default DPM++ SDE)</li>
<li><code>--seed SEED</code> – Deterministic seed (optional)</li>
</ul>
<p>LORAs: <code>&lt;lora:filename:weight&gt;</code></p>
<p>Requires a locally running Stable Diffusion API.</p>
</details>
"""