import base64
import mimetypes
from collections.abc import AsyncIterable, Awaitable, Callable
from typing import Literal

from pydantic_ai import (
    AgentStreamEvent,
    FunctionToolCallEvent,
    PartDeltaEvent,
    PartStartEvent,
    RunContext,
    TextPartDelta,
)

from app.agent.deps import Deps
from app.channels.whatsapp.formatting import markdown_to_whatsapp_text, normalize_newlines
from app.settings import MessagesSettings

SendJson = Callable[[dict[str, object]], Awaitable[None]]


class MessageStreamer:
    def __init__(
        self,
        send_json: SendJson,
        to_jid: str,
        messages: Callable[[], MessagesSettings] | None = None,
    ) -> None:
        self._send_json: SendJson = send_json
        self._to_jid: str = to_jid
        self._messages: Callable[[], MessagesSettings] = messages or (lambda: MessagesSettings())
        self._buffer: str = ""
        self._current_type: Literal["text", "tool"] = "text"

    async def start(self) -> None:
        if self._messages().display_tools:
            await self._send_json({"action": "send_typing", "to": self._to_jid})

    async def _flush(self) -> None:
        rendered_text = markdown_to_whatsapp_text(self._buffer.strip())
        if not rendered_text:
            return
        await self._send_json(
            {
                "action": "send_text",
                "to": self._to_jid,
                "text": rendered_text,
            }
        )
        self._buffer = ""

    async def finalize(self) -> None:
        await self._flush()

    async def send_file(self, fname: str, fdata: bytes) -> None:
        mime_type, _ = mimetypes.guess_type(fname)
        await self._send_json(
            {
                "action": "send_document",
                "to": self._to_jid,
                "filename": fname,
                "mimetype": mime_type or "application/octet-stream",
                "data_base64": base64.b64encode(fdata).decode("ascii"),
            }
        )

    async def event_handler(
        self,
        _ctx: RunContext[Deps],
        stream: AsyncIterable[AgentStreamEvent],
    ) -> None:
        async for event in stream:
            if isinstance(event, PartStartEvent):
                if event.part.part_kind == "text":
                    if self._current_type == "tool" and self._messages().display_tools:
                        await self._send_json({"action": "send_typing", "to": self._to_jid})
                    self._current_type = "text"
                    part_content: object = getattr(event.part, "content", None)
                    if part_content:
                        self._buffer += normalize_newlines(str(part_content))
                continue

            if isinstance(event, FunctionToolCallEvent):
                if self._current_type == "text" and self._buffer:
                    await self._flush()
                self._current_type = "tool"
                if self._messages().display_tools:
                    await self._send_json(
                        {
                            "action": "send_text",
                            "to": self._to_jid,
                            "text": f"⚙️ *Ferramenta:* {event.part.tool_name}",
                        }
                    )
                continue

            if isinstance(event, PartDeltaEvent) and isinstance(event.delta, TextPartDelta):
                if self._current_type == "tool" and self._messages().display_tools:
                    await self._send_json({"action": "send_typing", "to": self._to_jid})
                self._current_type = "text"
                self._buffer += normalize_newlines(str(event.delta.content_delta or ""))
