# ruff: noqa: E712
import asyncio
import json
import logging
import re
from asyncio import Lock
from datetime import UTC, datetime, timedelta
from io import BytesIO
from pathlib import Path
from random import choice as random_choice
from traceback import format_exc
from typing import Any

import jinja2 as j2
from disnake import (
    ApplicationCommandInteraction,
    Colour,
    DMChannel,
    Embed,
    File,
    Forbidden,
    GroupChannel,
    InteractionContextTypes,
    Member,
    Message,
    MessageCommandInteraction,
    MessageInteraction,
    NotFound,
    TextChannel,
    Thread,
    User,
)
from disnake.ext import commands, tasks
from fastapi import HTTPException
from Levenshtein import distance as lev_distance
from sqlalchemy import select
from sqlalchemy.orm import load_only
from transformers import AutoTokenizer, BatchEncoding

import logsnake
from ai.constants import AI_DATA_DIR, AI_LOG_DIR, AI_LOG_FORMAT
from ai.eyes import DiscoEyes
from ai.imagen import Imagen
from ai.settings import (
    BotMode,
    BotParameters,
    GuildSettings,
    GuildSettingsList,
    LMApiConfig,
    Prompt,
    get_ai_settings,
)
from ai.tokenizers import PreTrainedTokenizerBase, extract_tokenizer
from ai.types import AiResponse, ModelInfo, NamedSnowflake
from ai.ui import AiParam, AiStatusEmbed, convert_param, set_choices
from ai.utils import (
    LruDict,
    MentionMixin,
    cleanup_thoughts,
    get_message_author_name,
    get_prompt_datetime,
    member_in_any_role,
)
from ai.utils.regex import (
    re_angle_bracket,
    re_consecutive_newline,
    re_detect_url,
    re_mention,
    re_nonword,
    re_start_expression,
    re_unescape_md,
)
from ai.web import GradioUi
from cogs.privacy import PrivacyEmbed, PrivacyView, get_policy_text
from db import DiscordUser, Session, SessionType
from db.ai import AiResponseLog
from disco_snake import checks
from disco_snake.blacklist import Blacklist
from disco_snake.bot import DiscoSnake
from shimeji.chatbot import ChatBot
from shimeji.model_provider import OobaGenRequest, OobaModel

COG_UID = "ai"

ai_logger = logsnake.setup_logger(
    name=COG_UID,
    level=logging.DEBUG,
    isRootLogger=False,
    formatter=logsnake.LogFormatter(fmt=AI_LOG_FORMAT, datefmt="%Y-%m-%d %H:%M:%S"),
    logfile=AI_LOG_DIR.joinpath(f"{COG_UID}.log"),
    fileLoglevel=logging.DEBUG,
    maxBytes=1 * (2**20),
    backupCount=1,
    propagate=True,
)
logger = logging.getLogger(__name__)


