from __future__ import annotations

import base64
import json
from typing import Any, cast

import aiohttp

from micro_agents import (
    Conversation,
    LlmMessage,
    Media,
    Tool,
    ToolCallMessage,
    ToolCallResponseMessage,
    UserMessage,
)
from micro_agents.conversation import Message, ReasoningMessage
from micro_agents.llm import OpenAIBasedLLM


class ChatCompletionsLLM(OpenAIBasedLLM):
    """micro-agents adapter for OpenAI-compatible Chat Completions APIs."""

    async def ask(
        self,
        conversation: Conversation,
        tools: list[Tool],
        *,
        instructions: str | None = None,
        thinking: bool | None = None,
    ) -> list[Message]:
        messages = self._serialize_conversation(conversation, instructions)
        serialized_tools = self._serialize_tools(tools)
        request: dict[str, Any] = {
            "model": self._model,
            "messages": messages,
            "stream": False,
            "extra_body": self._extra_body(thinking),
        }
        if serialized_tools:
            request["tools"] = serialized_tools

        response = await self._ai.chat.completions.create(**cast(Any, request))
        if not response.choices:
            raise RuntimeError("Chat Completions API returned no choices")

        message = response.choices[0].message
        result: list[Message] = []
        reasoning = getattr(message, "reasoning_content", None)
        if reasoning:
            result.append(ReasoningMessage(str(reasoning)))
        if message.content:
            result.append(LlmMessage(message.content))

        for call in message.tool_calls or ():
            try:
                arguments = json.loads(call.function.arguments)
            except (TypeError, ValueError) as exc:
                raise ValueError("Model returned invalid tool-call arguments") from exc
            if not isinstance(arguments, dict):
                raise ValueError("Model tool-call arguments must be an object")
            result.append(ToolCallMessage(call.id, call.function.name, arguments))

        return result

    @staticmethod
    def _extra_body(thinking: bool | None) -> dict[str, object]:
        return {"thinking": {"type": "enabled" if thinking else "disabled"}}

    @staticmethod
    def _reasoning_fields(message: ReasoningMessage) -> dict[str, object]:
        return {"reasoning_content": message.text()}

    def _serialize_tools(self, tools: list[Tool]) -> list[dict[str, object]]:
        response_tools = super()._serialize_tools(tools)
        result: list[dict[str, object]] = []
        for tool in response_tools:
            function = {
                key: tool[key]
                for key in ("name", "description", "parameters")
                if key in tool and tool[key] is not None
            }
            result.append({"type": "function", "function": function})
        return result

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

        pending_assistant: dict[str, object] | None = None
        pending_attachments: list[Media] = []

        def flush_assistant() -> None:
            nonlocal pending_assistant
            if pending_assistant is None:
                return
            pending_assistant.setdefault("content", "")
            result.append(pending_assistant)
            pending_assistant = None

        def flush_attachments() -> None:
            if not pending_attachments:
                return
            result.append(
                {
                    "role": "user",
                    "content": cls._serialize_user_content(
                        "Attachments returned by tools:",
                        pending_attachments,
                    ),
                }
            )
            pending_attachments.clear()

        for message in conversation.messages:
            if isinstance(message, ToolCallResponseMessage):
                flush_assistant()
                result.append(
                    {
                        "role": "tool",
                        "tool_call_id": message.unique_id(),
                        "content": message.text(),
                    }
                )
                pending_attachments.extend(message.attachments())
                continue

            flush_attachments()
            if isinstance(message, UserMessage):
                flush_assistant()
                result.append(
                    {
                        "role": "user",
                        "content": cls._serialize_user_content(
                            message.message(),
                            message.media(),
                        ),
                    }
                )
            elif isinstance(message, ReasoningMessage):
                flush_assistant()
                pending_assistant = {"role": "assistant", **cls._reasoning_fields(message)}
            elif isinstance(message, LlmMessage):
                if pending_assistant is None:
                    pending_assistant = {"role": "assistant"}
                pending_assistant["content"] = message.message()
            elif isinstance(message, ToolCallMessage):
                if pending_assistant is None:
                    pending_assistant = {"role": "assistant"}
                calls = cast(list[dict[str, object]], pending_assistant.setdefault("tool_calls", []))
                calls.append(
                    {
                        "id": message.unique_id(),
                        "type": "function",
                        "function": {
                            "name": message.tool_name(),
                            "arguments": json.dumps(
                                message.arguments(),
                                ensure_ascii=False,
                                separators=(",", ":"),
                            ),
                        },
                    }
                )
            else:
                raise ValueError(f"Unknown conversation message: {type(message).__name__}")

        flush_assistant()
        flush_attachments()
        return result

    @classmethod
    def _serialize_user_content(
        cls,
        text: str,
        media_items: list[Media],
    ) -> str | list[dict[str, object]]:
        if not media_items:
            return text
        content: list[dict[str, object]] = []
        if text:
            content.append({"type": "text", "text": text})
        content.extend(cls._serialize_media(media) for media in media_items)
        return content

    @staticmethod
    def _serialize_media(media: Media) -> dict[str, object]:
        if media.kind == "image":
            if media.url is not None:
                url = media.url
            elif media.data is not None:
                mime_type = media.mime_type or "image/jpeg"
                encoded = base64.b64encode(media.data).decode("ascii")
                url = f"data:{mime_type};base64,{encoded}"
            else:
                raise ValueError("Chat Completions image input requires data or URL")
            return {
                "type": "image_url",
                "image_url": {"url": url, "detail": media.detail},
            }

        if media.kind == "audio":
            if media.data is None:
                raise ValueError("Chat Completions audio input requires inline data")
            mime_type = media.mime_type or ""
            audio_format = "mp3" if mime_type == "audio/mpeg" else "wav"
            return {
                "type": "input_audio",
                "input_audio": {
                    "data": base64.b64encode(media.data).decode("ascii"),
                    "format": audio_format,
                },
            }

        file: dict[str, object] = {}
        if media.file_id is not None:
            file["file_id"] = media.file_id
        elif media.data is not None:
            mime_type = media.mime_type or "application/octet-stream"
            encoded = base64.b64encode(media.data).decode("ascii")
            file["file_data"] = f"data:{mime_type};base64,{encoded}"
            file["filename"] = media.filename or "file"
        elif media.url is not None:
            file["file_data"] = media.url
            file["filename"] = media.filename or "file"
        return {"type": "file", "file": file}


