from __future__ import annotations

import json
from abc import ABC, abstractmethod
from typing import Any, Literal, override

from terminaluse.lib.types.llm_messages import (
    AssistantMessage,
    Message,
    ToolCall,
    ToolCallRequest,
    ToolMessage,
    UserMessage,
)
from terminaluse.types.ui_message import UiMessage
from terminaluse.types.ui_message_part_wrapper import (
    UiMessagePartWrapper_Data,
    UiMessagePartWrapper_Text,
    UiMessagePartWrapper_Tool,
)


class UiMessageConverter(ABC):
    """
    Abstract base class for converting a UiMessage to an LLM Message.

    Each converter should be responsible for one content type.
    """

    @abstractmethod
    def convert(self, ui_message: UiMessage) -> Message:
        """
        Convert a UiMessage to an LLM Message.

        Args:
            ui_message: The UiMessage to convert

        Returns:
            A Message (Pydantic model)
        """
        pass


class DefaultTextPartConverter(UiMessageConverter):
    """Converter for UiMessage with text parts."""

    @override
    def convert(self, ui_message: UiMessage) -> Message:
        """Convert text parts to UserMessage or AssistantMessage based on role."""
        text_parts = []
        for part in ui_message.parts:
            if isinstance(part, UiMessagePartWrapper_Text):
                text_parts.append(part.text)

        content = "\n".join(text_parts) if text_parts else ""

        if ui_message.role == "user":
            return UserMessage(content=content)
        else:  # assistant or other
            return AssistantMessage(content=content)


class DefaultToolPartConverter(UiMessageConverter):
    """Converter for UiMessage with tool parts."""

    @override
    def convert(self, ui_message: UiMessage) -> Message:
        """Convert tool parts to AssistantMessage with tool_calls or ToolMessage."""
        tool_parts = [p for p in ui_message.parts if isinstance(p, UiMessagePartWrapper_Tool)]

        if not tool_parts:
            return AssistantMessage(content="")

        # If there's output, this is a tool response
        first_tool = tool_parts[0]
        if first_tool.output is not None:
            return ToolMessage(
                content=str(first_tool.output),
                tool_call_id=first_tool.tool_call_id,
                name=first_tool.tool_name,
            )

        # Otherwise it's a tool request
        tool_calls = []
        for part in tool_parts:
            arguments_str = json.dumps(part.input) if part.input else "{}"
            tool_call = ToolCallRequest(
                id=part.tool_call_id,
                function=ToolCall(name=part.tool_name, arguments=arguments_str),
            )
            tool_calls.append(tool_call)

        return AssistantMessage(content=None, tool_calls=tool_calls)


class DefaultDataPartConverter(UiMessageConverter):
    """Converter for UiMessage with data parts."""

    @override
    def convert(self, ui_message: UiMessage) -> Message:
        """Convert data parts to UserMessage or AssistantMessage based on role."""
        data_parts = []
        for part in ui_message.parts:
            if isinstance(part, UiMessagePartWrapper_Data):
                data_parts.append(str(part.data))

        content = "\n".join(data_parts) if data_parts else ""

        if ui_message.role == "user":
            return UserMessage(content=content)
        else:  # assistant or other
            return AssistantMessage(content=content)


class DefaultMixedContentConverter(UiMessageConverter):
    """Converter for UiMessage with mixed content types."""

    @override
    def convert(self, ui_message: UiMessage) -> Message:
        """Convert mixed content to a message, prioritizing text."""
        text_parts = []
        for part in ui_message.parts:
            if isinstance(part, UiMessagePartWrapper_Text):
                text_parts.append(part.text)
            elif isinstance(part, UiMessagePartWrapper_Data):
                text_parts.append(str(part.data))

        content = "\n".join(text_parts) if text_parts else ""

        if ui_message.role == "user":
            return UserMessage(content=content)
        else:
            return AssistantMessage(content=content)


def convert_ui_message_to_llm_message(
    ui_message: UiMessage,
    output_mode: Literal["pydantic", "dict"] = "pydantic",
    text_converter: UiMessageConverter | None = None,
    tool_converter: UiMessageConverter | None = None,
    data_converter: UiMessageConverter | None = None,
    mixed_converter: UiMessageConverter | None = None,
) -> Message | dict[str, Any]:
    """
    Convert a UiMessage to an LLM Message format.

    Args:
        ui_message: The UiMessage to convert
        output_mode: Whether to return a Pydantic model or dict
        text_converter: Optional converter for text content.
        tool_converter: Optional converter for tool content.
        data_converter: Optional converter for data content.
        mixed_converter: Optional converter for mixed content.

    Returns:
        Either a Message (Pydantic model) or dict representation
    """
    # Determine the primary content type
    has_text = any(isinstance(p, UiMessagePartWrapper_Text) for p in ui_message.parts)
    has_tool = any(isinstance(p, UiMessagePartWrapper_Tool) for p in ui_message.parts)
    has_data = any(isinstance(p, UiMessagePartWrapper_Data) for p in ui_message.parts)

    # Select converter based on content
    if has_tool and not has_text and not has_data:
        converter = tool_converter or DefaultToolPartConverter()
    elif has_text and not has_tool and not has_data:
        converter = text_converter or DefaultTextPartConverter()
    elif has_data and not has_tool and not has_text:
        converter = data_converter or DefaultDataPartConverter()
    else:
        # Mixed content
        converter = mixed_converter or DefaultMixedContentConverter()

    message = converter.convert(ui_message)

    if output_mode == "dict":
        return message.model_dump()
    return message


def convert_ui_messages_to_llm_messages(
    ui_messages: list[UiMessage],
    output_mode: Literal["pydantic", "dict"] = "pydantic",
    text_converter: UiMessageConverter | None = None,
    tool_converter: UiMessageConverter | None = None,
    data_converter: UiMessageConverter | None = None,
    mixed_converter: UiMessageConverter | None = None,
) -> list[Message | dict[str, Any]]:
    """
    Convert a list of UiMessages to LLM Message format.

    Args:
        ui_messages: List of UiMessages to convert
        output_mode: Whether to return Pydantic models or dicts
        text_converter: Optional converter for text content.
        tool_converter: Optional converter for tool content.
        data_converter: Optional converter for data content.
        mixed_converter: Optional converter for mixed content.

    Returns:
        List of either Messages (Pydantic models) or dicts
    """
    return [
        convert_ui_message_to_llm_message(
            ui_message,
            output_mode,
            text_converter,
            tool_converter,
            data_converter,
            mixed_converter,
        )
        for ui_message in ui_messages
    ]


# Backward compatibility aliases
TaskMessageConverter = UiMessageConverter
convert_task_message_to_llm_messages = convert_ui_message_to_llm_message
convert_task_messages_to_llm_messages = convert_ui_messages_to_llm_messages
