import itertools

from langchain.agents import create_agent
from langchain.tools import tool
from langchain_core.messages import AIMessageChunk, ToolCall
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_openai import ChatOpenAI

import mlflow


class FakeOpenAI(ChatOpenAI, extra="allow"):
    # In normal LangChain tests, we use the fake OpenAI server to mock the OpenAI REST API.
    # The fake server returns the input payload as it is. However, for agent tests, the
    # response should be a specific format so that the agent can parse it correctly.
    # Also, mocking with mock.patch does not work for testing model serving (as the server
    # will run in a separate process).
    # Therefore, we mock the OpenAI client in the model definition here.
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Using itertools.cycle to create an infinite iterator
        self._responses = itertools.cycle([
            AIMessageChunk(
                content="",
                tool_calls=[ToolCall(name="multiply", args={"a": 2, "b": 3}, id="123")],
            ),
            AIMessageChunk(content="The result of 2 * 3 is 6."),
        ])

    def _generate(self, *args, **kwargs):
        return ChatResult(generations=[ChatGeneration(message=next(self._responses))])

    def _stream(self, *args, **kwargs):
        yield ChatGenerationChunk(message=next(self._responses))


@tool
def add(a: int, b: int) -> int:
    """Add two numbers."""
    return a + b


@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers."""
    return a * b


llm = FakeOpenAI()
agent = create_agent(llm, [add, multiply], system_prompt="You are a helpful assistant")
mlflow.models.set_model(agent)