import html
import logging
from collections.abc import AsyncIterable, Callable
from typing import Literal

from pydantic_ai import (
    AgentStreamEvent,
    FunctionToolCallEvent,
    FunctionToolResultEvent,
    PartDeltaEvent,
    PartStartEvent,
    RunContext,
    TextPartDelta,
    ThinkingPartDelta,
)
from telegram import Bot, Message
from telegram.error import BadRequest

from app.agent.deps import Deps
from app.channels.telegram.formatting import (
    markdown_to_telegram_html,
    normalize_newlines,
    split_for_telegram,
)
from app.settings import MessagesSettings

logger = logging.getLogger(__name__)


class MessageStreamer:
    def __init__(
        self,
        bot: Bot,
        chat_id: int,
        messages: Callable[[], MessagesSettings] | None = None,
    ) -> None:
        self._bot = bot
        self._chat_id = chat_id
        self._messages: Callable[[], MessagesSettings] = messages or (lambda: MessagesSettings())
        self._current_type: Literal["thinking", "text", "tool"] = "thinking"
        self._accumulated_text: str = ""
        self._current_msg: Message | None = None
        self._tool_names: list[str] = []

    async def start(self) -> None:
        if self._messages().display_tools:
            await self._new_message("💭 <i>Pensando...</i>")

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

    async def send_file(self, fname: str, fdata: bytes) -> None:
        await self._bot.send_document(chat_id=self._chat_id, document=fdata, filename=fname)

    async def event_handler(
        self,
        _ctx: RunContext[Deps],
        stream: AsyncIterable[AgentStreamEvent],
    ) -> None:
        async for event in stream:
            if isinstance(event, PartStartEvent):
                part_content: object = getattr(event.part, "content", None)
                if event.part.part_kind == "text":
                    if self._current_type == "tool":
                        self._current_msg = None
                        self._tool_names = []
                    elif self._current_type == "text" and self._accumulated_text:
                        await self._flush()
                        self._current_msg = None
                    self._current_type = "text"
                    if part_content:
                        self._accumulated_text += normalize_newlines(str(part_content))
                continue

            if isinstance(event, FunctionToolCallEvent):
                if self._current_type == "text" and self._accumulated_text:
                    await self._flush()
                    self._current_msg = None
                if self._current_type != "tool":
                    self._tool_names = []
                self._current_type = "tool"
                self._tool_names.append(event.part.tool_name)
                if self._messages().display_tools:
                    tools_text = " > ".join(html.escape(n) for n in self._tool_names)
                    label = "Ferramenta" if len(self._tool_names) == 1 else "Ferramentas"
                    msg_text = f"⚙️ <b>{label}:</b> <code>{tools_text}</code>"
                    try:
                        if self._current_msg is not None:
                            await self._current_msg.edit_text(msg_text, parse_mode="HTML")
                        else:
                            self._current_msg = await self._bot.send_message(
                                chat_id=self._chat_id, text=msg_text, parse_mode="HTML"
                            )
                    except Exception:
                        logger.exception("error updating tool call message")
                continue

            if isinstance(event, FunctionToolResultEvent):
                continue

            if isinstance(event, PartDeltaEvent):
                delta = event.delta

                if isinstance(delta, ThinkingPartDelta):
                    if self._current_type != "thinking":
                        if self._current_type == "tool":
                            self._current_msg = None
                            self._tool_names = []
                        self._current_type = "thinking"
                        if self._messages().display_tools:
                            try:
                                if self._current_msg is not None:
                                    await self._current_msg.edit_text(
                                        "💭 <i>Pensando...</i>",
                                        parse_mode="HTML",
                                    )
                            except Exception:
                                logger.exception("error updating thinking message")
                    continue

                if isinstance(delta, TextPartDelta):
                    if self._current_type == "tool":
                        self._current_msg = None
                        self._tool_names = []
                    self._current_type = "text"
                    self._accumulated_text += normalize_newlines(str(delta.content_delta or ""))

    async def _flush(self) -> None:
        if not self._accumulated_text:
            return

        try:
            chunks = split_for_telegram(self._accumulated_text)
            for index, chunk in enumerate(chunks):
                rendered_chunk = markdown_to_telegram_html(chunk)
                if index == 0 and self._current_msg is not None:
                    await self._current_msg.edit_text(rendered_chunk, parse_mode="HTML")
                else:
                    self._current_msg = await self._bot.send_message(
                        chat_id=self._chat_id, text=rendered_chunk, parse_mode="HTML"
                    )
        except BadRequest as error:
            if "Message is not modified" not in str(error):
                logger.warning("flush BadRequest: %s", error)
        except Exception:
            logger.exception("flush error")
        finally:
            self._accumulated_text = ""

    async def _new_message(self, placeholder: str) -> None:
        self._current_msg = await self._bot.send_message(
            chat_id=self._chat_id, text=placeholder, parse_mode="HTML"
        )
