ForcePilot/backend/package/yuxi/channels/adapters/telegram/streaming.py

from __future__ import annotations

import asyncio
import re
import time
from typing import Any

from telegram import Bot
from telegram.error import BadRequest, RetryAfter

from yuxi.channels.adapters.telegram.format import markdown_to_html, strip_all_tags
from yuxi.channels.adapters.telegram.stream.lane_delivery import LaneDelivery, LaneType
from yuxi.channels.adapters.telegram.stream.reasoning_lane_coordinator import ReasoningLaneCoordinator
from yuxi.utils.logging_config import logger

_STREAM_STORAGE: dict[str, dict] = {}
_STREAM_TTL_SECONDS = 300
_STREAM_STORAGE_LOCK = asyncio.Lock()

_lane_delivery = LaneDelivery()
_reasoning_coordinator = ReasoningLaneCoordinator()

_BREAK_BOUNDARY_SENTENCE = re.compile(r"[。！？.!?\n]")
_BREAK_BOUNDARY_WORD = re.compile(r"\s")


def streaming_mode_from_config(config: dict[str, Any]) -> str:
    streaming_cfg = config.get("streaming", {})
    if isinstance(streaming_cfg, dict):
        return streaming_cfg.get("mode", config.get("streaming_mode", "partial"))
    if isinstance(streaming_cfg, str):
        return streaming_cfg
    return config.get("streaming_mode", "partial")


def _resolve_stream_config(config: dict[str, Any]) -> dict[str, Any]:
    streaming_cfg = config.get("streaming", {})
    if not isinstance(streaming_cfg, dict):
        streaming_cfg = {}
    preview_cfg = streaming_cfg.get("preview", {})
    if not isinstance(preview_cfg, dict):
        preview_cfg = {}
    chunk_cfg = preview_cfg.get("chunk", {})
    if not isinstance(chunk_cfg, dict):
        chunk_cfg = {}
    block_cfg = streaming_cfg.get("block", {})
    if not isinstance(block_cfg, dict):
        block_cfg = {}

    return {
        "preview_chunk_min_chars": chunk_cfg.get("min_chars", chunk_cfg.get("minChars", 50)),
        "preview_chunk_max_chars": chunk_cfg.get("max_chars", chunk_cfg.get("maxChars", 200)),
        "preview_chunk_break_preference": chunk_cfg.get(
            "break_preference", chunk_cfg.get("breakPreference", "newline")
        ),
        "preview_tool_progress": preview_cfg.get("tool_progress", preview_cfg.get("toolProgress", True)),
        "block_enabled": block_cfg.get("enabled", True),
        "block_coalesce": block_cfg.get("coalesce", False),
    }


def _cleanup_expired(now: float) -> None:
    expired = [key for key, state in _STREAM_STORAGE.items() if now - state.get("created_at", 0) > _STREAM_TTL_SECONDS]
    for key in expired:
        del _STREAM_STORAGE[key]


def _format_block(content: str, max_chars: int = 4096) -> str:
    if len(content) <= max_chars:
        return content
    return content[: max_chars - 4] + "..."


