from typing import Any

from langchain.agents import create_agent
from langchain.messages import AIMessage, ToolMessage
from langchain.tools import tool
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import BaseMessage
from langchain_core.outputs import ChatGeneration, ChatResult


@tool
def convert_celsius_to_fahrenheit(celsius: float) -> float:
    """Convert a Celsius temperature to Fahrenheit."""
    return celsius * 9 / 5 + 32


class ToolResultModel(BaseChatModel):
    @property
    def _llm_type(self) -> str:
        return "tool-result-model"

    def bind_tools(self, tools, *, tool_choice=None, **kwargs):
        return self

    def _generate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: CallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ) -> ChatResult:
        tool_message = next(
            (
                message
                for message in reversed(messages)
                if isinstance(message, ToolMessage)
            ),
            None,
        )

        if tool_message is None:
            response = AIMessage(
                content="",
                tool_calls=[
                    {
                        "name": "convert_celsius_to_fahrenheit",
                        "args": {"celsius": 20},
                        "id": "call_temperature",
                    }
                ],
            )
        else:
            fahrenheit = float(tool_message.content)
            response = AIMessage(
                content=(
                    "20 degrees Celsius equals "
                    f"{fahrenheit:g} degrees Fahrenheit."
                )
            )

        return ChatResult(generations=[ChatGeneration(message=response)])


model = ToolResultModel()

agent = create_agent(
    model=model,
    tools=[convert_celsius_to_fahrenheit],
    system_prompt="Use the temperature tool for conversions.",
)

result = agent.invoke(
    {"messages": [{"role": "user", "content": "Convert 20 C to Fahrenheit."}]}
)

tool_call = next(
    call
    for message in result["messages"]
    if isinstance(message, AIMessage)
    for call in message.tool_calls
)
tool_message = next(
    message for message in result["messages"] if isinstance(message, ToolMessage)
)
final_reply = result["messages"][-1].content
expected_reply = (
    "20 degrees Celsius equals "
    f"{float(tool_message.content):g} degrees Fahrenheit."
)
assert final_reply == expected_reply, "Agent reply did not use the tool result."

print(f"Tool called: {tool_call['name']}")
print(f"Tool result: {tool_message.content}")
print(f"Agent reply: {final_reply}")