class OpenRouterChatLLM(ChatCompletionsLLM):
    """Chat Completions adapter for OpenRouter.

    OpenRouter rejects Z.AI's thinking body. Reasoning is requested with
    OpenRouter's reasoning.effort parameter, and reasoning_details are echoed
    back on later tool turns.
    """

    async def ask(
        self,
        conversation: Conversation,
        tools: list[Tool],
        *,
        instructions: str | None = None,
        thinking: bool | None = None,
    ) -> list[Message]:
        messages = self._serialize_conversation(conversation, instructions)
        serialized_tools = self._serialize_tools(tools)
        request: dict[str, Any] = {
            "model": self._model,
            "messages": messages,
            "stream": False,
        }
        extra_body = self._extra_body(thinking)
        if extra_body:
            request["extra_body"] = extra_body
        if serialized_tools:
            request["tools"] = serialized_tools

        response = await self._ai.chat.completions.create(**cast(Any, request))
        if not response.choices:
            raise RuntimeError("Chat Completions API returned no choices")
        return self._parse_message(response.choices[0].message)

    @staticmethod
    def _extra_body(thinking: bool | None) -> dict[str, object]:
        if not thinking:
            return {}
        return {"reasoning": {"effort": "medium"}}

    @staticmethod
    def _reasoning_fields(message: ReasoningMessage) -> dict[str, object]:
        fields: dict[str, object] = {}
        item = message.provider_item()
        if isinstance(item, dict):
            details = item.get("reasoning_details")
            if isinstance(details, list) and details:
                fields["reasoning_details"] = details
        text = message.text()
        if text:
            fields["reasoning"] = text
        return fields

    @classmethod
    def _parse_message(cls, message: Any) -> list[Message]:
        result: list[Message] = []
        extra = getattr(message, "model_extra", None)
        if not isinstance(extra, dict):
            extra = {}
        reasoning = (
            getattr(message, "reasoning", None)
            or getattr(message, "reasoning_content", None)
            or extra.get("reasoning")
            or extra.get("reasoning_content")
        )
        details = getattr(message, "reasoning_details", None)
        if details is None:
            details = extra.get("reasoning_details")
        provider_item = None
        if details:
            provider_item = {"reasoning_details": cls._jsonable(details)}
        if reasoning or provider_item:
            result.append(ReasoningMessage(str(reasoning or ""), provider_item=provider_item))
        if message.content:
            result.append(LlmMessage(message.content))

        for call in message.tool_calls or ():
            try:
                arguments = json.loads(call.function.arguments)
            except (TypeError, ValueError) as exc:
                raise ValueError("Model returned invalid tool-call arguments") from exc
            if not isinstance(arguments, dict):
                raise ValueError("Model tool-call arguments must be an object")
            result.append(ToolCallMessage(call.id, call.function.name, arguments))
        return result

    @classmethod
    def _jsonable(cls, value: Any) -> Any:
        if hasattr(value, "model_dump"):
            return value.model_dump(mode="json")
        if isinstance(value, list):
            return [cls._jsonable(item) for item in value]
        if isinstance(value, dict):
            return {str(key): cls._jsonable(item) for key, item in value.items()}
        return value


