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.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 import TaskMessageContent
from terminaluse.types.task_message_content import (
    TaskMessageContent_Data,
    TaskMessageContent_Reasoning,
    TaskMessageContent_Text,
    TaskMessageContent_ToolRequest,
    TaskMessageContent_ToolResponse,
)
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,
    ) -> 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,
            )
        except Exception as e:
            logger.error(f"Failed to publish UI message: task_id={task_id}, error={e}")

    async def send(
        self,
        task_id: str,
        content: TaskMessageContent,
        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: TaskMessageContent,
        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[TaskMessageContent],
        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, TaskMessageContent],
        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: TaskMessageContent) -> list[dict[str, Any]]:
        """Convert TaskMessageContent to UIMessage parts."""
        if isinstance(content, TaskMessageContent_Text):
            return [{"type": "text", "text": content.content or ""}]
        elif isinstance(content, TaskMessageContent_Data):
            return [{"type": "data", "data": content.data}]
        elif isinstance(content, TaskMessageContent_ToolRequest):
            return [
                {
                    "type": "tool",
                    "tool_name": content.name,
                    "tool_call_id": content.tool_call_id,
                    "arguments": content.arguments,
                }
            ]
        elif isinstance(content, TaskMessageContent_ToolResponse):
            return [
                {
                    "type": "tool",
                    "tool_name": content.name,
                    "tool_call_id": content.tool_call_id,
                    "content": content.content,
                }
            ]
        elif isinstance(content, TaskMessageContent_Reasoning):
            return [
                {
                    "type": "reasoning",
                    "text": "\n".join(content.summary or []),
                }
            ]
        # Fallback: try to serialize as text
        if hasattr(content, "content"):
            return [{"type": "text", "text": str(content.content)}]
        return [{"type": "text", "text": ""}]

    def _get_role_from_content(self, content: TaskMessageContent) -> str:
        """Get role from content author field."""
        if hasattr(content, "author"):
            author = content.author
            if author == "agent":
                return "assistant"
            return str(author)
        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 which returns UIMessages
            response = await self._terminaluse_client.messages_v2.list(
                task_id=task_id,
                limit=limit,
            )
            if span:
                span.output = [ui_msg.model_dump() for ui_msg in response.data]
            return list(response.data)
