"""Agent utils."""

import json

from llama_index.core.llms import ChatMessage, TextBlock
from typing import List, Type, Dict, Any, Optional, cast
from llama_index.core.bridge.pydantic import BaseModel
from llama_index.core.base.llms.types import ChatMessage
from llama_index.core.llms import LLM


def messages_to_xml_format(messages: List[ChatMessage]) -> List[ChatMessage]:
    blocks = [TextBlock(text="<current_conversation>\n")]
    system_msg: Optional[ChatMessage] = None
    for message in messages:
        if message.role.value == "system":
            system_msg = message
        blocks.append(TextBlock(text=f"\t<{message.role.value}>\n"))
        for block in message.blocks:
            if isinstance(block, TextBlock):
                blocks.append(TextBlock(text=f"\t\t<message>{block.text}</message>\n"))
        blocks.append(TextBlock(text=f"\t</{message.role.value}>\n"))
    blocks.append(TextBlock(text="</current_conversation>\n\n"))
    blocks.append(
        TextBlock(
            text="Given the conversation, format the output according to the provided schema."
        )
    )
    messages = [ChatMessage(role="user", blocks=blocks)]
    if system_msg:
        messages.insert(0, system_msg)
    return messages


async def generate_structured_response(
    messages: List[ChatMessage], llm: LLM, output_cls: Type[BaseModel]
) -> Dict[str, Any]:
    xml_message = messages_to_xml_format(messages)
    structured_response = await llm.as_structured_llm(
        output_cls,
    ).achat(messages=xml_message)
    return cast(Dict[str, Any], json.loads(structured_response.message.content))