class Ai(MentionMixin, commands.Cog, name=COG_UID):
    def __init__(self, bot: DiscoSnake, *args, **kwargs):
        self.bot: DiscoSnake = bot
        self.timezone = self.bot.timezone
        self.last_response = datetime.now(UTC) - timedelta(minutes=10)

        # init the MentionMixin cache and other stuff
        super().__init__(*args, **kwargs)

        # Load config file
        self.config = get_ai_settings()
        self.lm_lock = Lock()  # used to stop multiple responses from happening at once

        # Parse config file
        self.provider_config: LMApiConfig = self.config.model_provider

        # Load config params up into top level properties
        self.params: BotParameters = self.config.params

        self.prompt: Prompt = self.config.prompt
        self.max_retries = self.params.max_retries

        self.idle_enable: bool = self.params.idle_enable
        self.idle_channels: list[int] = self.params.get_idle_channels()

        self.ctxbreak = self.params.ctxbreak
        self.guilds: GuildSettingsList = self.params.guilds
        self.dm_user_ids: list[int] = self.params.dm_user_ids
        self.dm_user_ids.extend(self.bot.config.admin_ids)
        self.tos_reject_ids: set[int] = set()

        # we will populate these later during async init
        self.model_provider: OobaModel = None  # type: ignore
        self.model_provider_name: str = self.provider_config.provider
        self.model_info: ModelInfo | None = None
        self.chatbot: ChatBot = None  # type: ignore
        self.logging_channel: TextChannel | DMChannel | None = None
        self.tokenizer_type = self.provider_config.modeltype
        self.tokenizer: PreTrainedTokenizerBase = None  # type: ignore
        self.bad_words: list[str] = self.config.bad_words

        # database client (async init)
        self.db_client: SessionType = Session
        self.blacklist: Blacklist = Blacklist()

        ## addon modules
        self.imagen = Imagen(cog=self)  # selfietron
        self.eyes = DiscoEyes(cog=self)  # image caption engine
        self.webui = GradioUi(cog=self, config=self.config.gradio)  # gradio ui

        # debugging related
        self.debug: bool = self.params.debug
        self.debug_dir: Path = AI_LOG_DIR.joinpath("ai")
        self.debug_dir.mkdir(parents=True, exist_ok=True)
        self.debug_dir.joinpath("dm").mkdir(parents=True, exist_ok=True)
        self._last_debug_log: dict[str, Any] = {}
        self.logging_channel_id: int = self.params.logging_channel_id  # type: ignore

        # cache dicts and props
        self._trigger_cache: LruDict = LruDict(max_size=100)
        self._n_prompt_tokens: int = None  # type: ignore
        # other cached vars
        self._chat_template: j2.Template | None = None

    # Getters for config object sub-properties
    @property
    def name(self) -> str:
        return self.config.name.split(" ")[0]

    @property
    def lm_request(self) -> OobaGenRequest:
        return self.provider_config.gensettings

    @property
    def context_size(self) -> int:
        if self.params.context_size < 0:
            return (
                self.lm_request.truncation_length
                - (
                    self.lm_request.max_tokens
                    if self.lm_request.max_tokens and self.lm_request.max_tokens > 0
                    else (self.lm_request.truncation_length // 2)
                )
                - 32
            )
        return self.params.context_size

    @property
    def siblings(self) -> list[NamedSnowflake]:
        return [x for x in self.params.siblings]

    @property
    def sibling_ids(self) -> list[int]:
        return self.params.siblings.ids

    @property
    def n_prompt_tokens(self):
        if self.prompt.disco_mode:
            # no prompt tokens in disco mode
            self._n_prompt_tokens = 0
        if self._n_prompt_tokens is None:
            _, _, prompt_tokens = self._render_prompt_context()
            self._n_prompt_tokens = prompt_tokens
        return self._n_prompt_tokens

    @property
    def chat_template(self) -> j2.Template | None:
        return self.tokenizer.get_chat_template()

    def _count_tokens(self, text: str) -> int:
        """
        Count tokens for a rendered string using the configured tokenizer.
        """
        if self.tokenizer is None:
            raise ValueError("Tokenizer not initialized")
        encoded: BatchEncoding = self.tokenizer.batch_encode_plus([text], return_length=True)
        return encoded.get("length", [0]).pop(0)

    def _normalize_chat_messages(
        self, conversation: Any, *, default_role: str = "user"
    ) -> list[dict[str, str]]:
        """
        Coerce prompt or conversation data into a flat list of chat messages.
        """
        if conversation is None:
            return []
        if hasattr(conversation, "model_dump"):
            return self._normalize_chat_messages(conversation.model_dump(), default_role=default_role)
        if isinstance(conversation, dict):
            role = conversation.get("role", default_role).strip()
            content = conversation.get("content", "")
            return [{"role": role, "content": content.strip()}]
        if isinstance(conversation, list):
            messages: list[dict[str, str]] = []
            for entry in conversation:
                messages.extend(self._normalize_chat_messages(entry, default_role=default_role))
            return messages
        if isinstance(conversation, str):
            return [{"role": default_role, "content": conversation.strip()}]
        raise TypeError(f"Unsupported conversation type: {type(conversation)}")

    def apply_chat_template(
        self,
        conversation: list[dict[str, str]] | list[list[dict[str, str]]],
        **kwargs,
    ) -> str:
        if not self.tokenizer.chat_template:
            raise ValueError("No chat template defined")
        try:
            return self.tokenizer.apply_chat_template(
                conversation=conversation,
                add_generation_prompt=self.prompt.add_generation_prompt,
                enable_thinking=self.prompt.thinking,
                thinking=self.prompt.thinking,
                keep_cots=self.prompt.thinking,
                tokenize=False,
                **kwargs,
            )
        except (j2.exceptions.TemplateError, AttributeError) as e:
            logger.exception(f"Error rendering chat template: {e}")
            raise e

    def _render_prompt_context(
        self, ctx: MessageInteraction | Message | None = None
    ) -> tuple[list[dict[str, str]], str, int]:
        """
        Render the configured prompt through the chat template and measure its token length.
        """
        prompt_data = self.get_prompt(ctx=ctx)
        prompt_messages = self._normalize_chat_messages(prompt_data, default_role="system")
        if len(prompt_messages) == 0:
            return [], "", 0
        prompt_text = self.apply_chat_template(prompt_messages)
        prompt_tokens = self._count_tokens(prompt_text)
        return prompt_messages, prompt_text, prompt_tokens

    def _update_gensettings_from_model_info(self) -> None:
        if self.model_info:
            logger.info("Got model info from API server, updating length settings")
            if self.model_info.parameters.max_seq_len < 2048:
                raise ValueError("Model max_seq_len is too small (<2048 tokens), good luck...")
            self.lm_request.truncation_length = self.model_info.parameters.max_seq_len

            if (not self.lm_request.max_tokens) or self.lm_request.max_tokens <= 0:
                # max_tokens is not set or is explicitly set to auto, so set it to half of truncation length
                self.lm_request.max_tokens = self.lm_request.truncation_length // 2
            else:
                # max_tokens was explicitly set, ensure it does not exceed truncation length/2
                self.lm_request.max_tokens = min(
                    self.lm_request.max_tokens,
                    self.lm_request.truncation_length // 2,
                )
        else:
            # set max_tokens to -1 to indicate "however many fit"
            self.lm_request.max_tokens = -1

        logger.info("Dynamic generation settings:")
        logger.info(f"  truncation_length: {self.lm_request.truncation_length}")
        logger.info(f"  max_tokens: {self.lm_request.max_tokens}")
        logger.info(f"  context_size: {self.context_size}")

    async def cog_load(self) -> None:
        logger.info("AI engine initializing, please wait...")

        tokenizer_dir = AI_DATA_DIR.joinpath(f"tokenizers/{self.tokenizer_type}")
        tokenizer_dir.mkdir(parents=True, exist_ok=True)
        if not tokenizer_dir.joinpath("tokenizer.json").is_file():
            extract_tokenizer(name=self.tokenizer_type, target_dir=tokenizer_dir)

        self.tokenizer = AutoTokenizer.from_pretrained(
            pretrained_model_name_or_path=tokenizer_dir,
            local_files_only=True,
            trust_remote_code=True,
            use_fast=True,
            model_max_length=self.lm_request.truncation_length,
            chat_template=self.prompt.template_str,
        )
        logger.info("Tokenizer initialized.")

        match self.model_provider_name:
            case "ooba":
                self.model_provider = OobaModel(
                    endpoint_url=self.provider_config.endpoint,
                    default_args=self.lm_request,
                    api_key=self.provider_config.api_key,
                    auth_header=self.provider_config.auth_header,
                    tokenizer=self.tokenizer,
                    api_v2=self.provider_config.api_v2,
                    debug=self.params.debug,
                )
                if model_info := self.model_provider.model_info:
                    self.model_info = model_info
                try:
                    self._update_gensettings_from_model_info()
                except Exception:
                    logger.exception("Error updating generation settings from model info")
            case _:
                raise ValueError(f"Unknown model provider type: {self.model_provider_name}")
        logger.info("Model Provider initialized.")

        self.chatbot = ChatBot(
            name=self.name,
            model_provider=self.model_provider,
            preprocessors=[],
            postprocessors=[],
        )
        logger.info("ChatBot object initialized.")

        _logchannel = self.bot.get_channel(self.logging_channel_id)
        self.logging_channel = _logchannel if isinstance(_logchannel, (TextChannel, DMChannel)) else None

        logger.info("Starting WebUI (if enabled)...")
        await self.webui.launch()

        logger.debug("Known sibling bots:")
        for sibling in self.siblings:
            logger.debug(f" - {sibling.id=} {sibling.name=}")

        if self.eyes is not None:
            logger.info("Starting DiscoEyes...")
            await self.eyes.start()

        if not self.update_ignored.is_running():  # type: ignore
            logger.debug("Starting update_ignored task...")
            self.update_ignored.start()

        if self.idle_enable is True:
            logger.info("Starting idle messaging loop")
            self.idle_loop.start()

        if not self.log_archive_loop.is_running():
            logger.info("Starting log archive loop")
            self.log_archive_loop.start()

        _ = self.chat_template  # force template compilation

        logger.info("AI engine initialized... probably? yolo!")

    def cog_unload(self) -> None:
        if self.eyes is not None:
            logger.info("Shutting down DiscoEyes...")
            self.eyes.shutdown()

        if self.webui is not None:
            logger.info("Shutting down WebUI...")
            self.webui.shutdown()

        logger.info("AI engine unloaded.")

    @commands.Cog.listener("on_ready")
    async def on_ready(self):
        if self.logging_channel is None and self.logging_channel_id is not None:
            self.logging_channel = self.bot.get_channel(self.logging_channel_id)  # type: ignore
            logger.info(f"Logging channel: {self.logging_channel}")

        logger.debug("DMs are enabled for the following users:")
        for user_id in self.dm_user_ids:
            try:
                user_obj = self.bot.get_user(user_id)
                logger.debug(f"  - {user_obj} ({user_id})")
            except Exception:
                logger.exception(f"Failed to get user object for {user_id}")
        # done
        logger.info("AI engine is ready, uwu~")

    @commands.Cog.listener("on_message")
    async def on_message(self, message: Message):
        # ignore messages from ourselves, blacklisted users, ToS rejecters, and bots (if the last trigger was a bot)
        if message.author.id == self.bot.user.id:
            return  # ignore messages from self
        if message.author.id in self.blacklist:
            return  # ignore messages from blacklisted users
        if message.author.id in self.tos_reject_ids:
            return  # ignore messages from users who have rejected the ToS
        if message.author.bot is True and self.last_trigger_was_bot(message) is True:
            return  # ignore consecutive messages from bots

        if message.content.startswith(", "):
            return  # ignore messages that start with a comma-space
        if len(message.content.strip()) == 1:
            return  # ignore messages with only a single character
        if self.config.skip_codeblocks and "```" in message.content:
            return  # ignore messages with code blocks

        direct_message = isinstance(message.channel, DMChannel)
        guild_settings = self.params.get_guild_settings(message.guild.id) if message.guild else None

        # Ignore messages from unapproved users in DMs/groups
        if direct_message is True:
            if message.author.id not in self.dm_user_ids:
                logger.debug(f"DM from {message.author} (ID {message.author.id}) ignored")
                return
        # Ignore threads, group DMs, and messages with no guild (e.g. system messages)
        elif isinstance(message.channel, (Thread, GroupChannel)) or message.guild is None:
            return

        if not direct_message:
            if guild_settings is None:
                logger.debug(f"Got message in unknown guild {message.guild=}")
                return  # ignore messages from guilds that don't have settings
            if guild_settings.enabled is False:
                # logger.debug(f"Got message in disabled guild {message.guild=}")
                return  # ignore messages from guilds that have the bot disabled
            if guild_settings.channel_enabled(message.channel.id) is False:
                # logger.debug(f"Got message in disabled channel {message.channel=}")
                return  # ignore messages from channels that have the bot disabled
            if message.guild is not None:
                if message.channel.permissions_for(message.guild.me).send_messages is False:
                    # channel is enabled but we don't have permission to respond
                    logger.info(f"Got message in {message.channel} but don't have permission to respond.")
                    return

        if message.author.bot is True:
            # bots can't and don't need to accept the ToS
            tos_accepted = True
        else:
            # Check if the user has accepted the ToS
            tos_accepted = await self.user_accepted_tos(message.author)

        if tos_accepted is False:
            return  # ignore messages from users who have rejected the ToS

        # Now that we've filtered out messages we don't care about, let's get the message content
        content = self.get_msg_content_clean(message)
        if content is None:
            return  # ignore empty messages

        trigger: str | None = None  # response trigger reason (tfw no StrEnum in 3.10)
        conversation = None  # message's conversation context
        append = None  # optional masked message to append to the response

        if self.lm_lock.locked():
            logger.info("Already processing a message, skipping this one to avoid overlap.")
            return

        try:
            async with self.lm_lock:
                if direct_message:
                    logger.debug(f"DM from {message.author}: {content}")
                    trigger = "DM"

                elif self.check_mention(message):
                    logger.debug(f"Mentioned in {message.channel}: {content}")
                    trigger = "mention"

                if trigger is not None:
                    await message.channel.trigger_typing()
                    if conversation is None:
                        conversation = await self.get_message_context(message, guild_settings=guild_settings)

                    # save whether we're responding to a bot or not to prevent loops
                    self._trigger_cache.update({message.channel.id: message.author.bot is True})

                    # send ToS if it's not been accepted or rejected yet
                    if tos_accepted is None:
                        append = "Please check your DMs for a privacy notice!"
                        await self.check_send_tos(message)

                    # actually respond
                    await self.generate_response(conversation, message, trigger, append=append)
                    return

        except Exception as e:
            logger.exception(e)
            if isinstance(e, HTTPException):
                return  # ignore HTTP errors
            msg_timestamp = message.created_at.astimezone(tz=self.bot.timezone)
            exc_class = e.__class__.__name__
            exc_desc = str(f"**``{exc_class}``**\n```{format_exc()}```")
            err_filename = f"error-{msg_timestamp.strftime('%Y-%m-%d-%H:%M:%S%z')}.txt".replace(" ", "")
            self.debug_dir.joinpath(err_filename).write_text(exc_desc)

    # retrieve the LM prompt and inject name, time, etc.
    def get_prompt(
        self,
        ctx: MessageInteraction | Message | None = None,
        return_str: bool = False,
    ) -> str | list[dict[str, str]]:
        if ctx is None:
            location_context = "and friends in a Discord server"
        elif hasattr(ctx, "guild") and ctx.guild is not None:
            location_context = f'and friends in the "{ctx.guild.name}" Discord server'
        elif hasattr(ctx, "author") and ctx.author is not None:
            location_context = f"and {ctx.author.name} in a Discord DM"
        else:
            location_context = "and a friend in a Discord DM"

        replace_tokens = {
            "bot_name": self.name,
            "location_context": location_context,
            "current_time": get_prompt_datetime(with_date=self.prompt.with_date),
            "time_type": "time and date" if self.prompt.with_date else "time",
        }

        def apply_replacements(msg_dict: dict[str, str]) -> dict[str, str]:
            content: str = msg_dict["content"]
            for token, replacement in replace_tokens.items():
                content = content.replace("{" + token + "}", replacement)
            msg_dict["content"] = content
            return msg_dict

        sys_prompt = self.prompt.system.full
        if isinstance(sys_prompt, str):
            sys_prompt = [{"role": "system", "content": sys_prompt}]

        char_prompt = self.prompt.character.full
        if isinstance(char_prompt, str):
            char_prompt = [{"role": "system", "content": char_prompt}]

        # now both are list[dict[str, str]]
        prompt: list[dict[str, str]] = []

        if sys_prompt:
            sys_prompt = map(apply_replacements, sys_prompt)
            prompt.extend(sys_prompt)
        if char_prompt:
            char_prompt = map(apply_replacements, char_prompt)
            prompt.extend(char_prompt)

        if return_str is True:
            return "\n".join([msg["content"] for msg in prompt])
        else:
            return prompt

    # get the last N messages in a channel, up to and including the message that triggered the response
    async def get_message_context(
        self,
        message: Message,
        *,
        guild_settings: GuildSettings | None = None,
        max_messages: int | None = None,
    ) -> list[str] | list[dict[str, str]]:
        if max_messages is None:
            max_messages = self.params.context_messages

        messages = await message.channel.history(limit=max_messages, before=message).flatten()
        messages.insert(0, message)  # add the message that triggered the response

        # magic tag to break context chain to get bot out of librarian mode
        for idx, message in enumerate(messages):
            if self.check_ctxbreak(message):
                logger.debug("Found context break tag, breaking context chain")
                messages = messages[0 : idx + 1]
                break

        if guild_settings is None:
            # attempt to retrieve guild settings if we don't have them
            guild_settings = self.params.get_guild_settings(message.guild.id) if message.guild else None

        if guild_settings is not None:
            # get bot mode for this channel
            bot_mode = guild_settings.channel_bot_mode(message.channel.id)
        else:
            # default to siblings mode if we still don't have guild settings
            bot_mode = BotMode.Siblings

        def wrap_message(content: str, role: str = "user", name: str = "") -> dict[str, str]:
            if role == "assistant" and not name:
                name = f"{self.name}:"

            if name and not content.startswith(name):
                content = name.split(":")[0] + ": " + content.lstrip()

            return {"role": role.strip(), "content": content.strip()}

        skip = []  # cache for replies we've reordered
        chain: list[dict[str, str]] = []  # the context chain
        for msg in reversed(messages):
            if msg.id in skip:
                continue  # skip messages we've already reordered for reply-chasing
            if msg.author.id in self.tos_reject_ids:
                continue  # skip users who rejected the privacy policy
            if msg.content is not None:
                if any(msg.content.startswith(x) for x in ["/", ", "]):
                    continue  # skip messages that start with a command prefix or a comma-space
                if self.config.skip_codeblocks and "```" in msg.content:
                    continue  # skip messages with code blocks

            if msg.author.id == self.bot.user.id:
                if msg.content.startswith("This is an AI chatbot made by"):
                    continue  # this is a TOS message in a DM so lets skip it
                if msg.content.startswith(self.config.analysis_header):
                    continue
                if msg.embeds and msg.embeds[0].type == "rich":
                    # skip rich embeds from self (usually UI/command/analysis messages)
                    continue
            elif msg.author.bot is True:  # bots who aren't us
                if bot_mode == BotMode.Strip:
                    logger.debug(f"Stripping bot message {msg.id} from context chain")
                    continue  # skip all other bot messages if we're in strip mode
                elif bot_mode == BotMode.Siblings:
                    if msg.author.id not in self.sibling_ids:
                        logger.debug(f"Stripping non-sibling bot message {msg.id} from context chain")
                        continue  # skip non-sibling bot messages

            # set up author name
            if msg.author.id == self.bot.user.id:
                msg_role = "assistant"
                author_name = f"{self.name}:"
            else:
                msg_role = "user"
                author_name = get_message_author_name(msg, suffix=":")

            if len(msg.embeds) > 0:
                nitro_bs = self.handle_stupid_fucking_embed(msg)
                if nitro_bs is not None:
                    chain.append(wrap_message(f"{author_name} {nitro_bs}", msg_role))
                    continue
                for embed in msg.embeds:
                    if embed.type == "image":
                        try:
                            caption = await self.eyes.perceive_url(embed.thumbnail.url, msg.id)  # type: ignore
                            if caption is not None:
                                chain.append(wrap_message(f"{author_name} [image: {caption}]", msg_role))
                        except Exception as e:
                            logger.exception(e)
                            chain.append(wrap_message(f"{author_name} [image: loading error]", msg_role))
                    elif embed.description is not None and msg.author.bot is False:
                        chain.append(wrap_message(f"{author_name} [embed: {embed.description}]", msg_role))
                    elif embed.title is not None and msg.author.id != self.bot.user.id:
                        chain.append(wrap_message(f"{author_name} [embed: {embed.title}]", msg_role))

            if msg.content and len(msg.embeds) == 0:
                content = self.get_msg_content_clean(msg)
                if content is None:
                    logger.debug(f"Message {msg.id} was empty after cleaning, skipping...")
                    continue  # skip empty-after-cleaning messages

                if content.startswith("http"):
                    chain.append(wrap_message(f"{author_name} [a link to a webpage]", msg_role))
                    continue

                if msg.author.id == self.bot.user.id:
                    content = re_start_expression.sub("", content)
                    if len(content) >= 0:
                        chain.append(wrap_message(f"{author_name} {content}", msg_role))
                    else:
                        logger.debug(f"Self-sent message {msg.id} was empty after cleaning, skipping...")
                        continue
                else:
                    chain.append(wrap_message(f"{author_name} {content}", msg_role))

            for attachment in msg.attachments:
                try:
                    if attachment.content_type is None:
                        logger.debug(f"got no content-type for attachment: '{attachment.filename}', skipping")
                        continue
                    if not attachment.content_type.startswith("image/"):
                        logger.debug(f"got non-image content-type: '{attachment.content_type}', skipping")
                        continue
                    # caption image
                    caption = await self.eyes.perceive_attachment(attachment)
                    if caption is not None:
                        chain.append(wrap_message(f"{author_name} [image: {caption}]", msg_role))
                    else:
                        chain.append(wrap_message(f"{author_name} [image: unknown content]", msg_role))
                    continue
                except Exception as e:
                    logger.exception(e)
                    chain.append(wrap_message(f"{author_name} [image: loading error]", msg_role))

        return chain

    # assemble a prompt/context for the model
    async def process_context(self, conversation: list[str] | list[dict[str, str]], message: Message):
        """
        Apply the chat template to the prompt and history, trimming on whole-message
        boundaries until the context token budget is reached.
        """
        if self.prompt.disco_mode:
            logger.debug("Disco mode enabled, returning only the message that triggered the response")
            context = message.content
            context = re_mention.sub("", context)  # remove mentions
            return context.strip()

        prompt_messages, prompt_text, prompt_tokens = self._render_prompt_context(message)
        self._n_prompt_tokens = prompt_tokens

        if prompt_tokens > self.context_size:
            logger.warning("Prompt alone exceeds configured context size; sending prompt without history")

        history_messages = self._normalize_chat_messages(conversation, default_role="user")
        if len(history_messages) == 0:
            context = prompt_text
            context = re_consecutive_newline.sub("\n", context)
            return context.strip()

        context_size = self.context_size
        context_text = self.apply_chat_template(prompt_messages + history_messages[-1:])

        for idx in range(len(history_messages) - 2, -1, -1):
            candidate_messages = history_messages[idx:]
            candidate_text = self.apply_chat_template(prompt_messages + candidate_messages)
            candidate_tokens = self._count_tokens(candidate_text)
            if candidate_tokens <= context_size:
                context_text = candidate_text
            else:
                break

        context = re_consecutive_newline.sub("\n", context_text)
        context = context.replace("|> ", "|>")  # remove any spurious whitespace from the roles
        return context.strip()

    # actual response logic
    async def generate_response(
        self,
        conversation: list[str] | list[dict[str, str]],
        message: Message,
        trigger: str | None = None,
        append: str | None = None,
    ) -> str | None:
        async with message.channel.typing():
            debug_data: dict[str, Any] = {}

            response = ""
            resp_img = None
            should_reply = False
            author_name = get_message_author_name(message, suffix=":")

            msg_timestamp = message.created_at.astimezone(tz=self.bot.timezone).strftime(
                "%Y-%m-%d-%H:%M:%S%z"
            )
            msg_trigger = trigger.lower() if trigger is not None else "unknown"

            # build debug data
            debug_data["id"] = message.id
            debug_data["app_id"] = self.bot.user.id
            debug_data["instance"] = self.name.title()
            debug_data["message"] = {
                "id": message.id,
                "timestamp": msg_timestamp,
                "guild_id": message.guild.id if message.guild else None,
                "guild": str(message.guild) if message.guild else None,
                "author_id": message.author.id,
                "author": str(message.author),
                "channel_id": message.channel.id or None,
                "channel": message.channel.name if hasattr(message.channel, "name") else "DM",
                "trigger": msg_trigger,
                "author_name": author_name,
                "content": message.content,
            }
            # make conversation be a top level key
            if isinstance(conversation[0], dict):
                debug_data["conversation"] = [f"{x['content']}" for x in conversation]  # type: ignore
            else:
                debug_data["conversation"] = conversation

            try:
                debug_data["parameters"] = self.lm_request.model_dump(exclude_none=True)
            except Exception:
                logger.exception("Failed to get gensettings")

            # Build context
            context = await self.process_context(conversation, message)
            context = context.rstrip()
            debug_data["context"] = context.splitlines()
            debug_data["n_prompt_tokens"] = self._n_prompt_tokens
            debug_data["n_context_tokens"] = self._count_tokens(context)

            bad_words = []
            try:
                # Generate the response, and retry if it contains bad words (up to self.max_retries times)
                for attempt in range(self.max_retries):
                    attempt = attempt + 1  # deal with range() starting from 0
                    try:
                        response = await self.chatbot.respond_async(context)
                    except Exception as e:
                        logger.exception(e)
                        await message.add_reaction("❌")
                        return

                    # save raw response for debugging
                    debug_data["response_raw"] = response.splitlines(keepends=True)

                    # check for think tag and trim
                    if self.prompt.think_start in response:
                        logger.debug("Found chain of thought in response")
                        thoughts, response = response.rsplit(self.prompt.think_end, 1)
                        # debug log thought lines
                        thoughts = cleanup_thoughts([x.strip() for x in thoughts.splitlines()])
                        debug_data["thoughts"] = thoughts
                        logger.debug(f"Thoughts: {debug_data['thoughts']}")

                    rtemp = []
                    # first, strip any leading lines that are empty, and trim the bot's name from the start of each line if present
                    for line in response.splitlines():
                        # trim bot name from start of line if present
                        if self._starts_with_bot_name(line, ":"):
                            line = line.split(":", 1)[1].lstrip()
                        # if we're past the first non-empty line, keep all lines
                        if len(rtemp) > 0:
                            rtemp.append(line)
                            continue
                        # otherwise, only keep non-empty lines
                        if len(line.strip()) > 1:
                            rtemp.append(line)
                            continue
                    response = "\n".join(rtemp)
                    del rtemp

                    bad_words = self.find_bad_words(response)
                    if bad_words:
                        logger.info(f"Bad words in attempt #{attempt}: {bad_words}\nRetrying...")
                        continue
                    else:
                        if response.lower().startswith("i'm sorry, but"):
                            logger.info("Response was a ChatGPT apology, retrying...")
                            continue
                        if "as a language model" in response.lower():
                            logger.info("Response admits to being an AI, retrying...")
                            continue
                        if self.config.reject_urls and re_detect_url.search(response):
                            logger.info("Response contains a URL, retrying...")
                            continue
                        break  # no bad words, we're good, break out of the loop to avoid executing the else:

                else:  # ran out of retries...
                    if resp_img is not None:
                        response = ""  # we have a pic to send, so send it without a comment
                    else:
                        logger.warning(f"Final response contained bad words: {bad_words!r}, giving up.")
                        response = ""

                # see if there's an image description in the response
                description, response = self.imagen.find_image_desc(response)  # type: ignore
                if description is not None and "loading error" not in description.lower():
                    # there is, so we should reply with the image
                    logger.info("i'm feeling creative, let's make an image...")
                    try:
                        resp_img, img_debug_data = await self.take_pic(message=message)
                        debug_data.update(img_debug_data)
                        should_reply = True
                    except Exception:
                        logger.exception("Failed to generate image response")
                        pass

                # if we haven't described an image, but the user asked for one, send one anyway
                elif self.imagen.should_take_pic(message.content):
                    logger.info("the people have spoken and they want catgirl selfies")
                    try:
                        resp_img, img_debug_data = await self.take_pic(message=message)
                        debug_data.update(img_debug_data)
                        should_reply = True
                    except Exception:
                        logger.exception("Failed to generate image response")
                        pass

                # Clean response - trim left whitespace and fix emojis and pings
                response = self.restore_mentions_emoji(text=response, message=message)
                # Unescape markdown
                response = re_unescape_md.sub(r"\1", response)
                # scream into the void
                response = response.replace("\\r", "").replace("\\n", "\n")
                response = response.replace("\\u00a0", "\n")
                # em-dashes can go, thanks
                response = response.replace("—", " - ")

                if self.prompt.disco_mode:
                    if response == "":
                        logger.warning("Response was empty in disco mode...")
                    else:
                        logger.debug("Prepending prompt to response (disco mode)")
                        response = f"{context} {response}"

                # construct AiResponse object, don't use `response` variable name to avoid confusion
                resp = AiResponse(ctx=message, content=response, image=resp_img, is_reply=should_reply)
                del response  # avoid confusion

                if len(resp.content) > 1900:
                    if self.prompt.disco_mode:
                        logger.debug("Overlength response in disco mode, will send as file")
                        buf = BytesIO(resp.content.encode("utf-8"))
                        resp.file = File(buf, filename="story.txt")
                    else:
                        logger.debug("Response is too long, trimming...")
                        resp.content = resp.content[:1900].strip() + "-"

                if append is not None and len(append.strip()) > 0:
                    resp.content = f"{resp.content} < {append.strip()} >"
                debug_data["response"] = resp.content

                # Send response if not empty
                resp, sent = await self.send_response(resp)
                if isinstance(sent, Message):
                    debug_data["response_id"] = sent.id

                self.last_response = datetime.now(UTC)

                # update timestamp for this channel in the image caption engine
                if self.eyes.enabled and self.eyes.background:
                    self.eyes.watch(message.channel)  # type: ignore  # bad type hint in disnake

                # update webui state if it's enabled
                if self.webui is not None:
                    logger.debug("Updating webui")
                    self.webui.lm_update(
                        prompt=context,
                        message=f"{author_name}: {message.content}",
                        response=str(debug_data["response_raw"]).lstrip(),
                    )

            except Exception as e:
                logger.exception(e)
            finally:
                if self.debug:
                    dump_file = f"msg-{self.name.lower()}-{message.id}-{msg_timestamp}.json"
                    if debug_data.get("trigger", None) == "dm":
                        dump_file = self.debug_dir.joinpath("dm", dump_file)
                        dump_file.parent.mkdir(parents=True, exist_ok=True)
                    else:
                        dump_file = self.debug_dir.joinpath(dump_file)

                    with dump_file.open("w", encoding="utf-8") as f:
                        json.dump(debug_data, f, indent=4, skipkeys=True, default=str, ensure_ascii=False)
                    logger.debug(f"Dumped message debug data to {dump_file.name}")

                await self.log_response(debug_data)

    async def send_response(
        self, response: AiResponse, check: bool = True
    ) -> tuple[AiResponse, Message | None]:
        # check to see if the message was deleted or significantly edited before we could respond
        if check is True:
            message = await self.validate_message(response.ctx, strict=True)
            if message is False:
                logger.info("Context message was deleted before we could respond, aborting send")
                return response, None
            elif isinstance(message, Message):
                logger.info("Context message was edited before we could respond, aborting send")
                return response, None

        log_msg = "Response: "
        if response.content:
            if response.file:
                # response was too long so we made it a file
                response.content = " ".join(response.content.split(" ")[:20]) + "... <too long, attached>"
            log_msg += f"{response.content}"
        else:
            log_msg += "<no text>"

        if response.image is not None:
            # if we have an image to send, send it
            if not isinstance(response.image, File):
                logger.warning("Response image is not a File object, will attempt to convert")
                response.image = File(response.image, filename="image.png")
            log_msg += " (with image)"

        if response.file is not None:
            # if we have a file to send, send it
            if not isinstance(response.file, File):
                logger.warning("Response file is not a File object, will attempt to convert")
                response.file = File(response.file, filename="message.txt")
            log_msg += " (with file)"

            # short-circuit for empty responses
        if not response.content and not response.file and not response.image:
            # empty response, just react
            logger.info(log_msg + " -- empty response, adding reaction only")
            await response.add_reaction(self.config.empty_react)
            return response, None
        else:
            logger.info(log_msg)
            sent = await response.send()
            return response, sent

    async def log_response(self, debug_data: dict) -> None:
        debug_data["timestamp"] = datetime.now(tz=self.bot.timezone)
        try:
            log_entry = AiResponseLog(**debug_data)
            async with self.db_client.begin() as session:
                session.add(log_entry)
            logger.debug(f"Logged response to database: {log_entry.id}")
        except Exception:
            logger.exception("Failed to log response to database, continuing anyway")

    async def validate_message(
        self,
        message: Message,
        *,
        threshold: float = 0.1,
        strict: bool = False,
    ) -> Message | bool:
        """
        Check if a message was deleted or edited.
        Returns False if deleted, the new message if it was edited, or True if unchanged.
        """
        try:
            new_msg = await message.channel.fetch_message(message.id)
            if new_msg.edited_at != message.edited_at:
                # ignore capitalization and leading/trailing whitespace when checking for edits
                old_content = message.content.lower().strip()
                new_content = new_msg.content.lower().strip()
                # if threshold is 0.0 or less, any edit counts as significant
                if threshold <= 0.0:
                    return old_content == new_content  # return True if unchanged, False if edited
                else:
                    # use Levenshtein distance to determine if the edit was significant
                    distance = lev_distance(old_content, new_content, weights=(1, 1, 1))
                    if distance / max(len(old_content), 1) > threshold:
                        return new_msg

            return True  # message is unchanged
        except (NotFound, Forbidden):
            return False  # message was deleted or is otherwise inaccessible
        except Exception:
            logger.exception("Failed to check message status")
            # strict=True means we assume the message was deleted on all fetch errors
            return message

    @commands.message_command(
        name="Get Chain of Thought",
        contexts=InteractionContextTypes(bot_dm=True, guild=True, private_channel=True),
    )
    async def get_response_cot(self, ctx: MessageCommandInteraction):
        """Get the chain of thought behind a message we sent"""
        if ctx.target.author.id != self.bot.user.id:
            await ctx.send("I didn't send this message and I'm not psychic", ephemeral=True)
            return
        try:
            thoughts = await self.get_response_thoughts(ctx.target.id)
            if thoughts is None or len(thoughts) == 0:
                await ctx.send("Didn't find any thoughts for this message", ephemeral=True)
                return

            thoughts_file = None
            thoughts_str = "\n".join([self.config.analysis_header] + [f"> {x}" for x in thoughts])

            if len(thoughts_str) > 2000:
                buf = BytesIO("\n".join(thoughts).encode("utf-8"))
                thoughts_file = File(buf, filename=f"thoughts-{ctx.target.id}.txt")
                thoughts_str = (
                    self.config.analysis_header + "\n" + thoughts[0][:40] + "... (too long, attached)"
                )

            # thoughts_embed = Embed(title="Chain of Thought", description=thoughts_str, color=0x8D56EE)
            if thoughts_file:
                await ctx.send(thoughts_str, file=thoughts_file)
            else:
                await ctx.send(thoughts_str)

        except Exception as e:
            logger.exception(e)
            raise e

    # Get the chain of thought behind a message we sent
    async def get_response_thoughts(self, response_id: int) -> list[str] | None:
        """Get the chain of thought for a message ID"""
        try:
            async with self.db_client.begin() as session:
                query = select(AiResponseLog).where(AiResponseLog.response_id == response_id).limit(1)
                result = await session.scalars(query)
                response: AiResponseLog | None = result.first()
                if response is None:
                    logger.debug(f"No chain of thought found for message ID {response_id}")
                    return None
        except Exception:
            logger.exception("Failed to get response log")
            return None
        return response.thoughts

    # Idle loop, broken rn, need to factor on_message logic out into functions or smth
    @tasks.loop(seconds=21)
    async def idle_loop(self) -> None:
        if self.idle_enable is True:
            # get last message in a random priority channel
            channel = self.bot.get_channel(random_choice(self.idle_channels))
            if isinstance(channel, (TextChannel, DMChannel, GroupChannel)):
                messages = await channel.history(limit=1).flatten()
                message: Message = messages.pop()
                if message.author.bot is True:
                    return
                idle_sec = (datetime.now(tz=UTC) - message.created_at).total_seconds()
                if idle_sec >= 2**31:
                    if self.get_msg_content_clean(message) is None:
                        return
                    logger.debug(f"Running idle response to message {message.id}")
                    # if it's been more than <idle_interval> sec, send a response
                    async with self.lm_lock:
                        conversation = await self.get_message_context(message)
                        await self.generate_response(conversation, message)
        else:
            self.idle_loop.stop()
            return

    @idle_loop.before_loop
    async def before_idle_loop(self):
        logger.info("Idle loop waiting for bot to be ready")
        await self.bot.wait_until_ready()
        await asyncio.sleep(10)
        logger.info("Idle loop running!")

    ## Helper functions
    # Check and refresh the list of users who rejected the ToS
    @tasks.loop(seconds=60)
    async def update_ignored(self):
        """Check the database for users who have rejected the ToS and get a list of their IDs for message filtering"""
        async with self.db_client.begin() as session:
            query = (
                select(DiscordUser)
                .where(DiscordUser.tos_rejected == True)
                .with_only_columns(DiscordUser.id, DiscordUser.tos_accepted, DiscordUser.tos_rejected)
            )
            results = await session.scalars(query)
            users: list[DiscordUser] = results.all()  # type: ignore
        self.tos_reject_ids = set([x.id for x in users])

    # Check if a user has accepted the ToS
    async def user_accepted_tos(self, user: User | Member | int) -> bool | None:  # type: ignore
        """Checks if a user has accepted the ToS, rejected, or not completed the process.
        Returns True if accepted, False if rejected, None if not completed.
        """
        user_id = user.id if isinstance(user, (User, Member)) else user
        async with self.db_client.begin() as session:
            user: DiscordUser | None = await session.get(DiscordUser, user_id)  # type: ignore
            if user is None:
                logger.debug(f"User {user_id} has not completed the ToS process.")
                return None
            elif user.tos_rejected is True:
                return False
            elif user.tos_accepted is True:
                return True

    # get a list of users who have accepted the ToS
    async def tos_accepted_users(self) -> set[int]:
        async with self.db_client.begin() as session:
            query = (
                select(DiscordUser)
                .options(load_only("id", "tos_accepted", "tos_rejected", raiseload=True))  # type: ignore
                .filter(DiscordUser.tos_accepted is True)  # type: ignore
            )
            result = await session.scalars(query)
            users = result.all()
        return {x.id for x in users}

    # check if the person triggering the bot has agreed to ToS or not and send it if not
    async def check_send_tos(self, message: Message) -> bool:
        logger.debug(f"Checking ToS for {message.author} ({message.author.id})")
        try:
            async with self.db_client.begin() as session:
                user: DiscordUser | None = await session.get(DiscordUser, message.author.id)
                if user is None:
                    logger.info(f"User {message.author} not found in database, creating entry")
                    user = DiscordUser.from_discord(message.author)
                    await session.merge(user)
                    await session.flush()
                    user = await session.get(DiscordUser, message.author.id)
                if user is None:
                    raise Exception("Failed to create or retrieve user")

            # send the form in a DM if they haven't accepted or rejected it yet
            if user.tos_accepted is not True:
                logger.debug(f"Sending privacy embed to {message.author} ({message.author.id})")
                invite = await self.bot.support_invite()
                embed = PrivacyEmbed(
                    author=message.author, support_guild=self.bot.support_guild, user=user, invite=invite
                )
                view = PrivacyView(user=user)
                content = get_policy_text(self.bot)
                view.message = await message.author.send(content=content, embed=embed, view=view)
                logger.debug(f"Sent privacy embed to {message.author} ({message.author.id})")
                return False
            else:
                logger.debug(f"User {message.author} ({message.author.id}) has accepted ToS")
                return True
        except Exception:
            logger.exception("error checking tos")
            return False

    # clean up a message's content
    def get_msg_content_clean(self, message: Message, content: str | None = None) -> str | None:
        # if no content is provided, use the message's content
        content = content or message.content
        # if there's a codeblock in there, return None
        if "```" in content:
            return None
        # convert mentions and emoji from user IDs to to plain text
        content = self.stringify_mentions_emoji(content, message=message)
        # eliminate any and all content between angle brackets <>
        content = re_angle_bracket.sub("", content)
        # strip trailing whitespace
        content = content.rstrip()
        # return the cleaned content, or None if it's empty
        return content if len(content) > 0 else None

    # search for bad words in a message
    def find_bad_words(self, input: str) -> list[str]:
        found_words: list[str] = []
        input_words: str = input.split()  # type: ignore

        for word in input_words:
            word_stripped: str = re_nonword.sub("", word).lower()
            word_len = len(word_stripped)
            if word_len > 2:
                threshold = 2 if word_len > 6 else 1 if word_len > 4 else 0
                for bad_word in self.bad_words:
                    if lev_distance(word_stripped, bad_word) <= threshold:
                        found_words.append(f"{word} (like {bad_word})")
        if len(found_words) > 0:
            logger.warning(f"Found bad words: {' | '.join(found_words)}")
        return found_words

    # check if a line starts with the bot's name or a nickname, with an optional suffix
    def _starts_with_bot_name(self, line: str, suffix: str = "") -> bool:
        bot_names = set([self.name] + self.params.nicknames + self.params.nicknames_quiet)
        line = line.lower().lstrip()
        suffix = suffix.lower()
        for name in bot_names:
            if line.startswith(name.lower() + suffix):
                return True
        return False

    # check if the bot's name or a nickname is in the message
    def _bot_name_in_text(self, text: str, with_quiet: bool = False) -> bool:
        bot_names = set(
            [self.name] + self.params.nicknames + (self.params.nicknames_quiet if with_quiet else [])
        )
        for name in bot_names:
            if re.search(rf"\b({name})({name[-1]}*)\b", text, re.I + re.M):
                return True
        return False

    # check if the bot was mentioned in a message
    def check_mention(self, message: Message) -> bool:
        """Checks if the bot was mentioned in a message"""
        if message.guild is None:
            return False
        if self.bot.user in message.mentions or message.guild.me in message.mentions:
            return True
        if self._bot_name_in_text(message.content, with_quiet=False):
            return True
        return False

    # check if a given message is a context break message
    def check_ctxbreak(self, message: Message) -> bool:
        content = message.content.lower()
        if not content.startswith("<ctxbreak>"):
            return False  # not a context break message

        if self.config.params.ctxbreak_restrict is False:
            return True  # ctxbreak is unrestricted

        if isinstance(message.channel, DMChannel):
            return True  # people can always break context in DMs
        if message.author.id in (self.ctxbreak.user_ids + self.bot.config.admin_ids):
            return True  # admins and ctxbreak users can always break context
        if member_in_any_role(message.author, self.ctxbreak.role_ids):  # type: ignore
            return True  # ctxbreak role havers can always break context
        return False  # otherwise, no

    # check if the last message we responded to in this context was from a bot
    def last_trigger_was_bot(self, message: Message) -> bool:
        """Check if the last message we responded to in this context was from a bot."""
        cache_entry = self._trigger_cache.get(message.channel.id, False)
        return cache_entry

    # deal with fake nitro bullshit
    def handle_stupid_fucking_embed(self, message: Message) -> str | None:
        """Check if a message is a stupid fucking spec-beaking fake nitro emote embed
        Returns the emote name if it is, otherwise returns None.
        """
        if message.content.startswith("https://cdn.discordapp.com/emojis/"):
            logger.debug("Found a stupid fake nitro emote")
            name = message.content.split("name=", 1)[-1]
            name = name.split("&", 1)[0]
            return f":{name}:"
        else:
            return None

    ## Loop to archive debug logs
    @tasks.loop(hours=1)
    async def log_archive_loop(self) -> None:
        """
        Move message log files over 24h old to the archive directory
        """
        for file in self.debug_dir.glob("*.json"):
            file_age = (datetime.now() - datetime.fromtimestamp(file.stat().st_mtime)).total_seconds()
            if file_age >= (3600 * 24):
                logger.debug(f"Archiving debug log {file.name}")
                file = file.rename(file.parent / "archive" / file.name)

    @log_archive_loop.before_loop
    async def before_archive_loop(self):
        logger.info("Archive loop waiting for bot to be ready")
        await self.bot.wait_until_ready()
        self.debug_dir.joinpath("archive").mkdir(exist_ok=True, parents=True)
        logger.info("Archive loop running!")

    ## Image generation stuff
    async def take_pic(self, message: Message | str) -> tuple[File | None, dict[str, Any]]:
        """
        hold up, let me take a selfie
        """
        # get the message content
        if isinstance(message, Message):
            content: str | None = self.get_msg_content_clean(message)  # type: ignore
        elif isinstance(message, str):
            content: str = message
        else:
            raise ValueError("take_pic got an unexpected message type (not str or Message)")

        if not content:
            logger.debug("Message was empty after cleaning.")
            return None, {}

        debug_data: dict = {}

        # lowercase the content
        content = content.lower().strip()
        # get the image description and remove newlines
        content = content.replace("\n", " ").replace(f"{self.name} ", "")

        # build the LLM prompt for the image
        lm_trigger = self.imagen.strip_take_pic(content)
        lm_prompt, user_request = self.imagen.get_lm_prompt(lm_trigger)
        debug_data["imagen_llm_prompt"] = lm_prompt
        logger.info(f"LLM Request: {user_request}")

        # get the LLM to create tags for the image
        lm_tag_string = await self.imagen.submit_lm_prompt(lm_prompt)
        lm_tag_string = lm_tag_string.strip('",').lower()
        debug_data["imagen_llm_tags"] = lm_tag_string
        logger.info(f"LLM Response: {lm_tag_string}")

        # build the SD API request
        sdapi_request = self.imagen.build_request(lm_tag_string, content)
        debug_data["imagen_sdapi_request"] = json.dumps(sdapi_request)
        logger.info(f"SD API Request: {sdapi_request!r}")

        # submit it
        try:
            result_path = await self.imagen.submit_request(sdapi_request)
            debug_data["imagen_result_path"] = str(result_path)
            logger.info(f"Image generated at {result_path}")
            # return a discord File object to upstream
            result_file = File(result_path, filename=result_path.name)

            if self.webui.config.enabled:
                try:
                    self.webui.imagen_update(lm_trigger, lm_tag_string, result_path)
                except Exception:
                    logger.exception("Failed to update webui with image generation data")
            return result_file, debug_data  # type: ignore  # bad type hint in disnake
        except RuntimeError as e:
            logger.exception("Image generation failed")
            raise e

    ## UI stuff
    @commands.slash_command(name="ai", description="AI Management")
    @checks.not_blacklisted()
    async def ai_group(self, ctx: ApplicationCommandInteraction):
        pass

    @ai_group.sub_command(name="status", description="Get AI status")
    async def ai_status(
        self,
        ctx: ApplicationCommandInteraction,
        verbose: bool = commands.Param(default=False, name="verbose", description="Verbose output"),
    ):
        embed = AiStatusEmbed(self, ctx.author, verbose=verbose)
        await ctx.send(embed=embed, ephemeral=True)

    @ai_group.sub_command(name="set", description="Set AI parameters")
    @checks.is_admin()
    async def set_parameter(
        self,
        ctx: ApplicationCommandInteraction,
        param: AiParam = commands.Param(
            description="Parameter", choices=set_choices, converter=convert_param
        ),
        value: str = commands.Param(...),
    ) -> None:
        await ctx.response.defer(ephemeral=True)
        embed = Embed(title="Set Parameter", description="Success!", color=Colour.green())
        embed.set_author(name=ctx.author.display_name, icon_url=ctx.author.display_avatar.url)
        embed.add_field(name="Parameter", value=f"{param.name}", inline=False)

        try:
            if not hasattr(self.lm_request, param.id):
                raise ValueError(f"Unknown parameter: {param} (got {param.id}, {param.kind})")

            # cast the value to the correct type using the class from the tuple
            new_value = param.kind(value)
            # save the old value and set the new one
            old_value = getattr(self.lm_request, param.id)
            setattr(self.lm_request, param.id, new_value)

            # add the old and new values to the embed
            embed.add_field(name="Old", value=f"{old_value}")
            embed.add_field(name="New", value=f"{new_value}")

        except Exception as e:
            logger.exception(e)
            embed.description = "Failed!"
            embed.color = Colour.red()
            if isinstance(e, ValueError):
                embed.add_field(name="Error", value=f"Invalid value: {value}", inline=False)
            else:
                embed.add_field(name="Error", value="An unknown error occurred", inline=False)
            embed.add_field(name="Exception", value=e, inline=False)
        finally:
            # send the embed
            await ctx.send(embed=embed, ephemeral=True)


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