import asyncio
import os
import re
import uuid
from functools import partial
from typing import Dict, List, Optional, Union

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 slack_bolt.adapter.fastapi.async_handler import AsyncSlackRequestHandler
from slack_bolt.async_app import AsyncApp


class SlackEmitter(BaseChainlitEmitter):
    def __init__(
        self,
        session: HTTPSession,
        app: AsyncApp,
        channel_id: str,
        say,
        enabled=False,
        thread_ts: Optional[str] = None,
    ):
        super().__init__(session)
        self.app = app
        self.channel_id = channel_id
        self.say = say
        self.enabled = enabled
        self.thread_ts = thread_ts

    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[bytes, str]] = None

        if persisted_file:
            file = str(persisted_file["path"])
        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 = response.content

        if not file:
            return

        await self.app.client.files_upload_v2(
            channel=self.channel_id,
            thread_ts=self.thread_ts,
            file=file,
            title=element_dict.get("name"),
        )

    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

        enable_feedback = not step_dict.get("disableFeedback") and get_data_layer()
        blocks: List[Dict] = [
            {
                "type": "section",
                "text": {"type": "mrkdwn", "text": step_dict["output"]},
            }
        ]
        if enable_feedback:
            blocks.append(
                {
                    "type": "actions",
                    "elements": [
                        {
                            "action_id": "thumbdown",
                            "type": "button",
                            "text": {
                                "type": "plain_text",
                                "emoji": True,
                                "text": ":thumbsdown:",
                            },
                            "value": step_dict.get("id"),
                        },
                        {
                            "action_id": "thumbup",
                            "type": "button",
                            "text": {
                                "type": "plain_text",
                                "emoji": True,
                                "text": ":thumbsup:",
                            },
                            "value": step_dict.get("id"),
                        },
                    ],
                }
            )
        await self.say(
            text=step_dict["output"], blocks=blocks, thread_ts=self.thread_ts
        )

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

        await self.send_step(step_dict)


slack_app = AsyncApp(
    token=os.environ.get("SLACK_BOT_TOKEN"),
    signing_secret=os.environ.get("SLACK_SIGNING_SECRET"),
)


@trace
def init_slack_context(
    session: HTTPSession,
    slack_channel_id: str,
    event,
    say,
    thread_ts: Optional[str] = None,
) -> ChainlitContext:
    emitter = SlackEmitter(
        session=session,
        app=slack_app,
        channel_id=slack_channel_id,
        say=say,
        thread_ts=thread_ts,
    )
    context = ChainlitContext(session=session, emitter=emitter)
    context_var.set(context)
    user_session.set("slack_event", event)
    user_session.set(
        "fetch_slack_message_history",
        partial(
            fetch_message_history, channel_id=slack_channel_id, thread_ts=thread_ts
        ),
    )
    return context


slack_app_handler = AsyncSlackRequestHandler(slack_app)

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

USER_PREFIX = "slack_"


def clean_content(message: str):
    cleaned_text = re.sub(r"<@[\w]+>", "", message).strip()
    return cleaned_text


async def get_user(slack_user_id: str):
    if slack_user_id in users_by_slack_id:
        return users_by_slack_id[slack_user_id]

    slack_user = await slack_app.client.users_info(user=slack_user_id)
    slack_user_profile = slack_user["user"]["profile"]

    user_identifier = slack_user_profile.get("email") or slack_user_id
    user = User(identifier=USER_PREFIX + user_identifier, metadata=slack_user_profile)

    users_by_slack_id[slack_user_id] = user

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

    return users_by_slack_id[slack_user_id]


async def fetch_message_history(
    channel_id: str, thread_ts: Optional[str] = None, limit=30
):
    if not thread_ts:
        result = await slack_app.client.conversations_history(
            channel=channel_id, limit=limit
        )
    else:
        result = await slack_app.client.conversations_replies(
            channel=channel_id, ts=thread_ts, limit=limit
        )
    if result["ok"]:
        messages = result["messages"]
        return messages
    else:
        raise Exception(f"Failed to fetch messages: {result['error']}")