class AnthropicMessagesLLM(OpenAIBasedLLM):
    """micro-agents adapter for the Anthropic Messages API."""

    _MAX_TOKENS = 16000

    def __init__(self, api_url: str, api_key: str, model: str):
        super().__init__(api_url, api_key, model)
        self._api_key = api_key
        self._api_url = api_url.rstrip("/")

    async def ask(
        self,
        conversation: Conversation,
        tools: list[Tool],
        *,
        instructions: str | None = None,
        thinking: bool | None = None,
    ) -> list[Message]:
        payload = self._build_payload(conversation, tools, instructions, thinking)
        body = await self._post(payload)
        return self._parse_response(body)

    def _messages_url(self) -> str:
        if self._api_url.endswith("/messages"):
            return self._api_url
        if self._api_url.endswith("/v1"):
            return f"{self._api_url}/messages"
        return f"{self._api_url}/v1/messages"

    def _build_payload(
        self,
        conversation: Conversation,
        tools: list[Tool],
        instructions: str | None,
        thinking: bool | None,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "model": self._model,
            "max_tokens": self._MAX_TOKENS,
            "messages": self._serialize_messages(conversation),
            "thinking": (
                {"type": "adaptive", "display": "summarized"}
                if thinking
                else {"type": "disabled"}
            ),
        }
        if instructions:
            payload["system"] = instructions
        serialized_tools = self._serialize_anthropic_tools(tools)
        if serialized_tools:
            payload["tools"] = serialized_tools
        return payload

    def _serialize_anthropic_tools(self, tools: list[Tool]) -> list[dict[str, object]]:
        result: list[dict[str, object]] = []
        for tool in self._serialize_tools(tools):
            item: dict[str, object] = {
                "name": tool["name"],
                "input_schema": tool.get("parameters") or {"type": "object", "properties": {}},
            }
            description = tool.get("description")
            if description:
                item["description"] = description
            result.append(item)
        return result

    def _serialize_messages(self, conversation: Conversation) -> list[dict[str, object]]:
        messages: list[dict[str, object]] = []
        assistant_blocks: list[dict[str, object]] = []

        def flush_assistant() -> None:
            if not assistant_blocks:
                return
            messages.append({"role": "assistant", "content": list(assistant_blocks)})
            assistant_blocks.clear()

        def append_user_blocks(blocks: list[dict[str, object]]) -> None:
            flush_assistant()
            if not blocks:
                return
            if messages and messages[-1]["role"] == "user":
                content = messages[-1]["content"]
                if isinstance(content, str):
                    content = [{"type": "text", "text": content}]
                    messages[-1]["content"] = content
                if isinstance(content, list):
                    content.extend(blocks)
                    return
            messages.append({"role": "user", "content": blocks})

        for message in conversation.messages:
            if isinstance(message, ToolCallResponseMessage):
                append_user_blocks([self._serialize_tool_result(message)])
                continue
            if isinstance(message, UserMessage):
                append_user_blocks(self._serialize_user_blocks(message))
                continue
            if isinstance(message, ReasoningMessage):
                block = self._serialize_reasoning_block(message)
                if block is not None:
                    assistant_blocks.append(block)
                continue
            if isinstance(message, LlmMessage):
                text = message.message()
                if text:
                    assistant_blocks.append({"type": "text", "text": text})
                continue
            if isinstance(message, ToolCallMessage):
                assistant_blocks.append(
                    {
                        "type": "tool_use",
                        "id": message.unique_id(),
                        "name": message.tool_name(),
                        "input": message.arguments(),
                    }
                )
                continue
            raise ValueError(f"Unknown conversation message: {type(message).__name__}")

        flush_assistant()
        return messages

    @staticmethod
    def _serialize_reasoning_block(message: ReasoningMessage) -> dict[str, object] | None:
        item = message.provider_item()
        if isinstance(item, dict) and item.get("type") in {"thinking", "redacted_thinking"}:
            return dict(item)
        text = message.text()
        if not text:
            return None
        return {"type": "text", "text": text}

    def _serialize_user_blocks(self, message: UserMessage) -> list[dict[str, object]]:
        blocks: list[dict[str, object]] = []
        if message.message():
            blocks.append({"type": "text", "text": message.message()})
        blocks.extend(self._serialize_media_block(media) for media in message.media())
        return blocks

    def _serialize_tool_result(self, message: ToolCallResponseMessage) -> dict[str, object]:
        attachments = message.attachments()
        if attachments:
            content: object = []
            if message.text():
                cast(list[dict[str, object]], content).append(
                    {"type": "text", "text": message.text()}
                )
            cast(list[dict[str, object]], content).extend(
                self._serialize_media_block(media) for media in attachments
            )
        else:
            content = message.text()
        block: dict[str, object] = {
            "type": "tool_result",
            "tool_use_id": message.unique_id(),
            "content": content,
        }
        if message.is_error():
            block["is_error"] = True
        return block

    @staticmethod
    def _serialize_media_block(media: Media) -> dict[str, object]:
        if media.file_id is not None:
            raise ValueError("Claude media input does not support file ids")
        if media.kind == "audio":
            raise ValueError("Claude Messages API does not support audio input")
        if media.kind == "image":
            return {
                "type": "image",
                "source": AnthropicMessagesLLM._media_source(media, "image/jpeg"),
            }
        mime_type = media.mime_type or "application/octet-stream"
        if mime_type.startswith("text/") and media.data is not None:
            return {
                "type": "text",
                "text": media.data.decode("utf-8", errors="replace"),
            }
        block: dict[str, object] = {
            "type": "document",
            "source": AnthropicMessagesLLM._media_source(media, "application/pdf"),
        }
        if media.filename:
            block["title"] = media.filename
        return block

    @staticmethod
    def _media_source(media: Media, default_mime: str) -> dict[str, object]:
        if media.url is not None:
            return {"type": "url", "url": media.url}
        if media.data is None:
            raise ValueError("Claude media input requires data or URL")
        return {
            "type": "base64",
            "media_type": media.mime_type or default_mime,
            "data": base64.b64encode(media.data).decode("ascii"),
        }

    async def _post(self, payload: dict[str, Any]) -> dict[str, Any]:
        timeout = aiohttp.ClientTimeout(total=300)
        headers = {
            "x-api-key": self._api_key,
            "anthropic-version": "2023-06-01",
            "content-type": "application/json",
        }
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.post(self._messages_url(), headers=headers, json=payload) as response:
                raw = await response.text()
                try:
                    body = json.loads(raw) if raw else {}
                except json.JSONDecodeError as exc:
                    raise RuntimeError(
                        f"Claude API returned invalid JSON (HTTP {response.status})"
                    ) from exc
                if response.status >= 400:
                    error = body.get("error") if isinstance(body, dict) else None
                    message = error.get("message") if isinstance(error, dict) else None
                    raise RuntimeError(message or f"Claude API returned HTTP {response.status}")
        if not isinstance(body, dict):
            raise RuntimeError("Claude API returned an unexpected response")
        return body

    @staticmethod
    def _parse_response(body: dict[str, Any]) -> list[Message]:
        content = body.get("content")
        if not isinstance(content, list):
            raise RuntimeError("Claude API returned no content")

        result: list[Message] = []
        for block in content:
            if not isinstance(block, dict):
                raise RuntimeError("Claude API returned an invalid content block")
            block_type = block.get("type")
            if block_type in {"thinking", "redacted_thinking"}:
                text = block.get("thinking") if block_type == "thinking" else ""
                result.append(ReasoningMessage(str(text or ""), provider_item=block))
            elif block_type == "text":
                text = block.get("text")
                if isinstance(text, str) and text:
                    result.append(LlmMessage(text))
            elif block_type == "tool_use":
                call_id = block.get("id")
                name = block.get("name")
                arguments = block.get("input")
                if not isinstance(call_id, str) or not isinstance(name, str):
                    raise ValueError("Claude returned an invalid tool call")
                if not isinstance(arguments, dict):
                    raise ValueError("Claude tool-call arguments must be an object")
                result.append(ToolCallMessage(call_id, name, arguments))
            else:
                continue

        if not result:
            raise RuntimeError("Claude API returned no message or tool call")
        return result