# ruff: noqa: I001
from __future__ import annotations
from dataclasses import asdict, is_dataclass
from datetime import timedelta
from typing import Any, Union

from temporalio.common import RetryPolicy

from terminaluse.lib.adk.utils._modules.client import create_async_terminaluse_client
from terminaluse.lib.core.services.adk.messages import MessagesService
from terminaluse.lib.core.temporal.activities.activity_helpers import ActivityHelpers
from terminaluse.lib.core.temporal.activities.adk.messages_activities import (
    CreateMessageParams,
    CreateMessagesBatchParams,
    ListMessagesParams,
    MessagesActivityName,
    UpdateMessageParams,
    UpdateMessagesBatchParams,
)
from terminaluse.lib.core.tracing.tracer import AsyncTracer
from terminaluse.lib.utils.logging import make_logger
from terminaluse.lib.utils.temporal import in_temporal_workflow
from terminaluse.lib.utils.sdk_version import SDK_TYPE_CLAUDE, get_claude_sdk_version
from terminaluse.types import TaskMessageContent
from terminaluse.types.ui_message import UiMessage
from claude_agent_sdk.types import Message as ClaudeMessage

logger = make_logger(__name__)
DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)


class MessagesModule:
    """
    Module for managing task messages.

    v2 Architecture: Claude messages are forwarded directly to the backend via raw_events.
    The backend's ClaudeAdapter handles all transformation and streaming.
    Platform messages are sent as UIMessages via TerminalUseAdapter.

    All send/update methods return None - the backend handles message creation asynchronously.
    Use list() to retrieve messages.
    """

    def __init__(self, messages_service: MessagesService | None = None):
        if messages_service is None:
            self._terminaluse_client = create_async_terminaluse_client()
            tracer = AsyncTracer(self._terminaluse_client)
            self._messages_service = MessagesService(terminaluse_client=self._terminaluse_client, tracer=tracer)
        else:
            self._messages_service = messages_service
            self._terminaluse_client = messages_service._terminaluse_client

    async def _publish_claude_event(
        self,
        task_id: str,
        raw_event: dict[str, Any],
    ) -> None:
        """
        Publish a Claude SDK event directly to the backend via raw_events API.

        Args:
            task_id: The task ID
            raw_event: The raw Claude SDK event
        """
        event_type = raw_event.get("type", "unknown")
        sdk_version = get_claude_sdk_version()

        try:
            await self._terminaluse_client.raw_events.publish_raw_event(
                task_id=task_id,
                sdk_type=SDK_TYPE_CLAUDE,
                sdk_version=sdk_version,
                raw_event=raw_event,
            )
        except Exception as e:
            # Log with full context and stack trace for debugging
            logger.error(
                f"Failed to publish Claude event: task_id={task_id}, "
                f"event_type={event_type}, sdk_version={sdk_version}, error={e}",
                exc_info=True,
            )

    async def send(
        self,
        task_id: str,
        content: Union[TaskMessageContent, ClaudeMessage],
        emit_updates: bool = True,
        store_raw_claude_message: bool = False,
        trace_id: str | None = None,
        parent_span_id: str | None = None,
        start_to_close_timeout: timedelta = timedelta(seconds=5),
        heartbeat_timeout: timedelta = timedelta(seconds=5),
        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
    ) -> None:
        """Send a message. Accepts platform TaskMessageContent or Claude SDK message types.

        Messages are sent asynchronously to the backend. Use list() to retrieve messages.
        """
        # Check if this is a Claude SDK message type
        if isinstance(content, ClaudeMessage):
            await self._handle_claude_message(
                task_id=task_id,
                message=content,
            )
            return

        # Standard platform content type handling
        await self._send_platform_message(
            task_id=task_id,
            content=content,
            emit_updates=emit_updates,
            trace_id=trace_id,
            parent_span_id=parent_span_id,
            start_to_close_timeout=start_to_close_timeout,
            heartbeat_timeout=heartbeat_timeout,
            retry_policy=retry_policy,
        )

    async def _send_platform_message(
        self,
        task_id: str,
        content: TaskMessageContent,
        emit_updates: bool = True,
        trace_id: str | None = None,
        parent_span_id: str | None = None,
        start_to_close_timeout: timedelta = timedelta(seconds=5),
        heartbeat_timeout: timedelta = timedelta(seconds=5),
        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
    ) -> None:
        """
        Send a new message for a task.

        Args:
            task_id: The ID of the task.
            content: The message content to send.
            emit_updates: Whether to emit updates.
            trace_id: The trace ID for tracing.
            parent_span_id: The parent span ID for tracing.
            start_to_close_timeout: The start to close timeout.
            heartbeat_timeout: The heartbeat timeout.
            retry_policy: The retry policy.
        """
        params = CreateMessageParams(
            trace_id=trace_id,
            parent_span_id=parent_span_id,
            task_id=task_id,
            content=content,
            emit_updates=emit_updates,
        )
        if in_temporal_workflow():
            await ActivityHelpers.execute_activity(
                activity_name=MessagesActivityName.CREATE_MESSAGE,
                request=params,
                response_type=type(None),
                start_to_close_timeout=start_to_close_timeout,
                retry_policy=retry_policy,
                heartbeat_timeout=heartbeat_timeout,
            )
            return
        await self._messages_service.send(task_id=task_id, content=content, emit_updates=emit_updates)

    async def _handle_claude_message(
        self,
        task_id: str,
        message: ClaudeMessage,
    ) -> None:
        """
        Forward Claude SDK messages directly to the backend via raw_events.

        v2 Architecture: All Claude messages are published as raw events.
        The backend's ClaudeAdapter handles transformation, streaming, and storage.
        """
        # Convert Claude message to raw event dict
        raw_event = self._claude_message_to_raw_event(message)

        # Publish to backend via raw_events
        await self._publish_claude_event(
            task_id=task_id,
            raw_event=raw_event,
        )

    def _claude_message_to_raw_event(self, message: ClaudeMessage) -> dict[str, Any]:
        """
        Convert a Claude SDK message to raw event format for the backend.

        Uses dataclasses.asdict() for recursive serialization of nested dataclasses
        (e.g., UserMessage containing ToolResultBlock).

        The backend's ClaudeAdapter uses duck typing to infer content block types
        from their fields (e.g., 'tool_use_id' → tool_result).
        """
        message_type = type(message).__name__

        # Serialize message data
        # Claude SDK types are dataclasses - use asdict() for recursive serialization
        if hasattr(message, "model_dump"):
            data = message.model_dump(mode="json")
        elif is_dataclass(message):
            data = asdict(message)
        else:
            # Fallback for non-dataclass types
            data = vars(message)

        return {
            "type": message_type,
            "data": data,
        }

    async def update(
        self,
        task_id: str,
        message_id: str,
        content: TaskMessageContent,
        trace_id: str | None = None,
        parent_span_id: str | None = None,
        start_to_close_timeout: timedelta = timedelta(seconds=5),
        heartbeat_timeout: timedelta = timedelta(seconds=5),
        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
    ) -> None:
        """
        Update a message for a task.

        Args:
            task_id: The ID of the task.
            message_id: The ID of the message.
            content: The new message content.
            trace_id: The trace ID for tracing.
            parent_span_id: The parent span ID for tracing.
            start_to_close_timeout: The start to close timeout.
            heartbeat_timeout: The heartbeat timeout.
            retry_policy: The retry policy.
        """
        params = UpdateMessageParams(
            task_id=task_id, message_id=message_id, content=content, trace_id=trace_id, parent_span_id=parent_span_id
        )
        if in_temporal_workflow():
            await ActivityHelpers.execute_activity(
                activity_name=MessagesActivityName.UPDATE_MESSAGE,
                request=params,
                response_type=type(None),
                start_to_close_timeout=start_to_close_timeout,
                retry_policy=retry_policy,
                heartbeat_timeout=heartbeat_timeout,
            )
            return
        await self._messages_service.update_message(task_id=task_id, message_id=message_id, content=content)

    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,
        start_to_close_timeout: timedelta = timedelta(seconds=5),
        heartbeat_timeout: timedelta = timedelta(seconds=5),
        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
    ) -> None:
        """
        Send a batch of messages for a task.

        Args:
            task_id: The ID of the task.
            contents: The messages to send.
            emit_updates: Whether to emit updates.
            trace_id: The trace ID for tracing.
            parent_span_id: The parent span ID for tracing.
            start_to_close_timeout: The start to close timeout.
            heartbeat_timeout: The heartbeat timeout.
            retry_policy: The retry policy.
        """
        params = CreateMessagesBatchParams(
            task_id=task_id,
            contents=contents,
            emit_updates=emit_updates,
            trace_id=trace_id,
            parent_span_id=parent_span_id,
        )
        if in_temporal_workflow():
            await ActivityHelpers.execute_activity(
                activity_name=MessagesActivityName.CREATE_MESSAGES_BATCH,
                request=params,
                response_type=type(None),
                start_to_close_timeout=start_to_close_timeout,
                retry_policy=retry_policy,
                heartbeat_timeout=heartbeat_timeout,
            )
            return
        await self._messages_service.send_batch(task_id=task_id, contents=contents, emit_updates=emit_updates)

    async def update_batch(
        self,
        task_id: str,
        updates: dict[str, TaskMessageContent],
        trace_id: str | None = None,
        parent_span_id: str | None = None,
        start_to_close_timeout: timedelta = timedelta(seconds=5),
        heartbeat_timeout: timedelta = timedelta(seconds=5),
        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
    ) -> None:
        """
        Update a batch of messages for a task.

        Args:
            task_id: The ID of the task.
            updates: Dict mapping message_id to new content.
            trace_id: The trace ID for tracing.
            parent_span_id: The parent span ID for tracing.
            start_to_close_timeout: The start to close timeout.
            heartbeat_timeout: The heartbeat timeout.
            retry_policy: The retry policy.
        """
        params = UpdateMessagesBatchParams(
            task_id=task_id, updates=updates, trace_id=trace_id, parent_span_id=parent_span_id
        )
        if in_temporal_workflow():
            await ActivityHelpers.execute_activity(
                activity_name=MessagesActivityName.UPDATE_MESSAGES_BATCH,
                request=params,
                response_type=type(None),
                start_to_close_timeout=start_to_close_timeout,
                retry_policy=retry_policy,
                heartbeat_timeout=heartbeat_timeout,
            )
            return
        await self._messages_service.update_messages_batch(task_id=task_id, updates=updates)

    async def list(
        self,
        task_id: str,
        limit: int | None = None,
        trace_id: str | None = None,
        parent_span_id: str | None = None,
        start_to_close_timeout: timedelta = timedelta(seconds=5),
        heartbeat_timeout: timedelta = timedelta(seconds=5),
        retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
    ) -> list[UiMessage]:
        """
        List messages for a task.

        Args:
            task_id: The ID of the task.
            limit: The maximum number of messages to return.
            trace_id: The trace ID for tracing.
            parent_span_id: The parent span ID for tracing.
            start_to_close_timeout: The start to close timeout.
            heartbeat_timeout: The heartbeat timeout.
            retry_policy: The retry policy.

        Returns:
            List of UiMessage objects.
        """
        params = ListMessagesParams(task_id=task_id, limit=limit, trace_id=trace_id, parent_span_id=parent_span_id)
        if in_temporal_workflow():
            return await ActivityHelpers.execute_activity(
                activity_name=MessagesActivityName.LIST_MESSAGES,
                request=params,
                response_type=list[UiMessage],
                start_to_close_timeout=start_to_close_timeout,
                retry_policy=retry_policy,
                heartbeat_timeout=heartbeat_timeout,
            )
        return await self._messages_service.list_messages(task_id=task_id, limit=limit)