async def download_slack_file(url, token):
    headers = {"Authorization": f"Bearer {token}"}
    async with httpx.AsyncClient() as client:
        response = await client.get(url, headers=headers)
        if response.status_code == 200:
            return response.content
        else:
            return None


async def download_slack_files(session: HTTPSession, files, token):
    download_coros = [
        download_slack_file(file.get("url_private"), token) for file in files
    ]
    file_bytes_list = await asyncio.gather(*download_coros)
    file_refs = []
    for idx, file_bytes in enumerate(file_bytes_list):
        if file_bytes:
            name = files[idx].get("name")
            mime_type = files[idx].get("mimetype")
            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


async def process_slack_message(
    event,
    say,
    thread_name: Optional[str] = None,
    bind_thread_to_user=False,
    thread_ts: Optional[str] = None,
):
    user = await get_user(event["user"])

    channel_id = event["channel"]
    thread_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, thread_ts or channel_id))

    text = event.get("text")
    slack_files = event.get("files", [])

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

    ctx = init_slack_context(
        session=session,
        slack_channel_id=channel_id,
        event=event,
        say=say,
        thread_ts=thread_ts,
    )

    file_elements = await download_slack_files(
        session, slack_files, slack_app.client.token
    )

    msg = Message(
        content=clean_content(text),
        elements=file_elements,
        type="user_message",
        author=user.metadata.get("real_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:
        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 or msg.content,
                metadata=ctx.session.to_persistable(),
                user_id=user_id,
            )
        except Exception as e:
            logger.error(f"Error updating thread: {e}")

    ctx.session.delete()


@slack_app.event("app_home_opened")
async def handle_app_home_opened(event, say):
    pass


@slack_app.event("app_mention")
async def handle_app_mentions(event, say):
    thread_ts = event.get("thread_ts", event["ts"])
    await process_slack_message(event, say, thread_ts=thread_ts)


@slack_app.event("message")
async def handle_message(message, say):
    user = await get_user(message["user"])
    thread_name = f"{user.identifier} Slack DM"
    await process_slack_message(message, say, thread_name)


@slack_app.block_action("thumbdown")
async def thumb_down(ack, context, body):
    await ack()
    step_id = body["actions"][0]["value"]

    if data_layer := get_data_layer():
        thread_id = context_var.get().session.thread_id
        feedback = Feedback(forId=step_id, threadId=thread_id, value=0)
        await data_layer.upsert_feedback(feedback)

    text = body["message"]["text"]
    blocks = body["message"]["blocks"]
    updated_blocks = [block for block in blocks if block["type"] != "actions"]
    updated_blocks.append(
        {
            "type": "section",
            "text": {"type": "mrkdwn", "text": ":thumbsdown: Feedback received."},
        }
    )
    await context.client.chat_update(
        channel=body["channel"]["id"],
        ts=body["container"]["message_ts"],
        text=text,
        blocks=updated_blocks,
    )


@slack_app.block_action("thumbup")
async def thumb_up(ack, context, body):
    await ack()
    step_id = body["actions"][0]["value"]

    if data_layer := get_data_layer():
        thread_id = context_var.get().session.thread_id
        feedback = Feedback(forId=step_id, threadId=thread_id, value=1)
        await data_layer.upsert_feedback(feedback)

    text = body["message"]["text"]
    blocks = body["message"]["blocks"]
    updated_blocks = [block for block in blocks if block["type"] != "actions"]
    updated_blocks.append(
        {
            "type": "section",
            "text": {"type": "mrkdwn", "text": ":thumbsup: Feedback received."},
        }
    )
    await context.client.chat_update(
        channel=body["channel"]["id"],
        ts=body["container"]["message_ts"],
        text=text,
        blocks=updated_blocks,
    )
