import asyncio
import base64
import mimetypes
import os
import uuid
from datetime import datetime
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union

import filetype

if TYPE_CHECKING:
    from botbuilder.core import TurnContext
    from botbuilder.schema import Activity

import httpx
from botbuilder.core import (
    BotFrameworkAdapter,
    BotFrameworkAdapterSettings,
    MessageFactory,
    TurnContext,
)
from botbuilder.schema import (
    ActionTypes,
    Activity,
    ActivityTypes,
    Attachment,
    CardAction,
    ChannelAccount,
    HeroCard,
)

from chainlit.config import config
from chainlit.context import ChainlitContext, HTTPSession, context, 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


class TeamsEmitter(BaseChainlitEmitter):
    def __init__(self, session: HTTPSession, turn_context: TurnContext):
        super().__init__(session)
        self.turn_context = turn_context

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

        persisted_file = self.session.files.get(element_dict.get("chainlitKey") or "")
        attachment: Optional[Attachment] = None
        mime: Optional[str] = None

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

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

        if persisted_file:
            mime = element_dict.get("mime")
            with open(persisted_file["path"], "rb") as file:
                dencoded_string = base64.b64encode(file.read()).decode()
                content_url = f"data:{mime};base64,{dencoded_string}"
                attachment = Attachment(
                    content_type=mime, content_url=content_url, name=element_name
                )

        elif url := element_dict.get("url"):
            attachment = Attachment(
                content_type=mime, content_url=url, name=element_name
            )

        if not attachment:
            return

        await self.turn_context.send_activity(Activity(attachments=[attachment]))

    async def send_step(self, step_dict: StepDict):
        if not step_dict["type"] == "assistant_message":
            return

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

        if is_empty_output or not is_message:
            return
        else:
            reply = MessageFactory.text(step_dict["output"])
            enable_feedback = get_data_layer()
            if enable_feedback:
                current_run = context.current_run
                scorable_id = current_run.id if current_run else step_dict["id"]
                like_button = CardAction(
                    type=ActionTypes.message_back,
                    title="👍",
                    text="like",
                    value={"feedback": "like", "step_id": scorable_id},
                )
                dislike_button = CardAction(
                    type=ActionTypes.message_back,
                    title="👎",
                    text="dislike",
                    value={"feedback": "dislike", "step_id": scorable_id},
                )
                card = HeroCard(buttons=[like_button, dislike_button])
                attachment = Attachment(
                    content_type="application/vnd.microsoft.card.hero", content=card
                )
                reply.attachments = [attachment]

            await self.turn_context.send_activity(reply)

    async def update_step(self, step_dict: StepDict):
        if not step_dict["type"] == "assistant_message":
            return

        await self.send_step(step_dict)


adapter_settings = BotFrameworkAdapterSettings(
    app_id=os.environ.get("TEAMS_APP_ID"),
    app_password=os.environ.get("TEAMS_APP_PASSWORD"),
)
adapter = BotFrameworkAdapter(adapter_settings)


@trace
def init_teams_context(
    session: HTTPSession,
    turn_context: TurnContext,
) -> ChainlitContext:
    emitter = TeamsEmitter(session=session, turn_context=turn_context)
    context = ChainlitContext(session=session, emitter=emitter)
    context_var.set(context)
    user_session.set("teams_turn_context", turn_context)
    return context


users_by_teams_id: Dict[str, Union[User, PersistedUser]] = {}

USER_PREFIX = "teams_"


async def get_user(teams_user: ChannelAccount):
    if teams_user.id in users_by_teams_id:
        return users_by_teams_id[teams_user.id]

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

    users_by_teams_id[teams_user.id] = user

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

    return users_by_teams_id[teams_user.id]


async def download_teams_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_teams_files(
    session: HTTPSession, attachments: Optional[List[Attachment]] = None
):
    if not attachments:
        return []

    attachments = [
        attachment for attachment in attachments if isinstance(attachment.content, dict)
    ]
    download_coros = [
        download_teams_file(attachment.content.get("downloadUrl"))
        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].name
            mime_type = filetype.guess_mime(file_bytes) 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
    ]

    elements = [
        Element.from_dict(
            {
                "id": file["id"],
                "name": file["name"],
                "path": str(file["path"]),
                "chainlitKey": file["id"],
                "display": "inline",
                "type": Element.infer_type_from_mime(file["type"]),
            }
        )
        for file in files_dicts
    ]

    return elements


def clean_content(activity: Activity):
    return activity.text.strip()


async def process_teams_message(
    turn_context: TurnContext,
    thread_name: str,
):
    user = await get_user(turn_context.activity.from_property)

    thread_id = str(
        uuid.uuid5(
            uuid.NAMESPACE_DNS,
            str(
                turn_context.activity.conversation.id
                + datetime.today().strftime("%Y-%m-%d")
            ),
        )
    )

    text = clean_content(turn_context.activity)
    teams_files = turn_context.activity.attachments

    session_id = str(uuid.uuid4())

    session = HTTPSession(
        id=session_id,
        thread_id=thread_id,
        user=user,
        client_type="teams",
    )

    ctx = init_teams_context(
        session=session,
        turn_context=turn_context,
    )

    file_elements = await download_teams_files(session, teams_files)

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

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

    await msg.send()

    if on_message := config.code.on_message:
        await on_message(msg)

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

    if data_layer := get_data_layer():
        if isinstance(user, PersistedUser):
            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}")

    await ctx.session.delete()


async def handle_message(turn_context: TurnContext):
    if turn_context.activity.type == ActivityTypes.message:
        if (
            turn_context.activity.text == "like"
            or turn_context.activity.text == "dislike"
        ):
            feedback_value: Literal[0, 1] = (
                0 if turn_context.activity.text == "dislike" else 1
            )
            step_id = turn_context.activity.value.get("step_id")
            if data_layer := get_data_layer():
                await data_layer.upsert_feedback(
                    Feedback(forId=step_id, value=feedback_value)
                )
            updated_text = "👍" if turn_context.activity.text == "like" else "👎"
            # Update the existing message to remove the buttons
            updated_message = Activity(
                type=ActivityTypes.message,
                id=turn_context.activity.reply_to_id,
                text=updated_text,
                attachments=[],
            )
            await turn_context.update_activity(updated_message)
        else:
            # Send typing activity
            typing_activity = Activity(
                type=ActivityTypes.typing,
                from_property=turn_context.activity.recipient,
                recipient=turn_context.activity.from_property,
                conversation=turn_context.activity.conversation,
            )
            await turn_context.send_activity(typing_activity)
            thread_name = f"{turn_context.activity.from_property.name} Teams DM {datetime.today().strftime('%Y-%m-%d')}"
            await process_teams_message(turn_context, thread_name)


async def on_turn(turn_context: TurnContext):
    await handle_message(turn_context)


# Create the main bot class
class TeamsBot:
    async def on_turn(self, turn_context: TurnContext):
        await on_turn(turn_context)


# Create the bot instance
bot = TeamsBot()
