import asyncio
import mimetypes
import re
import uuid
from io import BytesIO
from typing import TYPE_CHECKING, Dict, List, Optional, Union

if TYPE_CHECKING:
    from discord.abc import MessageableChannel

import discord
import filetype
import httpx
from chainlit.config import config
from chainlit.context import ChainlitContext, HTTPSession, context_var
from chainlit.data import get_data_layer
from chainlit.element import Element, ElementDict
from chainlit.emitter import BaseChainlitEmitter
from chainlit.logger import logger
from chainlit.message import Message, StepDict
from chainlit.telemetry import trace
from chainlit.types import Feedback
from chainlit.user import PersistedUser, User
from chainlit.user_session import user_session
from discord.ui import Button, View


class FeedbackView(View):
    def __init__(self, step_id: str):
        super().__init__(timeout=None)
        self.step_id = step_id

    @discord.ui.button(label="👎")
    async def thumbs_down(self, interaction: discord.Interaction, button: Button):
        if data_layer := get_data_layer():
            try:
                thread_id = context_var.get().session.thread_id
                feedback = Feedback(forId=self.step_id, threadId=thread_id, value=0)
                await data_layer.upsert_feedback(feedback)
            except Exception as e:
                logger.error(f"Error upserting feedback: {e}")
        if interaction.message:
            await interaction.message.edit(view=None)
            await interaction.message.add_reaction("👎")

    @discord.ui.button(label="👍")
    async def thumbs_up(self, interaction: discord.Interaction, button: Button):
        if data_layer := get_data_layer():
            try:
                thread_id = context_var.get().session.thread_id
                feedback = Feedback(forId=self.step_id, threadId=thread_id, value=1)
                await data_layer.upsert_feedback(feedback)
            except Exception as e:
                logger.error(f"Error upserting feedback: {e}")
        if interaction.message:
            await interaction.message.edit(view=None)
            await interaction.message.add_reaction("👍")


class DiscordEmitter(BaseChainlitEmitter):
    def __init__(
        self, session: HTTPSession, channel: "MessageableChannel", enabled=False
    ):
        super().__init__(session)
        self.channel = channel
        self.enabled = enabled

    async def send_element(self, element_dict: ElementDict):
        if not self.enabled or element_dict.get("display") != "inline":
            return

        persisted_file = self.session.files.get(element_dict.get("chainlitKey") or "")
        file: Optional[Union[BytesIO, str]] = None
        mime: Optional[str] = None

        if persisted_file:
            file = str(persisted_file["path"])
            mime = element_dict.get("mime")
        elif file_url := element_dict.get("url"):
            async with httpx.AsyncClient() as client:
                response = await client.get(file_url)
                if response.status_code == 200:
                    file = BytesIO(response.content)
                    mime = filetype.guess_mime(file)

        if not file:
            return

        element_name: str = element_dict.get("name", "Untitled")

        if mime:
            file_extension = mimetypes.guess_extension(mime)
            if file_extension:
                element_name += file_extension

        file_obj = discord.File(file, filename=element_name)
        await self.channel.send(file=file_obj)

    async def send_step(self, step_dict: StepDict):
        if not self.enabled:
            return

        step_type = step_dict.get("type")
        is_message = step_type in [
            "user_message",
            "assistant_message",
        ]
        is_chain_of_thought = bool(step_dict.get("parentId"))
        is_empty_output = not step_dict.get("output")

        if is_chain_of_thought or is_empty_output or not is_message:
            return
        else:
            enable_feedback = not step_dict.get("disableFeedback") and get_data_layer()
            message = await self.channel.send(step_dict["output"])

            if enable_feedback:
                view = FeedbackView(step_dict.get("id", ""))
                await message.edit(view=view)

    async def update_step(self, step_dict: StepDict):
        if not self.enabled:
            return

        await self.send_step(step_dict)


intents = discord.Intents.default()
intents.message_content = True

client = discord.Client(intents=intents)


@trace
def init_discord_context(
    session: HTTPSession,
    channel: "MessageableChannel",
    message: discord.Message,
) -> ChainlitContext:
    emitter = DiscordEmitter(session=session, channel=channel)
    context = ChainlitContext(session=session, emitter=emitter)
    context_var.set(context)
    user_session.set("discord_message", message)
    user_session.set("discord_channel", channel)
    return context


users_by_discord_id: Dict[int, Union[User, PersistedUser]] = {}

USER_PREFIX = "discord_"