def _format_preview_chunk(content: str, stream_cfg: dict[str, Any]) -> str:
    min_chars = stream_cfg.get("preview_chunk_min_chars", 50)
    max_chars = stream_cfg.get("preview_chunk_max_chars", 200)
    break_pref = stream_cfg.get("preview_chunk_break_preference", "newline")

    if len(content) <= max_chars:
        return content

    if break_pref == "any":
        return content[:max_chars] + "..."

    lo = max(min_chars, max_chars // 2)
    hi = max_chars

    if break_pref == "sentence":
        pattern = _BREAK_BOUNDARY_SENTENCE
    elif break_pref == "word":
        pattern = _BREAK_BOUNDARY_WORD
    else:
        pattern = re.compile(r"\n")

    candidates = [m.start() for m in pattern.finditer(content, lo, hi)]
    if candidates:
        return content[: candidates[-1] + 1].rstrip() + "..."

    return content[:max_chars] + "..."


def _format_progress(content: str, state: dict) -> str:
    counter = state.get("edit_counter", 0) + 1
    state["edit_counter"] = counter
    dots = "." * ((counter % 3) + 1)
    if len(content) > 3800:
        content = content[:3800]
    return f"{content}\n\n\u23f3 {dots}"


def _format_tool_progress(content: str, stream_cfg: dict[str, Any]) -> str:
    if not stream_cfg.get("preview_tool_progress", True):
        return content
    return f"\U0001f6e0\ufe0f {content}"


async def stream_send_chunked(
    bot: Bot,
    chat_id: str,
    content: str,
    message_id: str | None = None,
    finished: bool = False,
    config: dict[str, Any] | None = None,
) -> str | None:
    config = config or {}
    stream_cfg = _resolve_stream_config(config)

    mode = streaming_mode_from_config(config)
    edit_interval = config.get("stream_edit_interval_ms", 500) / 1000
    stream_key = f"{chat_id}:{message_id}" if message_id else None

    async with _STREAM_STORAGE_LOCK:
        state = _STREAM_STORAGE.get(stream_key) if stream_key else None
        _cleanup_expired(time.monotonic())

    if finished and state:
        async with _STREAM_STORAGE_LOCK:
            del _STREAM_STORAGE[stream_key]
        if state.get("message_id"):
            try:
                final_content = content
                if mode == "progress":
                    final_content = _format_progress(content, state).rstrip("\n\u23f3 .")
                elif mode == "block" and stream_cfg.get("block_coalesce"):
                    final_content = state.get("accumulated", "") + content

                formatted = markdown_to_html(final_content)
                await bot.edit_message_text(
                    chat_id=chat_id,
                    message_id=state["message_id"],
                    text=formatted,
                    parse_mode="HTML",
                )
            except BadRequest:
                try:
                    clean = strip_all_tags(content)
                    await bot.edit_message_text(
                        chat_id=chat_id,
                        message_id=state["message_id"],
                        text=clean,
                    )
                except Exception:
                    pass

            if state.get("reasoning_state"):
                combined = _reasoning_coordinator.get_combined_text(chat_id)
                if combined and state["reasoning_state"] != stream_key:
                    _reasoning_coordinator.cleanup(chat_id)

        return None

    if not message_id:
        try:
            if mode == "block":
                if stream_cfg.get("block_coalesce"):
                    preview = _format_preview_chunk(content, stream_cfg)
                else:
                    preview = _format_block(content, 100) + "..." if len(content) > 100 else content
            elif mode == "progress":
                preview = content[:50] + "..." if len(content) > 50 else content
            else:
                preview = _format_preview_chunk(content, stream_cfg)

            msg = await bot.send_message(chat_id=chat_id, text=preview)
            now = time.monotonic()
            async with _STREAM_STORAGE_LOCK:
                _STREAM_STORAGE[f"{chat_id}:{str(msg.message_id)}"] = {
                    "message_id": msg.message_id,
                    "last_edit": now,
                    "counter": 0,
                    "edit_counter": 0,
                    "created_at": now,
                    "mode": mode,
                    "accumulated": "",
                }
            return str(msg.message_id)
        except Exception:
            logger.exception("[Telegram] Failed to create stream placeholder")
            return None

    if not state:
        return message_id

    if stream_cfg.get("block_coalesce"):
        state["accumulated"] = (state.get("accumulated", "") + "\n" + content).strip()
        return message_id

    counter = state.get("counter", 0)
    now = time.monotonic()
    if counter < 3 and (now - state.get("last_edit", 0)) < edit_interval:
        state["counter"] = counter + 1
        return message_id

    state["counter"] = 0
    state["last_edit"] = now

    display_content = content
    if mode == "block":
        display_content = _format_block(content)
    elif mode == "progress":
        display_content = _format_progress(content, state)

    try:
        formatted = markdown_to_html(display_content)
        await bot.edit_message_text(
            chat_id=chat_id,
            message_id=state["message_id"],
            text=formatted,
            parse_mode="HTML",
        )
    except RetryAfter as e:
        await asyncio.sleep(e.retry_after)
    except BadRequest:
        pass
    except Exception:
        logger.exception(f"[Telegram] Stream edit failed for {chat_id}")

    return message_id


async def stream_send_with_lane(
    bot: Bot,
    chat_id: str,
    content: str,
    message_id: str | None = None,
    finished: bool = False,
    config: dict[str, Any] | None = None,
) -> str | None:
    config = config or {}
    lane_enabled = (
        config.get("streaming", {}).get("lane_enabled", False) if isinstance(config.get("streaming"), dict) else False
    )

    if not lane_enabled:
        return await stream_send_chunked(bot, chat_id, content, message_id, finished, config)

    reasoning_enabled = _lane_delivery.is_reasoning_enabled()

    if reasoning_enabled:
        lane = _lane_delivery.extract_lane(content)
        if lane == LaneType.REASONING:
            state = _reasoning_coordinator.get_state(chat_id)
            if not state:
                _reasoning_coordinator.start_reasoning(chat_id)
            _reasoning_coordinator.append_reasoning(chat_id, content)
            if not message_id:
                return await stream_send_chunked(bot, chat_id, content, None, False, config)
            return message_id

        if lane == LaneType.MAIN:
            combined = _reasoning_coordinator.get_combined_text(chat_id)
            if combined:
                content = combined

    if finished:
        _reasoning_coordinator.complete(chat_id)

    return await stream_send_chunked(bot, chat_id, content, message_id, finished, config)