from collections.abc import Callable
from typing import Any, Sequence

from pydantic import BaseModel, Field

from langchain_core.runnables import Runnable
from langchain_core.tools import BaseTool
from langchain_core.messages import BaseMessage
from langchain_core.prompt_values import PromptValue
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel

from langchain.tools import tool
from langchain.messages import AIMessage, ToolCall
from langchain.agents.structured_output import(
    StructuredOutputValidationError,
    ToolStrategy
)
from langchain.agents import create_agent


# To bind structured output tool to fake model
class FakeWithSO(GenericFakeChatModel):
    def bind_tools(
        self,
        tools: Sequence[dict[str, Any] | type | Callable[..., Any] | BaseTool],
        *,
        tool_choice: str | None = None, **kwargs: Any
        ) -> Runnable[PromptValue | str | Sequence[BaseMessage | list[str] | tuple[str, str] | str | dict[str, Any]], AIMessage]:
        return self.bind(tools=tools, tool_choice=tool_choice, **kwargs)


# Expected response format
class Parrot(BaseModel):
    name: str = Field(description="Parrot name.")
    species: str = Field(description="Parrot species.")
    response: str = Field(description="Parrot's response.")

# A fake tool
@tool
def fake_tool(x:str)-> str:
    """A fake tool."""
    return "fake!"



# Test for correct tool calls
correct_call_model = FakeWithSO(
    messages=iter(
        [
            AIMessage(
                content="A tool call",
                tool_calls=[
                    ToolCall(
                        name="Parrot",
                        args={
                            "name": "Polly",
                            "species": "Norwegian Blue",
                            "response": "...",
                        },
                        id="call_1",
                    )
                ],
            ),
        ]
    )
)

correct_call_agent = create_agent(
    model=correct_call_model,
    tools=[fake_tool],
    response_format=ToolStrategy(Parrot, handle_errors=False)
)
try:
    response = correct_call_agent.invoke(
        {"messages": [{"role": "user", "content": "Hello, Polly!!"}]}
    )
    print("No exception raised for correct tool call!")
    assert("structured_response" in response)
except StructuredOutputValidationError:
    print("Exception raised for correct tool call!")


# Test for bad tool calls (name -> NOTname)
bad_call_model = FakeWithSO(
    messages=iter(
        [
            AIMessage(
                content="A tool call",
                tool_calls=[
                    ToolCall(
                        name="Parrot",
                        args={
                            "NOTname": "Polly",
                            "species": "Norwegian Blue",
                            "response": "...",
                        },
                        id="call_1",
                    )
                ],
            ),
        ]
    )
)

bad_call_agent = create_agent(
    model=bad_call_model,
    tools=[fake_tool],
    response_format=ToolStrategy(Parrot, handle_errors=False)
)
try:
    response = bad_call_agent.invoke(
        {"messages": [{"role": "user", "content": "Hello, Polly!!"}]}
    )
    print("No exception raised for bad tool call!")
    assert("structured_response" in response)
except StructuredOutputValidationError:
    print("Exception raised for bad tool call!")


# Test for no tool calls
no_call_model = FakeWithSO(
    messages=iter(
        [
            AIMessage(
                content="NOT a tool call",
            ),
        ]
    )
)

no_call_agent = create_agent(
    model=no_call_model,
    tools=[fake_tool],
    response_format=ToolStrategy(Parrot, handle_errors=False)
)
try:
    response = no_call_agent.invoke(
        {"messages": [{"role": "user", "content": "Hello, Polly!!"}]}
    )
    print("No exception raised for no tool call!")
    assert("structured_response" in response)
except StructuredOutputValidationError:
    print("Exception raised for no tool call!")


# No exception raised for correct tool call!
# Exception raised for bad tool call!
# No exception raised for no tool call!

---

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[1], line 133
    129     response = no_call_agent.invoke(
    130         {"messages": [{"role": "user", "content": "Hello, Polly!!"}]}
    131     )
    132     print("No exception raised for no tool call!")
--> 133     assert("structured_response" in response)
    134 except StructuredOutputValidationError:
    135     print("Exception raised for no tool call!")

AssertionError:

---

# at langchain_v1/langchain/agents/structured_outputs.py
class NoStructuredOutputError(StructuredOutputError):
    """Raised when no structured output tool is called."""

    def __init__(self, ai_message: AIMessage) -> None:
        """Initialize `NoStructuredOutputError`.

        Args:
            ai_message: The AI message that contained the invalid structured output.
        """

        self.ai_message = ai_message
        super().__init__(f"No structured output tool called.")

---

elif (isinstance(effective_response_format, ToolStrategy)
            and isinstance(output, AIMessage)
        ):
            exception = NoStructuredOutputError(output)
            should_retry, error_message = _handle_structured_output_error(
                exception, effective_response_format
            )
            if not should_retry:
                raise exception
            
            return {
                "messages": [
                    output,
                    ToolMessage(
                    content=error_message,
                    ),
                ],
            }