import collections
import discord
from discord.ext import commands
from openai import AsyncOpenAI
import os
import base64
import aiohttp
import argparse
from typing import List, Dict, Any

# --- Configuration ---
OPENAI_API_KEY = "eh"
MODEL = "p620"
DEFAULT_SYSTEM_PROMPT = "you are a catboy named Aoi with dark blue fur and is a tsundere"
DEFAULT_NAME = "Aoi"
DEFAULT_AVATAR = "https://cdn.discordapp.com/avatars/1406466525858369716/f1dfeaf2a1c361dbf981e2e899c7f981?size=256"

# --- Data Structures ---
class Conversation:
    def __init__(self, prompt, name):
        self.history = [{"role": "system", "content": prompt}]
        self.bot_name = name

    def add_message_pair(self, user, assistant):
        self.history.extend([
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ])

# --- Command Line Arguments ---
parser = argparse.ArgumentParser(description="Aoi Discord Bot")
parser.add_argument(
    '--base_url',
    type=str,
    required=True,
    help='The base URL for the OpenAI API.',
)
parser.add_argument(
    '--discord_token',
    type=str,
    required=True,
    help='The Discord bot token.',
)
args = parser.parse_args()

# --- Bot Setup ---
intents = discord.Intents.default()
intents.messages = True
intents.message_content = True
bot = commands.Bot(command_prefix="/", intents=intents)

# --- Data Storage ---
# Keyed by channel ID
conversation_history: Dict[int, Conversation] = collections.defaultdict(
    lambda: Conversation(prompt=DEFAULT_SYSTEM_PROMPT, name=DEFAULT_NAME),
)
_webhooks = {}
async def webhook(channel):
    if channel.id not in _webhooks:
        _webhooks[channel.id] = await channel.create_webhook(name=f'aoi-{channel.id}')
    return _webhooks[channel.id]

# --- OpenAI Client ---
client = AsyncOpenAI(
    base_url=args.base_url,
    api_key=OPENAI_API_KEY,
)

# --- Helpers ---
async def get_user_from_id(ctx, userid):
    if ctx.guild:
        user = await ctx.guild.fetch_member(userid)
    else:
        user = await bot.fetch_user(userid)
    return user.display_name

async def get_user_from_mention(ctx, mention):
    match = re.findall(r"<@!?(\d+)>", mention)
    if not match:
        return mention
    return await get_user_from_id(ctx, int(match[0]))

# --- Bot Events ---
@bot.event
async def on_ready():
    print(f'Logged in as {bot.user.name}')
    print(f'Using OpenAI base URL: {args.base_url}')
    await bot.tree.sync()

@bot.event
async def on_message(message):
    if message.author == bot.user:
        return

    if bot.user.mentioned_in(message):
        bot_tag = f'<@{bot.user.id}>'
        channel = message.channel
        conversation = conversation_history[channel.id]
        user_message_text = message.content
        if user_message_text.startswith(bot_tag):
            user_message_text = user_message_text[len(bot_tag):]
        user_message_text = user_message_text.replace(bot_tag, conversation.bot_name).strip()
        print(f'> {message.author.name}: {user_message_text}')

        # Prepare content for OpenAI API
        openai_content = []
        if user_message_text:
            openai_content.append({"type": "text", "text": user_message_text})
        else:
            openai_content.append({"type": "text", "text": "."})

        if message.attachments:
            async with aiohttp.ClientSession() as session:
                for attachment in message.attachments:
                    if attachment.content_type and "image" in attachment.content_type:
                        async with channel.typing():
                            try:
                                async with session.get(attachment.url) as resp:
                                    if resp.status == 200:
                                        image_data = await resp.read()
                                        base64_image = base64.b64encode(image_data).decode('utf-8')
                                        image_url = f"data:{attachment.content_type};base64,{base64_image}"
                                        openai_content.append({
                                            "type": "image_url",
                                            "image_url": {"url": image_url}
                                        })
                            except Exception as e:
                                print(f"Error downloading or processing attachment: {e}")

        request = conversation.history + [{"role": "user", "content": openai_content}]
        try:
            async with channel.typing():
                response = await client.chat.completions.create(
                    model=MODEL, messages=request,
                )
                bot_response = response.choices[0].message.content
                conversation.add_message_pair(openai_content, bot_response)
                # Split into chunks for discord to prevent message too long
                chunks = [bot_response[i:i+2000] for i in range(0, len(bot_response), 2000)]
                for chunk in chunks:
                    if channel.guild:
                        hook = await webhook(channel)
                        await hook.send(
                            content=chunk,
                            username=conversation.bot_name,
                            avatar_url=DEFAULT_AVATAR,
                            wait=True,
                        )
                    else:
                        await channel.send(content=chunk)
        except Exception as e:
            print(f"An error occurred: {e}")
            await message.reply("Sorry, I had a little hiccup. Baka!")


# --- Slash Commands ---
@bot.tree.command(name="newchat", description="Start a new chat with a new system prompt.")
async def newchat(interaction: discord.Interaction, prompt: str = None):
    await interaction.response.defer()
    channel_id = interaction.channel_id
    prompt = prompt or DEFAULT_SYSTEM_PROMPT
    name_response = await client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": prompt},
            {"role": "user", "content": "reply with your name, nothing else, no punctuation"}
        ],
    )
    name = name_response.choices[0].message.content.split('\n')[0]
    print(f'$ name={name}')
    conversation_history[channel_id] = Conversation(prompt=prompt, name=name)
    await interaction.followup.send(f'Starting a new chat with the prompt: "{prompt}"')


# --- Running the Bot ---
if __name__ == "__main__":
    bot.run(args.discord_token)