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

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.types.message_parts import DataPart, TextPart
from terminaluse.lib.core.temporal.activities.activity_helpers import ActivityHelpers
from terminaluse.lib.core.temporal.activities.adk.messages_activities import (
    ListMessagesParams,
    MessagesActivityName,
)
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,
    SDK_TYPE_CODEX,
    get_claude_sdk_version,
    get_codex_sdk_version,
)
from terminaluse.types.ui_message import UiMessage

if TYPE_CHECKING:
    from agents.extensions.experimental.codex.events import ThreadEvent as CodexSdkMessage
    from claude_agent_sdk.types import Message as ClaudeAgentSdkMessage

logger = make_logger(__name__)
DEFAULT_RETRY_POLICY = RetryPolicy(maximum_attempts=1)
CODEX_THREAD_EVENT_TYPES = {
    "thread.started",
    "turn.started",
    "turn.completed",
    "turn.failed",
    "item.started",
    "item.updated",
    "item.completed",
    "error",
}


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

    v2 Architecture: SDK-native messages are forwarded directly to the backend via raw_events.
    The backend adapter (ClaudeAdapter, CodexAdapter, etc.) handles 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,
                idempotency_key=str(uuid4()),
            )
        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,
            )
            raise

    async def _publish_codex_event(
        self,
        task_id: str,
        raw_event: dict[str, Any],
        sdk_version: str | None = None,
    ) -> None:
        """
        Publish a Codex SDK thread event to the backend via raw_events API.

        Args:
            task_id: The task ID
            raw_event: The raw Codex thread event
            sdk_version: Optional explicit Codex SDK version
        """
        event_type = raw_event.get("type", "unknown")
        effective_sdk_version = sdk_version or get_codex_sdk_version()

        try:
            await self._terminaluse_client.raw_events.publish_raw_event(
                task_id=task_id,
                sdk_type=SDK_TYPE_CODEX,
                sdk_version=effective_sdk_version,
                raw_event=raw_event,
                idempotency_key=str(uuid4()),
            )
        except Exception as e:
            logger.error(
                f"Failed to publish Codex event: task_id={task_id}, "
                f"event_type={event_type}, sdk_version={effective_sdk_version}, error={e}",
                exc_info=True,
            )
            raise

    async def send(
        self,
        task_id: str,
        content: "ClaudeAgentSdkMessage | CodexSdkMessage | TextPart | DataPart | str",
        emit_updates: bool = True,
        trace_id: str | None = None,
        parent_span_id: str | None = None,
        sdk_version: str | None = None,
    ) -> None:
        """Send a message.

        Args:
            task_id: The task ID.
            content: Message content - can be:
                - A string (auto-wrapped to TextPart)
                - TextPart or DataPart for structured messages
                - ClaudeAgentSdkMessage for Claude SDK streaming
                - CodexSdkMessage with `as_dict()` for Codex SDK thread events
            emit_updates: Whether to emit updates.
            trace_id: Optional trace ID for tracing.
            parent_span_id: Optional parent span ID for tracing.
            sdk_version: Optional SDK version override (used for raw Codex events).
        """
        # Handle Claude SDK messages
        if self._is_claude_message(content):
            await self._handle_claude_message(
                task_id=task_id,
                message=content,
            )
            return

        # Handle raw Codex SDK events (CodexSdkMessage objects)
        codex_raw_event = self._coerce_codex_raw_event(content)
        if codex_raw_event is not None:
            await self._handle_codex_event(
                task_id=task_id,
                raw_event=codex_raw_event,
                sdk_version=sdk_version,
            )
            return

        # Auto-wrap strings to TextPart
        if isinstance(content, str):
            content = TextPart(text=content)

        if not isinstance(content, (TextPart, DataPart)):
            raise TypeError(
                "Unsupported content type. Expected TextPart, DataPart, str, ClaudeAgentSdkMessage, or CodexSdkMessage."
            )

        # Handle AgentMessagePart (TextPart, DataPart)
        await self._messages_service.send(
            task_id=task_id,
            content=content,
            emit_updates=emit_updates,
            trace_id=trace_id,
            parent_span_id=parent_span_id,
        )

    async def _handle_claude_message(
        self,
        task_id: str,
        message: "ClaudeAgentSdkMessage",
    ) -> 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,
        )

    async def _handle_codex_event(
        self,
        task_id: str,
        raw_event: dict[str, Any],
        sdk_version: str | None = None,
    ) -> None:
        """
        Forward Codex SDK thread events directly to the backend via raw_events.

        The backend CodexAdapter handles transformation, streaming, and storage.
        """
        await self._publish_codex_event(
            task_id=task_id,
            raw_event=raw_event,
            sdk_version=sdk_version,
        )

    def _coerce_codex_raw_event(self, content: Any) -> dict[str, Any] | None:
        """Normalize a CodexSdkMessage object to dict form when possible."""
        as_dict = getattr(content, "as_dict", None)
        if callable(as_dict):
            try:
                candidate = as_dict()
            except Exception:
                candidate = None
            if isinstance(candidate, dict) and self._is_codex_raw_event(candidate):
                return candidate

        return None

    def _is_codex_raw_event(self, raw_event: dict[str, Any]) -> bool:
        """Return True if payload matches Codex SDK thread event envelope."""
        event_type = raw_event.get("type")
        if not isinstance(event_type, str) or event_type not in CODEX_THREAD_EVENT_TYPES:
            return False

        if event_type == "thread.started":
            return isinstance(raw_event.get("thread_id"), str)
        if event_type == "turn.started":
            return True
        if event_type == "turn.completed":
            return isinstance(raw_event.get("usage"), dict)
        if event_type == "turn.failed":
            return isinstance(raw_event.get("error"), dict)
        if event_type in {"item.started", "item.updated", "item.completed"}:
            return isinstance(raw_event.get("item"), dict)
        if event_type == "error":
            message = raw_event.get("message")
            if isinstance(message, str):
                return True
            nested_error = raw_event.get("error")
            return isinstance(nested_error, dict) and isinstance(nested_error.get("message"), str)

        return False

    def _is_claude_message(self, content: Any) -> bool:
        """Detect Claude SDK message objects without importing claude_agent_sdk."""
        cls = type(content)
        module = getattr(cls, "__module__", "")
        return isinstance(module, str) and module.startswith("claude_agent_sdk")

    def _claude_message_to_raw_event(self, message: "ClaudeAgentSdkMessage") -> 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 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)