async def get_user(discord_user: Union[discord.User, discord.Member]):
    if discord_user.id in users_by_discord_id:
        return users_by_discord_id[discord_user.id]

    metadata = {
        "name": discord_user.name,
        "id": discord_user.id,
    }
    user = User(identifier=USER_PREFIX + str(discord_user.name), metadata=metadata)

    users_by_discord_id[discord_user.id] = user

    if data_layer := get_data_layer():
        try:
            persisted_user = await data_layer.create_user(user)
            if persisted_user:
                users_by_discord_id[discord_user.id] = persisted_user
        except Exception as e:
            logger.error(f"Error creating user: {e}")

    return users_by_discord_id[discord_user.id]


async def download_discord_file(url: str):
    async with httpx.AsyncClient() as client:
        response = await client.get(url)
        if response.status_code == 200:
            return response.content
        else:
            return None


async def download_discord_files(
    session: HTTPSession, attachments: List[discord.Attachment]
):
    download_coros = [
        download_discord_file(attachment.url) for attachment in attachments
    ]
    file_bytes_list = await asyncio.gather(*download_coros)
    file_refs = []
    for idx, file_bytes in enumerate(file_bytes_list):
        if file_bytes:
            name = attachments[idx].filename
            mime_type = attachments[idx].content_type or "application/octet-stream"
            file_ref = await session.persist_file(
                name=name, mime=mime_type, content=file_bytes
            )
            file_refs.append(file_ref)

    files_dicts = [
        session.files[file["id"]] for file in file_refs if file["id"] in session.files
    ]

    file_elements = [Element.from_dict(file_dict) for file_dict in files_dicts]

    return file_elements


def clean_content(message: discord.Message):
    if not client.user:
        return message.content

    # Regex to find mentions of the bot
    bot_mention = f"<@!?{client.user.id}>"
    # Replace the bot's mention with nothing
    return re.sub(bot_mention, "", message.content).strip()


async def process_discord_message(
    message: discord.Message,
    thread_name: str,
    channel: "MessageableChannel",
    bind_thread_to_user=False,
):
    user = await get_user(message.author)

    thread_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, str(channel.id)))

    text = clean_content(message)
    discord_files = message.attachments

    session_id = str(uuid.uuid4())
    session = HTTPSession(
        id=session_id,
        thread_id=thread_id,
        user=user,
        client_type="discord",
    )

    ctx = init_discord_context(
        session=session,
        channel=channel,
        message=message,
    )

    file_elements = await download_discord_files(session, discord_files)

    msg = Message(
        content=text,
        elements=file_elements,
        type="user_message",
        author=user.metadata.get("name"),
    )

    await msg.send()

    ctx.emitter.enabled = True

    if on_chat_start := config.code.on_chat_start:
        await on_chat_start()

    if on_message := config.code.on_message:
        async with channel.typing():
            await on_message(msg)

    if on_chat_end := config.code.on_chat_end:
        await on_chat_end()

    if data_layer := get_data_layer():
        user_id = None
        if isinstance(user, PersistedUser):
            user_id = user.id if bind_thread_to_user else None

        try:
            await data_layer.update_thread(
                thread_id=thread_id,
                name=thread_name,
                metadata=ctx.session.to_persistable(),
                user_id=user_id,
            )
        except Exception as e:
            logger.error(f"Error updating thread: {e}")

    ctx.session.delete()


@client.event
async def on_ready():
    logger.info(f"Logged in as {client.user}")


@client.event
async def on_message(message: discord.Message):
    if not client.user or message.author == client.user:
        return

    is_dm = isinstance(message.channel, discord.DMChannel)
    if not client.user.mentioned_in(message) and not is_dm:
        return

    thread_name: str = ""
    bind_thread_to_user = False
    channel = message.channel

    if isinstance(message.channel, discord.Thread):
        thread_name = f"{message.channel.name}"
    elif isinstance(message.channel, discord.ForumChannel):
        thread_name = f"{message.channel.name}"
    elif isinstance(message.channel, discord.DMChannel):
        thread_name = f"{message.author} Discord DM"
        bind_thread_to_user = True
    elif isinstance(message.channel, discord.GroupChannel):
        thread_name = f"{message.channel.name}"
    elif isinstance(message.channel, discord.TextChannel):
        # Discord limits thread names to 100 characters and does not create
        # threads from empty messages.
        discord_thread_name = clean_content(message)[:100] or "Untitled"
        channel = await message.channel.create_thread(
            name=discord_thread_name, message=message
        )
        thread_name = f"{channel.name}"
    else:
        logger.warning(f"Unsupported channel type: {message.channel.type}")
        return

    await process_discord_message(
        message=message,
        thread_name=thread_name,
        channel=channel,
        bind_thread_to_user=bind_thread_to_user,
    )
