from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
import json

def format_conversation_for_download(conversation: list) -> str:
    """Formats the conversation history for download."""
    formatted_messages = []
    tool_id_to_name_map = {}
    for message in conversation:
        if isinstance(message, AIMessage) and message.tool_calls:
            for tool_call in message.tool_calls:
                tool_id_to_name_map[tool_call["id"]] = tool_call["name"]

    for message in conversation:
        if isinstance(message, HumanMessage):
            formatted_messages.append({"role": "user", "content": message.content})
        elif isinstance(message, AIMessage):
            if not message.content and not message.tool_calls:
                continue
            
            ai_dict = {"role": "assistant"}
            if message.content:
                ai_dict["content"] = message.content

            if message.tool_calls:
                ai_dict["tool_calls"] = [
                    {
                        "type": "function",
                        "function": {
                            "name": tc["name"],
                            "arguments": tc["args"],
                        },
                    }
                    for tc in message.tool_calls
                ]
            formatted_messages.append(ai_dict)
        elif isinstance(message, ToolMessage):
            tool_name = tool_id_to_name_map.get(message.tool_call_id)
            formatted_messages.append(
                {
                    "role": "tool",
                    "name": tool_name,
                    "content": message.content,
                }
            )
            
    return json.dumps(formatted_messages, indent=4)