from __future__ import annotations

import logging
import time
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from yuxi.channel.extensions.rocketchat.outbound import RocketChatOutbound

logger = logging.getLogger(__name__)


class RocketChatEditStream:
    def __init__(
        self,
        outbound: RocketChatOutbound,
        room_id: str,
        thread_id: str | None = None,
        max_chars: int = 4000,
        throttle_ms: int = 1200,
    ):
        self.outbound = outbound
        self.room_id = room_id
        self.thread_id = thread_id
        self.max_chars = max_chars
        self.throttle_ms = max(throttle_ms, 250)
        self.stream_msg_id: str | None = None
        self.last_update_at: float = 0
        self._dirty: bool = False

    async def update(self, text: str) -> None:
        now = time.monotonic()
        elapsed = now - self.last_update_at
        if elapsed < self.throttle_ms / 1000:
            self._dirty = True
            return

        if self._dirty or elapsed >= self.throttle_ms / 1000:
            await self._flush(text)

    async def _flush(self, text: str) -> None:
        truncated = text[: self.max_chars]
        try:
            if self.stream_msg_id:
                await self.outbound.edit_message(
                    self.room_id,
                    self.stream_msg_id,
                    truncated,
                )
            else:
                resp = await self.outbound.client.post_message(
                    self.room_id,
                    truncated + "\n\n_Thinking…_",
                    thread_id=self.thread_id,
                )
                self.stream_msg_id = resp.get("message", {}).get("_id", "")
        except Exception as e:
            logger.debug("Stream flush error: %s", e)
        self.last_update_at = time.monotonic()
        self._dirty = False

    async def finalize(self, final_text: str) -> None:
        if self.stream_msg_id:
            try:
                await self.outbound.edit_message(
                    self.room_id,
                    self.stream_msg_id,
                    final_text[: self.max_chars],
                )
            except Exception as e:
                logger.debug("Stream finalize error: %s", e)
        else:
            try:
                await self.outbound.send_text(
                    self.room_id,
                    final_text,
                    thread_id=self.thread_id,
                )
            except Exception as e:
                logger.debug("Stream finalize send error: %s", e)

    async def cancel(self) -> None:
        if self.stream_msg_id:
            try:
                await self.outbound.delete_message(self.room_id, self.stream_msg_id)
            except Exception:
                pass
            self.stream_msg_id = None