import pytest

from llama_index.core.llms import ChatMessage
from llama_index.core.tools import ToolSelection
from llama_index.core.bridge.pydantic import BaseModel, ValidationError
from llama_index.core.agent.workflow.workflow_events import (
    AgentWorkflowStartEvent,
    AgentOutput,
    PydanticConversionWarning,
)
from llama_index.core.memory import Memory


@pytest.fixture()
def example_agent_output() -> dict:
    return {
        "response": ChatMessage(role="user", content="30 times 2 is 60."),
        "tool_calls": [
            ToolSelection(
                tool_id="1", tool_name="multiply", tool_kwargs={"i": 30, "j": 2}
            )
        ],
        "raw": '{"role": "user", "content": "30 times 2 is 60."}',
        "structured_response": {"operation": "30 times 2", "result": 60},
        "current_agent_name": "CalculatorAgent",
    }


class MathResult(BaseModel):
    operation: str
    result: int


class WrongMathResult(BaseModel):
    operation: str
    result: str


def test_agent_workflow_start_event():
    event = AgentWorkflowStartEvent(
        user_msg="Hello, world!",
        chat_history=[ChatMessage(role="user", content="Hello, world!")],
        max_iterations=10,
    )
    assert event.user_msg == "Hello, world!"
    assert event.chat_history[0].role.value == "user"
    assert event.chat_history[0].content == "Hello, world!"
    assert event.max_iterations == 10


def test_agent_workflow_start_event_with_dict():
    event = AgentWorkflowStartEvent(
        user_msg="Hello, world!",
        chat_history=[{"role": "user", "content": "Hello, world!"}],
        max_iterations=10,
    )
    assert event.user_msg == "Hello, world!"
    assert event.chat_history[0].role.value == "user"
    assert event.chat_history[0].content == "Hello, world!"
    assert event.max_iterations == 10


def test_agent_workflow_start_event_to_dict():
    event = AgentWorkflowStartEvent(
        user_msg="Hello, world!",
        chat_history=[ChatMessage(role="user", content="Hello, world!")],
        max_iterations=10,
        memory=Memory.from_defaults(),
    )

    # Memory is not included in the dump
    dump = event.model_dump()
    assert len(dump) == 3
    assert dump["user_msg"] == "Hello, world!"
    assert dump["chat_history"][0]["role"] == "user"
    assert dump["chat_history"][0]["blocks"][0]["text"] == "Hello, world!"
    assert dump["max_iterations"] == 10


def test_agent_output_with_structured_response(example_agent_output: dict) -> None:
    try:
        agent_output = AgentOutput.model_validate(example_agent_output)
        success = True
    except ValidationError:
        success = False
    assert success
    assert agent_output.get_pydantic_model(MathResult) == MathResult.model_validate(
        example_agent_output["structured_response"]
    )
    with pytest.warns(PydanticConversionWarning):
        a = agent_output.get_pydantic_model(WrongMathResult)
    assert a is None

from llama_index.core.llms import ChatMessage
from llama_index.core.agent.workflow.workflow_events import AgentWorkflowStartEvent
from llama_index.core.memory import Memory


def test_agent_workflow_start_event():
    event = AgentWorkflowStartEvent(
        user_msg="Hello, world!",
        chat_history=[ChatMessage(role="user", content="Hello, world!")],
        max_iterations=10,
    )
    assert event.user_msg == "Hello, world!"
    assert event.chat_history[0].role.value == "user"
    assert event.chat_history[0].content == "Hello, world!"
    assert event.max_iterations == 10


def test_agent_workflow_start_event_with_dict():
    event = AgentWorkflowStartEvent(
        user_msg="Hello, world!",
        chat_history=[{"role": "user", "content": "Hello, world!"}],
        max_iterations=10,
    )
    assert event.user_msg == "Hello, world!"
    assert event.chat_history[0].role.value == "user"
    assert event.chat_history[0].content == "Hello, world!"
    assert event.max_iterations == 10


def test_agent_workflow_start_event_to_dict():
    event = AgentWorkflowStartEvent(
        user_msg="Hello, world!",
        chat_history=[ChatMessage(role="user", content="Hello, world!")],
        max_iterations=10,
        memory=Memory.from_defaults(),
    )

    # Memory is not included in the dump
    dump = event.model_dump()
    assert len(dump) == 3
    assert dump["user_msg"] == "Hello, world!"
    assert dump["chat_history"][0]["role"] == "user"
    assert dump["chat_history"][0]["blocks"][0]["text"] == "Hello, world!"
    assert dump["max_iterations"] == 10
