from typing import Any
from typing import AsyncGenerator
from typing import Dict
from typing import List
from typing import Union

from openinference.semconv.trace import OpenInferenceSpanKindValues
from pydantic import BaseModel
from pydantic import ConfigDict

from grafi.common.models.default_id import default_id
from grafi.common.models.execution_context import ExecutionContext
from grafi.common.models.message import Message


class Tool(BaseModel):
    """
    A base class representing a tool in the agent.

    This class defines the interface for all tools. Each specific tool should
    inherit from this class and implement its methods.
    """

    tool_id: str = default_id
    name: str
    type: str
    oi_span_type: OpenInferenceSpanKindValues

    model_config = ConfigDict(arbitrary_types_allowed=True)

    class Builder:
        """Inner builder class for Tool construction."""

        def __init__(self):
            self._tool = self._init_tool()

        def _init_tool(self) -> "Tool":
            raise NotImplementedError

        def name(self, name: str):
            self._tool.name = name
            return self

        def type(self, type_name: str):
            self._tool.type = type_name
            return self

        def oi_span_type(self, oi_span_type: OpenInferenceSpanKindValues):
            self._tool.oi_span_type = oi_span_type
            return self

        def build(self) -> "Tool":
            return self._tool

    def execute(
        self,
        execution_context: ExecutionContext,
        input_data: Union[Message, List[Message]],
    ) -> Message:
        """
        Process the input data and return a response.

        This method should be implemented by all subclasses to define
        the specific behavior of each tool.

        Raises:
            NotImplementedError: If the method is not implemented by a subclass.
        """
        raise NotImplementedError("Subclasses must implement this method.")

    async def a_execute(
        self,
        execution_context: ExecutionContext,
        input_data: Union[Message, List[Message]],
    ) -> AsyncGenerator[Message, None]:
        raise NotImplementedError("Subclasses must implement this method.")

    def to_message(self, response: Any) -> Message:
        """
        Convert the tool's response to a Message object.

        Args:
            response (Any): The response generated by the tool.

        Returns:
            Message: The response converted to a Message object.

        Raises:
            NotImplementedError: If the method is not implemented by a subclass.
        """
        raise NotImplementedError("Subclasses must implement this method.")

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert the tool instance to a dictionary representation.

        Returns:
            Dict[str, Any]: A dictionary representation of the tool.
        """
        return {
            "tool_id": self.tool_id,
        }
