"""z.ai adapter for the provider-neutral micro-agents interface."""

from __future__ import annotations

from typing import Any

import orjson
from micro_agents import OpenAIBasedLLM
from micro_agents.conversation import (
    Conversation,
    LlmMessage,
    Message,
    ReasoningMessage,
    ToolCallMessage,
    ToolCallResponseMessage,
    UserMessage,
)
import micro_agents.tools as tools_package


class ZaiLLM(OpenAIBasedLLM):
    """Use z.ai's OpenAI-compatible Chat Completions endpoint.

    ``OpenAIBasedLLM`` in micro-agents targets OpenAI's Responses API, which
    z.ai does not expose.  This adapter retains its tool-schema inference while
    translating the agent conversation to the Chat Completions format.
    """

    async def ask(
        self,
        conversation: Conversation,
        tools: list[tools_package.Tool],
        *,
        instructions: str | None = None,
        thinking: bool | None = None,
    ) -> list[Message]:
        extra_body: dict[str, object] | None = None
        if thinking is not None:
            extra_body = {"thinking": {"type": "enabled" if thinking else "disabled"}}
        response = await self._ai.chat.completions.create(
            model=self._model,
            messages=self._serialize_conversation(conversation, instructions),
            tools=self._serialize_chat_tools(tools),
            stream=False,
            extra_body=extra_body,
        )
        if not response.choices:
            raise ValueError("z.ai returned no choices")
        message = response.choices[0].message
        tool_calls = message.tool_calls or []
        out: list[Message] = []

        # GLM returns this field when thinking is enabled.  Keeping it in the
        # conversation lets the next call include it with the tool invocation.
        reasoning = getattr(message, "reasoning_content", None)
        if isinstance(reasoning, str) and reasoning:
            out.append(ReasoningMessage(reasoning))

        # A Chat Completions response with tool calls is represented by one
        # assistant message.  The content is normally empty in that case; do
        # not create a second assistant message before the tool-call message.
        if not tool_calls and isinstance(message.content, str) and message.content:
            out.append(LlmMessage(message.content))
        for call in tool_calls:
            function = call.function
            if not isinstance(call.id, str) or not isinstance(function.name, str):
                raise ValueError("z.ai returned an invalid function call")
            try:
                arguments = orjson.loads(function.arguments)
            except orjson.JSONDecodeError as error:
                raise ValueError("z.ai returned invalid JSON function arguments") from error
            if not isinstance(arguments, dict):
                raise ValueError("z.ai function-call arguments must be an object")
            out.append(ToolCallMessage(call.id, function.name, arguments))
        return out

    def _serialize_chat_tools(self, tools: list[tools_package.Tool]) -> list[dict[str, object]]:
        result: list[dict[str, object]] = []
        for tool in tools:
            schema = self._serialize_tool(tool)
            result.append(
                {
                    "type": "function",
                    "function": {
                        "name": schema["name"],
                        "description": schema["description"],
                        "parameters": schema["parameters"],
                    },
                }
            )
        return result

    @staticmethod
    def _serialize_conversation(conversation: Conversation, instructions: str | None) -> list[dict[str, Any]]:
        messages: list[dict[str, Any]] = []
        if instructions:
            messages.append({"role": "system", "content": instructions})

        items = conversation.messages
        index = 0
        pending_reasoning: str | None = None
        while index < len(items):
            item = items[index]
            if isinstance(item, ReasoningMessage):
                pending_reasoning = item.text()
                index += 1
                continue
            if isinstance(item, UserMessage):
                messages.append({"role": "user", "content": item.message()})
            elif isinstance(item, LlmMessage):
                assistant: dict[str, Any] = {"role": "assistant", "content": item.message()}
                if pending_reasoning:
                    assistant["reasoning_content"] = pending_reasoning
                    pending_reasoning = None
                messages.append(assistant)
            elif isinstance(item, ToolCallMessage):
                calls: list[dict[str, object]] = []
                while index < len(items) and isinstance(items[index], ToolCallMessage):
                    call = items[index]
                    calls.append(
                        {
                            "id": call.unique_id(),
                            "type": "function",
                            "function": {
                                "name": call.tool_name(),
                                "arguments": orjson.dumps(call.arguments()).decode("utf-8"),
                            },
                        }
                    )
                    index += 1
                assistant = {"role": "assistant", "content": None, "tool_calls": calls}
                if pending_reasoning:
                    assistant["reasoning_content"] = pending_reasoning
                    pending_reasoning = None
                messages.append(assistant)
                continue
            elif isinstance(item, ToolCallResponseMessage):
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": item.unique_id(),
                        "content": item.text(),
                    }
                )
            else:
                raise ValueError(f"Unknown type of message {type(item).__name__}")
            index += 1
        return messages