import asyncio
import json
import logging
from datetime import datetime, timezone
from typing import Any, cast

from langchain_core.messages import (
    AIMessageChunk,
    BaseMessage,
    HumanMessage,
    TextContentBlock,
    ToolCall,
    ToolMessage,
)
from pydantic import ValidationError
from uipath.core.chat import (
    UiPathConversationContentPartChunkEvent,
    UiPathConversationContentPartEndEvent,
    UiPathConversationContentPartEvent,
    UiPathConversationContentPartStartEvent,
    UiPathConversationMessage,
    UiPathConversationMessageEndEvent,
    UiPathConversationMessageEvent,
    UiPathConversationMessageStartEvent,
    UiPathConversationToolCallEndEvent,
    UiPathConversationToolCallEvent,
    UiPathConversationToolCallStartEvent,
    UiPathInlineValue,
)
from uipath.runtime import UiPathRuntimeStorageProtocol

logger = logging.getLogger(__name__)

STORAGE_NAMESPACE_EVENT_MAPPER = "chat-event-mapper"
STORAGE_KEY_TOOL_CALL_ID_TO_MESSAGE_ID_MAP = "tool_call_map"


class UiPathChatMessagesMapper:
    """Stateful mapper that converts LangChain messages to UiPath message events.

    Maintains state across multiple message conversions to properly track:
    - The AI message ID associated with each tool call for proper correlation with ToolMessage
    """

    def __init__(self, runtime_id: str, storage: UiPathRuntimeStorageProtocol | None):
        """Initialize the mapper with empty state."""
        self.runtime_id = runtime_id
        self.storage = storage
        self.current_message: AIMessageChunk
        self.seen_message_ids: set[str] = set()
        self._storage_lock = asyncio.Lock()

    def _extract_text(self, content: Any) -> str:
        """Normalize LangGraph message.content to plain text."""
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            return "".join(
                part.get("text", "")
                for part in content
                if isinstance(part, dict) and part.get("type") == "text"
            )
        return str(content or "")

    def map_messages(self, messages: list[Any]) -> list[Any]:
        """Normalize any 'messages' list into LangChain messages.

        - If already BaseMessage instances: return as-is.
        - If UiPathConversationMessage: convert to HumanMessage.
        """
        if not isinstance(messages, list):
            raise TypeError("messages must be a list")

        if not messages:
            return []

        first = messages[0]

        # Case 1: already LangChain messages
        if isinstance(first, BaseMessage):
            return cast(list[BaseMessage], messages)

        # Case 2: UiPath messages -> convert to HumanMessage
        if isinstance(first, UiPathConversationMessage):
            if not all(isinstance(m, UiPathConversationMessage) for m in messages):
                raise TypeError("Mixed message types not supported")
            return self._map_messages_internal(
                cast(list[UiPathConversationMessage], messages)
            )

        # Case3: List[dict] -> parse to List[UiPathConversationMessage]
        if isinstance(first, dict):
            try:
                parsed_messages = [
                    UiPathConversationMessage.model_validate(message)
                    for message in messages
                ]
                return self._map_messages_internal(parsed_messages)
            except ValidationError:
                pass

        # Fallback: unknown type – just pass through
        return messages

    def _map_messages_internal(
        self, messages: list[UiPathConversationMessage]
    ) -> list[HumanMessage]:
        """
        Converts a UiPathConversationMessage into a list of HumanMessages for LangGraph.
        Supports multimodal content parts (text, external content) and preserves metadata.
        """
        human_messages: list[HumanMessage] = []

        for uipath_msg in messages:
            # Loop over each content part
            if uipath_msg.content_parts:
                for part in uipath_msg.content_parts:
                    data = part.data
                    content = ""
                    metadata: dict[str, Any] = {
                        "message_id": uipath_msg.message_id,
                        "content_part_id": part.content_part_id,
                        "mime_type": part.mime_type,
                        "created_at": uipath_msg.created_at,
                        "updated_at": uipath_msg.updated_at,
                    }

                    if isinstance(data, UiPathInlineValue):
                        content = str(data.inline)

                    # Append a HumanMessage for this content part
                    human_messages.append(
                        HumanMessage(content=content, metadata=metadata)
                    )

            # Handle the case where there are no content parts
            else:
                metadata = {
                    "message_id": uipath_msg.message_id,
                    "role": uipath_msg.role,
                    "created_at": uipath_msg.created_at,
                    "updated_at": uipath_msg.updated_at,
                }
                human_messages.append(HumanMessage(content="", metadata=metadata))

        return human_messages

    async def map_event(
        self,
        message: BaseMessage,
    ) -> list[UiPathConversationMessageEvent] | None:
        """Convert LangGraph BaseMessage (chunk or full) into a UiPathConversationMessageEvent.

        Args:
            message: The LangChain message to convert

        Returns:
            A UiPathConversationMessageEvent if the message should be emitted, None otherwise.
        """
        # --- Streaming AIMessageChunk ---
        if isinstance(message, AIMessageChunk):
            return await self.map_ai_message_chunk_to_events(message)

        # --- ToolMessage ---
        if isinstance(message, ToolMessage):
            return await self.map_tool_message_to_events(message)

        # Don't send events for system or user messages. Agent messages are handled above.
        return None

    def get_timestamp(self):
        """Format current time as ISO 8601 UTC with milliseconds: 2025-01-04T10:30:00.123Z"""
        return (
            datetime.now(timezone.utc)
            .isoformat(timespec="milliseconds")
            .replace("+00:00", "Z")
        )

    def get_content_part_id(self, message_id: str) -> str:
        return f"chunk-{message_id}-0"

    async def map_ai_message_chunk_to_events(
        self, message: AIMessageChunk
    ) -> list[UiPathConversationMessageEvent]:
        if message.id is None:  # Should we throw instead?
            return []

        events: list[UiPathConversationMessageEvent] = []

        # For every new message_id, start a new message
        if message.id not in self.seen_message_ids:
            self.current_message = message
            self.seen_message_ids.add(message.id)
            events.append(self.map_to_message_start_event(message.id))

        if message.content_blocks:
            # Generate events for each chunk
            for block in message.content_blocks:
                block_type = block.get("type")
                match block_type:
                    case "text":
                        events.append(
                            self.map_chunk_to_content_part_chunk_event(
                                message.id, cast(TextContentBlock, block)
                            )
                        )
                    case "tool_call_chunk":
                        # Accumulate the message chunk
                        self.current_message = self.current_message + message

        elif isinstance(message.content, str) and message.content:
            # Fallback: raw string content on the chunk (rare when using content_blocks)
            events.append(
                self.map_content_to_content_part_chunk_event(
                    message.id, message.content
                )
            )

        # Check if this is the last chunk by examining chunk_position, send end message event only if there are no pending tool calls
        if message.chunk_position == "last":
            if (
                self.current_message.tool_calls is not None
                and len(self.current_message.tool_calls) > 0
            ):
                events.extend(
                    await self.map_current_message_to_start_tool_call_events()
                )
            else:
                events.append(self.map_to_message_end_event(message.id))

        return events

    async def map_current_message_to_start_tool_call_events(self):
        events: list[UiPathConversationMessageEvent] = []
        if (
            self.current_message
            and self.current_message.id is not None
            and self.current_message.tool_calls
        ):
            async with self._storage_lock:
                if self.storage is not None:
                    tool_call_id_to_message_id_map: dict[
                        str, str
                    ] = await self.storage.get_value(
                        self.runtime_id,
                        STORAGE_NAMESPACE_EVENT_MAPPER,
                        STORAGE_KEY_TOOL_CALL_ID_TO_MESSAGE_ID_MAP,
                    )

                    if tool_call_id_to_message_id_map is None:
                        tool_call_id_to_message_id_map = {}
                else:
                    tool_call_id_to_message_id_map = {}

                for tool_call in self.current_message.tool_calls:
                    tool_call_id = tool_call["id"]
                    if tool_call_id is not None:
                        tool_call_id_to_message_id_map[tool_call_id] = (
                            self.current_message.id
                        )
                        events.append(
                            self.map_tool_call_to_tool_call_start_event(
                                self.current_message.id, tool_call
                            )
                        )

                if self.storage is not None:
                    await self.storage.set_value(
                        self.runtime_id,
                        STORAGE_NAMESPACE_EVENT_MAPPER,
                        STORAGE_KEY_TOOL_CALL_ID_TO_MESSAGE_ID_MAP,
                        tool_call_id_to_message_id_map,
                    )

        return events

    async def map_tool_message_to_events(
        self, message: ToolMessage
    ) -> list[UiPathConversationMessageEvent]:
        # Look up the AI message ID using the tool_call_id
        message_id, is_last_tool_call = await self.get_message_id_for_tool_call(
            message.tool_call_id
        )
        if message_id is None:
            logger.warning(
                f"Tool message {message.tool_call_id} has no associated AI message ID. Skipping."
            )
            return []

        content_value: Any = message.content
        if isinstance(content_value, str):
            try:
                content_value = json.loads(content_value)
            except (json.JSONDecodeError, TypeError):
                # Keep as string if not valid JSON
                pass

        events = [
            UiPathConversationMessageEvent(
                message_id=message_id,
                tool_call=UiPathConversationToolCallEvent(
                    tool_call_id=message.tool_call_id,
                    end=UiPathConversationToolCallEndEvent(
                        timestamp=self.get_timestamp(),
                        output=content_value,
                    ),
                ),
            )
        ]

        if is_last_tool_call:
            events.append(self.map_to_message_end_event(message_id))

        return events

    async def get_message_id_for_tool_call(
        self, tool_call_id: str
    ) -> tuple[str | None, bool]:
        if self.storage is None:
            logger.error(
                f"attempt to lookup tool call id {tool_call_id} when no storage provided"
            )
            return None, False

        async with self._storage_lock:
            tool_call_id_to_message_id_map: dict[
                str, str
            ] = await self.storage.get_value(
                self.runtime_id,
                STORAGE_NAMESPACE_EVENT_MAPPER,
                STORAGE_KEY_TOOL_CALL_ID_TO_MESSAGE_ID_MAP,
            )

            if tool_call_id_to_message_id_map is None:
                logger.error(
                    f"attempt to lookup tool call id {tool_call_id} when no map present in storage"
                )
                return None, False

            message_id = tool_call_id_to_message_id_map.get(tool_call_id)
            if message_id is None:
                logger.error(
                    f"tool call to message map does not contain tool call id {tool_call_id}"
                )
                return None, False

            del tool_call_id_to_message_id_map[tool_call_id]

            await self.storage.set_value(
                self.runtime_id,
                STORAGE_NAMESPACE_EVENT_MAPPER,
                STORAGE_KEY_TOOL_CALL_ID_TO_MESSAGE_ID_MAP,
                tool_call_id_to_message_id_map,
            )

            is_last = message_id not in tool_call_id_to_message_id_map.values()

        return message_id, is_last

    def map_tool_call_to_tool_call_start_event(
        self, message_id: str, tool_call: ToolCall
    ) -> UiPathConversationMessageEvent:
        return UiPathConversationMessageEvent(
            message_id=message_id,
            tool_call=UiPathConversationToolCallEvent(
                tool_call_id=tool_call["id"],
                start=UiPathConversationToolCallStartEvent(
                    tool_name=tool_call["name"],
                    timestamp=self.get_timestamp(),
                    input=tool_call["args"],
                ),
            ),
        )

    def map_chunk_to_content_part_chunk_event(
        self, message_id: str, block: TextContentBlock
    ) -> UiPathConversationMessageEvent:
        text = block["text"]
        return UiPathConversationMessageEvent(
            message_id=message_id,
            content_part=UiPathConversationContentPartEvent(
                content_part_id=self.get_content_part_id(message_id),
                chunk=UiPathConversationContentPartChunkEvent(
                    data=text,
                ),
            ),
        )

    def map_content_to_content_part_chunk_event(
        self, message_id: str, content: str
    ) -> UiPathConversationMessageEvent:
        return UiPathConversationMessageEvent(
            message_id=message_id,
            content_part=UiPathConversationContentPartEvent(
                content_part_id=self.get_content_part_id(message_id),
                chunk=UiPathConversationContentPartChunkEvent(
                    data=content,
                ),
            ),
        )

    def map_to_message_start_event(
        self, message_id: str
    ) -> UiPathConversationMessageEvent:
        return UiPathConversationMessageEvent(
            message_id=message_id,
            start=UiPathConversationMessageStartEvent(
                role="assistant", timestamp=self.get_timestamp()
            ),
            content_part=UiPathConversationContentPartEvent(
                content_part_id=self.get_content_part_id(message_id),
                start=UiPathConversationContentPartStartEvent(mime_type="text/plain"),
            ),
        )

    def map_to_message_end_event(
        self, message_id: str
    ) -> UiPathConversationMessageEvent:
        return UiPathConversationMessageEvent(
            message_id=message_id,
            end=UiPathConversationMessageEndEvent(),
            content_part=UiPathConversationContentPartEvent(
                content_part_id=self.get_content_part_id(message_id),
                end=UiPathConversationContentPartEndEvent(),
            ),
        )


__all__ = ["UiPathChatMessagesMapper"]
