message_handler.py

import discord
import re
import logging
import asyncio
import aiohttp
from io import BytesIO

logger = logging.getLogger(__name__)

class MessageHandler:
    def __init__(self, api_client=None, memory_manager=None, config=None, data_manager=None, bot=None):
        self.api_client = api_client
        self.memory_manager = memory_manager
        self.config = config
        self.data_manager = data_manager
        self.bot = bot

    async def handle_message(self, message):
        channel_id = str(message.channel.id)
        guild_id = str(message.guild.id) if message.guild else "DM"
        user_id = str(message.author.id)
        user_message = message.content
        if user_message.lower().startswith("!"):
            return
        self.memory_manager.add_user_message(channel_id, guild_id, user_id, user_message)
        user_model = self.memory_manager.get_user_model(guild_id, user_id)
        system_prompt = f"{self.config.system_instructions}\nYou are {user_model}."
        messages = [{"role": "system", "content": system_prompt}]
        channel_memories = self.memory_manager.channel_memories.get(channel_id, [])
        if channel_memories:
            messages.append({"role": "user", "content": "\n".join(channel_memories)})
        model_history = self.memory_manager.get_user_model_history(guild_id, user_id, user_model)
        for msg in model_history:
            if msg["content"].strip():
                role = "assistant" if msg["role"] == "ai" else msg["role"]
                messages.append({"role": role, "content": msg["content"]})
        logger.info(f"Preparing to send message for user {user_id} with model {user_model}")
        if not self.api_client:
            await message.channel.send(f"<@{user_id}> Error: API client not initialized")
            logger.error(f"API client is None for user {user_id}")
            return
        try:
            ai_response = await self.api_client.send_message(messages, user_model)
            if not ai_response or not ai_response.strip():
                await message.channel.send(f"<@{user_id}> Error: Empty response from API")
                return
        except Exception as e:
            await message.channel.send(f"<@{user_id}> Error: Failed to fetch response - {e}")
            return
        ai_response_clean = self.clean_response(ai_response)
        if not ai_response_clean:
            ai_response_clean = "Empty response from API—please try again later."
        final_message = self.build_message(ai_response_clean)
        try:
            await self._send_message(message, user_id, final_message, user_message.lower())
        except Exception as e:
            logger.error(f"Failed to send message for user {user_id}: {e}")
            await message.channel.send(f"<@{user_id}> Failed to send response - please try again.")

    def clean_response(self, response: str) -> str:
        response = response.strip()
        response = re.sub(r"\n{3,}", "\n\n", response)
        cleaned_lines = [line for line in response.split("\n") if not re.match(r"^---", line) and "Generated by" not in line and not re.search(r"\d{4}-\d{2}-\d{2}", line)]
        return "\n".join(cleaned_lines).strip()

    def build_message(self, content: str) -> dict:
        url_pattern = r"https?://\S+"
        image_urls = re.findall(url_pattern, content)
        content_without_urls = re.sub(url_pattern, "", content).strip()
        return {"content": content_without_urls, "image_urls": image_urls}

    async def _send_message(self, message, user_id, final_message, user_message):
        content = final_message["content"]
        image_urls = final_message["image_urls"]
        guild_id = str(message.guild.id) if message.guild else "DM"
        files = []
        if image_urls:
            async with aiohttp.ClientSession() as session:
                for url in image_urls:
                    try:
                        async with session.get(url) as resp:
                            if resp.status == 200:
                                data = await resp.read()
                                files.append(discord.File(BytesIO(data), filename="image.png"))
                            else:
                                logger.error(f"Failed to fetch image from {url}: status {resp.status}")
                    except Exception as e:
                        logger.error(f"Error fetching image from {url}: {e}")

        if content or files:
            prefixed_content = f"<@{user_id}> {content}" if content else f"<@{user_id}>"
            content_length = len(prefixed_content)
            logger.info(f"Response length for user {user_id}: {content_length} characters")
            if 'code' in user_message or 'script' in user_message:
                code_blocks = re.findall(r'\[CODE\]\s*(\w+)\s*([\s\S]*?)\[/CODE\]', prefixed_content, re.DOTALL)
                for lang, code in code_blocks:
                    code_content = f"```{lang}\n{code.strip()}\n```"
                    if len(code_content) <= 2000:
                        await message.channel.send(code_content)
                    elif len(code_content) <= 4096:
                        embed = discord.Embed(description=code_content[:4096])
                        await message.channel.send(embed=embed)
                    else:
                        buffer = BytesIO(code_content.encode('utf-8'))
                        file = discord.File(buffer, filename=f"{lang}_code.txt")
                        await message.channel.send(f"<@{user_id}> Code too long, attached as file:", file=file)
                prefixed_content = re.sub(r'\[CODE\][\s\S]*?\[/CODE\]', '', prefixed_content).strip()
                content_length = len(prefixed_content)

            if content_length > 0 or files:
                if content_length <= 2000:
                    await message.channel.send(prefixed_content, files=files if files else None)
                elif content_length <= 4096:
                    embed = discord.Embed(description=prefixed_content[:4096])
                    await message.channel.send(embed=embed, files=files if files else None)
                else:
                    buffer = BytesIO(prefixed_content.encode('utf-8'))
                    text_file = discord.File(buffer, filename="response.txt")
                    files.append(text_file)
                    await message.channel.send(f"<@{user_id}> Response too long, attached as file and images:", files=files)
            elif files:  # If no content but images exist
                await message.channel.send(f"<@{user_id}>", files=files)

        if content.strip():
            self.memory_manager.add_ai_message(str(message.channel.id), guild_id, user_id, content)