from langchain_core.tools import tool
from langchain_core.messages import ToolMessage
from langgraph.types import Command
from typing import Annotated, Any
from langgraph_swarm.handoff import _get_field, InjectedState
from loguru import logger


def create_end_conversation_tool(
    name: str = "end_conversation",
    description: str = "End the conversation naturally when the discussion has reached its conclusion or the user indicates they want to stop",
) -> callable:
    """Create a tool that can end the conversation.

    Args:
        name: Name of the tool to use for ending the conversation.
        description: Description for the end conversation tool.

    Returns:
        A tool that when called will end the conversation.
    """

    @tool(name, description=description)
    def end_conversation(
        reason: str,
        # Annotation is typed as Any instead of StateLike. StateLike
        # trigger validation issues from Pydantic / langchain_core interaction.
        # https://github.com/langchain-ai/langchain/issues/32067
        state: Annotated[Any, InjectedState],
        tool_call_id: str,
    ) -> Command:
        """End the conversation with the given reason.

        Args:
            reason: The reason for ending the conversation
        """
        logger.info(f"Conversation ended. Reason: {reason}")

        tool_message = ToolMessage(
            content=f"Conversation ended: {reason}",
            name=name,
            tool_call_id=tool_call_id,
        )

        # Return a Command that goes to END node
        return Command(
            # goto="END",
            graph=Command.PARENT,
            update={
                "messages": [*_get_field(state, "messages"), tool_message],
            },
        )

    return end_conversation


@tool
def function_name(
    input: str,
) -> str:
    """
    Mô tả chức năng của hàm này.
    """
    logger.info(f"Received input: {input}")
    # Thực hiện các thao tác cần thiết với input
    result: str = f"Processed: {input}"
    logger.info(f"Returning result: {result}")
    return result