import asyncio
import json
from typing import Any

from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.core.base.llms.types import (
    ChatMessage,
    ChatResponse,
    CompletionResponse,
    LLMMetadata,
    MessageRole,
)
from llama_index.core.llms.function_calling import FunctionCallingLLM
from llama_index.core.llms.llm import ToolSelection
from llama_index.core.workflow import Context, JsonSerializer

TOOL_CALL_LOG: list[str] = []


async def set_name(ctx: Context, name: str) -> str:
    """Save the user's preferred name in agent state."""
    async with ctx.store.edit_state() as ctx_state:
        ctx_state["state"]["name"] = name
    TOOL_CALL_LOG.append(f"set_name(name={name})")
    return name


async def read_name(ctx: Context) -> str:
    """Read the user's preferred name from agent state."""
    state = await ctx.store.get("state")
    TOOL_CALL_LOG.append("read_name()")
    return state["name"]


class StateAwareMockLLM(FunctionCallingLLM):
    @property
    def metadata(self) -> LLMMetadata:
        return LLMMetadata(
            is_chat_model=True,
            is_function_calling_model=True,
            model_name="state-aware-mock-llm",
        )

    def _prepare_chat_with_tools(
        self,
        tools: list[Any],
        user_msg: str | ChatMessage | None = None,
        chat_history: list[ChatMessage] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        messages = list(chat_history or [])
        if isinstance(user_msg, str):
            messages.append(ChatMessage(role=MessageRole.USER, content=user_msg))
        elif isinstance(user_msg, ChatMessage):
            messages.append(user_msg)
        return {"messages": messages}

    def get_tool_calls_from_response(
        self,
        response: ChatResponse,
        error_on_no_tool_call: bool = True,
        **kwargs: Any,
    ) -> list[ToolSelection]:
        calls = response.message.additional_kwargs.get("tool_calls", [])
        if not calls and error_on_no_tool_call:
            raise ValueError("No tool calls returned")
        return calls

    def chat(self, messages: list[ChatMessage], **kwargs: Any) -> ChatResponse:
        if messages and messages[-1].role == MessageRole.TOOL:
            return ChatResponse(
                message=ChatMessage(
                    role=MessageRole.ASSISTANT,
                    content=f"The saved name is {messages[-1].content}.",
                )
            )

        latest_user = next(
            str(message.content)
            for message in reversed(messages)
            if message.role == MessageRole.USER
        )
        if "What name did I ask" not in latest_user:
            tool_call = ToolSelection(
                tool_id="call_set_name",
                tool_name="set_name",
                tool_kwargs={"name": "Laurie"},
            )
        else:
            tool_call = ToolSelection(
                tool_id="call_read_name",
                tool_name="read_name",
                tool_kwargs={},
            )

        return ChatResponse(
            message=ChatMessage(
                role=MessageRole.ASSISTANT,
                content="",
                additional_kwargs={"tool_calls": [tool_call]},
            )
        )

    async def achat(
        self, messages: list[ChatMessage], **kwargs: Any
    ) -> ChatResponse:
        return self.chat(messages, **kwargs)

    def stream_chat(self, messages: list[ChatMessage], **kwargs: Any):
        yield self.chat(messages, **kwargs)

    async def astream_chat(self, messages: list[ChatMessage], **kwargs: Any):
        async def response_stream():
            yield self.chat(messages, **kwargs)

        return response_stream()

    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        return CompletionResponse(text="")

    async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        return self.complete(prompt, **kwargs)

    def stream_complete(self, prompt: str, **kwargs: Any):
        yield CompletionResponse(text="")

    async def astream_complete(self, prompt: str, **kwargs: Any):
        yield CompletionResponse(text="")


def build_workflow() -> AgentWorkflow:
    return AgentWorkflow.from_tools_or_functions(
        [set_name, read_name],
        llm=StateAwareMockLLM(),
        system_prompt="Use the state tools to save or retrieve the user's name.",
        initial_state={"name": "unset"},
    )


async def main() -> None:
    workflow = build_workflow()
    ctx = Context(workflow)

    saved_response = await workflow.run(
        user_msg="Remember that my name is Laurie.", ctx=ctx
    )
    print(f"saved_agent={saved_response}")

    ctx_dict = ctx.to_dict(serializer=JsonSerializer())
    with open("agent-state.json", "w", encoding="utf-8") as state_file:
        json.dump(ctx_dict, state_file)

    with open("agent-state.json", encoding="utf-8") as state_file:
        saved_ctx_dict = json.load(state_file)

    restored_ctx = Context.from_dict(
        workflow,
        saved_ctx_dict,
        serializer=JsonSerializer(),
    )
    restored_response = await workflow.run(
        user_msg="What name did I ask you to remember?", ctx=restored_ctx
    )
    print(f"restored_agent={restored_response}")
    print(f"tool_calls={TOOL_CALL_LOG}")


if __name__ == "__main__":
    asyncio.run(main())