from __future__ import annotations

import asyncio
from typing import Any, Coroutine
from uuid import uuid4

from terminaluse.lib.client import AsyncTerminalUse
from terminaluse.lib.core.tracing.tracer import AsyncTracer
from terminaluse.lib.types.message_parts import AgentMessagePart, DataPart, TextPart
from terminaluse.lib.utils.logging import make_logger
from terminaluse.lib.utils.sdk_version import DEFAULT_TERMINAL_USE_VERSION, SDK_TYPE_TERMINAL_USE
from terminaluse.lib.utils.temporal import heartbeat_if_in_workflow
from terminaluse.types.ui_message import UiMessage

logger = make_logger(__name__)


class MessagesService:
    """
    Service for managing task messages.

    v2 Architecture: Uses raw_events API for message creation/streaming.
    Messages are published as UIMessage format via TerminalUseAdapter.
    """

    def __init__(
        self,
        terminaluse_client: AsyncTerminalUse,
        tracer: AsyncTracer,
    ):
        self._terminaluse_client = terminaluse_client
        self._tracer = tracer

    async def _publish_ui_message(
        self,
        task_id: str,
        role: str,
        parts: list[dict[str, Any]],
        message_id: str | None = None,
        idempotency_key: str | None = None,
    ) -> None:
        """
        Publish a UIMessage directly to the backend via raw_events API.

        Args:
            task_id: The task ID
            role: Message role ("assistant", "user", etc.)
            parts: List of message parts (text, tool, data, etc.)
            message_id: Optional message ID (generated if not provided)
        """
        ui_message = {
            "id": message_id or str(uuid4()),
            "role": role,
            "parts": parts,
        }
        try:
            await self._terminaluse_client.raw_events.publish_raw_event(
                task_id=task_id,
                sdk_type=SDK_TYPE_TERMINAL_USE,
                sdk_version=DEFAULT_TERMINAL_USE_VERSION,
                raw_event=ui_message,
                idempotency_key=idempotency_key or str(uuid4()),
            )
        except Exception as e:
            logger.error(
                f"Failed to publish UI message: task_id={task_id}, error={e}",
                exc_info=True,
            )
            raise

    async def send(
        self,
        task_id: str,
        content: AgentMessagePart,
        emit_updates: bool = True,
        trace_id: str | None = None,
        parent_span_id: str | None = None,
    ) -> None:
        """
        Send a message for a task via v2 raw_events API.

        Creates a UIMessage and publishes it directly to the backend.
        """
        trace = self._tracer.trace(trace_id)
        async with trace.span(
            parent_id=parent_span_id,
            name="messages.send",
            input={"task_id": task_id, "message": content},
        ) as span:
            heartbeat_if_in_workflow("create message")

            # Generate message ID client-side
            message_id = str(uuid4())

            # Convert content to UIMessage parts
            parts = self._content_to_parts(content)
            role = self._get_role_from_content(content)

            # Publish via v2 raw_events
            await self._publish_ui_message(
                task_id=task_id,
                role=role,
                parts=parts,
                message_id=message_id,
            )

            if span:
                span.output = {"task_id": task_id, "message_id": message_id}

    async def update_message(
        self,
        task_id: str,
        message_id: str,
        content: AgentMessagePart,
        trace_id: str | None = None,
        parent_span_id: str | None = None,
    ) -> None:
        """
        Update a message via v2 raw_events API.

        Note: With v2, updates are published as new UIMessages with the same ID.
        The backend handles deduplication/update logic.
        """
        trace = self._tracer.trace(trace_id)
        async with trace.span(
            parent_id=parent_span_id,
            name="update_message",
            input={
                "task_id": task_id,
                "message_id": message_id,
                "message": content,
            },
        ) as span:
            heartbeat_if_in_workflow("update message")

            # Convert content to UIMessage parts
            parts = self._content_to_parts(content)
            role = self._get_role_from_content(content)

            # Publish update via v2 raw_events with existing message ID
            await self._publish_ui_message(
                task_id=task_id,
                role=role,
                parts=parts,
                message_id=message_id,
            )

            if span:
                span.output = {"task_id": task_id, "message_id": message_id}

    async def send_batch(
        self,
        task_id: str,
        contents: list[AgentMessagePart],
        emit_updates: bool = True,
        trace_id: str | None = None,
        parent_span_id: str | None = None,
    ) -> None:
        """
        Send a batch of messages via v2 raw_events API.

        Publishes each message as a separate UIMessage.
        """
        trace = self._tracer.trace(trace_id)
        async with trace.span(
            parent_id=parent_span_id,
            name="messages.send_batch",
            input={"task_id": task_id, "messages": contents},
        ) as span:
            heartbeat_if_in_workflow("create messages batch")

            publish_tasks: list[Coroutine[Any, Any, None]] = []
            message_ids: list[str] = []

            for content in contents:
                message_id = str(uuid4())
                message_ids.append(message_id)
                parts = self._content_to_parts(content)
                role = self._get_role_from_content(content)

                publish_tasks.append(
                    self._publish_ui_message(
                        task_id=task_id,
                        role=role,
                        parts=parts,
                        message_id=message_id,
                    )
                )

            # Publish all messages in parallel
            await asyncio.gather(*publish_tasks)

            if span:
                span.output = {"task_id": task_id, "message_ids": message_ids}

    async def update_messages_batch(
        self,
        task_id: str,
        updates: dict[str, AgentMessagePart],
        trace_id: str | None = None,
        parent_span_id: str | None = None,
    ) -> None:
        """
        Update a batch of messages via v2 raw_events API.

        Publishes each update as a UIMessage with the existing message ID.
        """
        trace = self._tracer.trace(trace_id)
        async with trace.span(
            parent_id=parent_span_id,
            name="update_messages_batch",
            input={"task_id": task_id, "updates": updates},
        ) as span:
            heartbeat_if_in_workflow("update messages batch")

            publish_tasks: list[Coroutine[Any, Any, None]] = []
            message_ids: list[str] = []

            for message_id, content in updates.items():
                message_ids.append(message_id)
                parts = self._content_to_parts(content)
                role = self._get_role_from_content(content)

                publish_tasks.append(
                    self._publish_ui_message(
                        task_id=task_id,
                        role=role,
                        parts=parts,
                        message_id=message_id,
                    )
                )

            # Publish all updates in parallel
            await asyncio.gather(*publish_tasks)

            if span:
                span.output = {"task_id": task_id, "message_ids": message_ids}

    def _content_to_parts(self, content: AgentMessagePart) -> list[dict[str, Any]]:
        """Convert AgentMessagePart to UIMessage parts."""
        if isinstance(content, TextPart):
            return [content.model_dump()]
        elif isinstance(content, DataPart):
            return [content.model_dump()]
        # Fallback: try to serialize as text
        if hasattr(content, "text"):
            return [{"type": "text", "text": str(content.text)}]
        return [{"type": "text", "text": ""}]

    def _get_role_from_content(self, content: AgentMessagePart) -> str:
        """Get role from content.

        AgentMessagePart (TextPart, DataPart) always returns "assistant"
        since these are sent by agents.
        """
        return "assistant"

    async def list_messages(
        self,
        task_id: str,
        limit: int | None = None,
        trace_id: str | None = None,
        parent_span_id: str | None = None,
    ) -> list[UiMessage]:
        """
        List messages for a task via v2 API.

        Returns UIMessages directly from the backend.
        """
        trace = self._tracer.trace(trace_id)
        async with trace.span(
            parent_id=parent_span_id,
            name="list_messages",
            input={"task_id": task_id, "limit": limit},
        ) as span:
            heartbeat_if_in_workflow("list messages")
            # Use v2 API pager and optionally exhaust all pages.
            pager = await self._terminaluse_client.messages_v2.list(
                task_id=task_id,
                limit=limit,
            )
            if limit is None:
                messages = [item async for item in pager]
            else:
                messages = list(pager.items or [])

            if span:
                span.output = [ui_msg.model_dump() for ui_msg in messages]
            return messages
