ForcePilot/backend/package/yuxi/channels/utils/streaming_fallback.py

from __future__ import annotations

import asyncio
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any

from yuxi.channels.models import DeliveryResult
from yuxi.utils.logging_config import logger

STREAM_CONFIG_DEFAULTS = {
    "mode": "partial",
    "preview": {
        "enabled": True,
        "chunk": {
            "min_chars": 50,
            "max_chars": 200,
            "break_preference": "newline",
        },
        "tool_progress": True,
    },
    "block": {
        "enabled": True,
        "coalesce": False,
        "coalesce_min_chars": 1500,
        "coalesce_idle_ms": 1000,
    },
    "typing_indicator": {
        "enabled": True,
        "duration_ms": 5000,
    },
    "fallback": {
        "on_edit_failure": "send_final",
        "on_rate_limit": "send_final",
    },
}


@dataclass
class StreamingConfig:
    mode: str = "off"
    preview_enabled: bool = True
    preview_chunk_min_chars: int = 50
    preview_chunk_max_chars: int = 200
    preview_chunk_break_preference: str = "newline"
    preview_tool_progress: bool = True
    block_enabled: bool = True
    block_coalesce: bool = False
    block_coalesce_min_chars: int = 1500
    block_coalesce_idle_ms: int = 1000
    typing_indicator_enabled: bool = True
    typing_indicator_duration_ms: int = 5000
    fallback_on_edit_failure: str = "send_final"
    fallback_on_rate_limit: str = "send_final"

    @classmethod
    def from_config(cls, config: dict[str, Any] | None) -> StreamingConfig:
        if not config:
            return cls()

        streaming = config.get("streaming", {})
        if not isinstance(streaming, dict):
            streaming = {"mode": str(streaming)} if streaming else {}

        mode = streaming.get("mode", config.get("streaming_mode", "off"))

        preview = streaming.get("preview", {})
        if not isinstance(preview, dict):
            preview = {}
        chunk = preview.get("chunk", {})
        if not isinstance(chunk, dict):
            chunk = {}

        block = streaming.get("block", {})
        if not isinstance(block, dict):
            block = {}

        typing_cfg = streaming.get("typing_indicator", {})
        if not isinstance(typing_cfg, dict):
            typing_cfg = {}

        fallback = streaming.get("fallback", {})
        if not isinstance(fallback, dict):
            fallback = {}

        return cls(
            mode=mode,
            preview_enabled=preview.get("enabled", True),
            preview_chunk_min_chars=chunk.get("min_chars", 50),
            preview_chunk_max_chars=chunk.get("max_chars", 200),
            preview_chunk_break_preference=chunk.get("break_preference", "newline"),
            preview_tool_progress=preview.get("tool_progress", True),
            block_enabled=block.get("enabled", True),
            block_coalesce=block.get("coalesce", False),
            block_coalesce_min_chars=block.get("coalesce_min_chars", 1500),
            block_coalesce_idle_ms=block.get("coalesce_idle_ms", 1000),
            typing_indicator_enabled=typing_cfg.get("enabled", True),
            typing_indicator_duration_ms=typing_cfg.get("duration_ms", 5000),
            fallback_on_edit_failure=fallback.get("on_edit_failure", "send_final"),
            fallback_on_rate_limit=fallback.get("on_rate_limit", "send_final"),
        )


SendFn = Callable[[str], Awaitable[DeliveryResult]]
TypingFn = Callable[[str, bool], Awaitable[DeliveryResult | None]]
EditFn = Callable[[str, str, str, bool], Awaitable[DeliveryResult]]


class SmartCoalesceBuffer:
    def __init__(
        self,
        min_chars: int = 1500,
        idle_ms: int = 1000,
        max_wait_ms: int = 5000,
    ):
        self._min_chars = min_chars
        self._idle_ms = idle_ms
        self._max_wait_ms = max_wait_ms
        self._buffer: str = ""
        self._last_feed: float = 0.0
        self._first_feed: float = 0.0
        self._has_pending: bool = False

    def feed(self, text: str) -> str | None:
        now = time.monotonic()
        if not self._buffer:
            self._first_feed = now
            self._buffer = text
            self._last_feed = now
            self._has_pending = True
            return None

        self._buffer += text
        self._last_feed = now

        if len(self._buffer) >= self._min_chars:
            result = self._buffer
            self._buffer = ""
            self._has_pending = False
            return result

        elapsed_since_first = (now - self._first_feed) * 1000
        if elapsed_since_first >= self._max_wait_ms:
            result = self._buffer
            self._buffer = ""
            self._has_pending = False
            return result

        return None

    def peek(self) -> str:
        return self._buffer

    def flush(self) -> str:
        result = self._buffer
        self._buffer = ""
        self._has_pending = False
        return result

    @property
    def pending(self) -> bool:
        return self._has_pending

    def reset(self) -> None:
        self._buffer = ""
        self._first_feed = 0.0
        self._last_feed = 0.0
        self._has_pending = False


class StreamPhase(StrEnum):
    INIT = "init"
    TYPING = "typing"
    STREAMING = "streaming"
    FINALIZING = "finalizing"
    ERROR = "error"
    FALLBACK_SEND = "fallback_send"
    DONE = "done"


@dataclass
class StreamContext:
    chat_id: str
    msg_id: str = ""
    phase: StreamPhase = StreamPhase.INIT
    accumulated_text: str = ""
    started_at: float = field(default_factory=time.monotonic)
    last_update_at: float = 0.0
    error: str = ""
    fallback_used: bool = False


class StreamStateMachine:
    def __init__(self):
        self._ctx: dict[str, StreamContext] = {}

    def get(self, chat_id: str) -> StreamContext | None:
        return self._ctx.get(chat_id)

    def start(self, chat_id: str, msg_id: str = "") -> StreamContext:
        ctx = StreamContext(chat_id=chat_id, msg_id=msg_id)
        self._ctx[chat_id] = ctx
        return ctx

    async def transition_typing(self, chat_id: str, typing_fn: TypingFn | None = None) -> None:
        ctx = self._ctx.get(chat_id)
        if ctx is None:
            return
        ctx.phase = StreamPhase.TYPING
        if typing_fn:
            try:
                await typing_fn(chat_id, True)
            except Exception:
                pass

    async def transition_streaming(
        self,
        chat_id: str,
        send_fn: SendFn | None = None,
        initial_text: str = "",
    ) -> None:
        ctx = self._ctx.get(chat_id)
        if ctx is None:
            return
        ctx.phase = StreamPhase.STREAMING
        ctx.accumulated_text = initial_text
        ctx.last_update_at = time.monotonic()

    async def transition_error(
        self,
        chat_id: str,
        error: str,
        fallback_fn: SendFn | None = None,
        final_text: str = "",
    ) -> DeliveryResult:
        ctx = self._ctx.get(chat_id)
        if ctx is None:
            return DeliveryResult(success=False, error=error)

        ctx.phase = StreamPhase.ERROR
        ctx.error = error
        ctx.fallback_used = True

        if fallback_fn and final_text:
            logger.warning(f"Stream error for {chat_id}: {error}, falling back to send_final")
            result = await fallback_fn(final_text)
            ctx.phase = StreamPhase.DONE
            return result

        return DeliveryResult(success=False, error=error)

    async def finalize(
        self,
        chat_id: str,
        typing_fn: TypingFn | None = None,
    ) -> StreamContext | None:
        ctx = self._ctx.get(chat_id)
        if ctx is None:
            return None

        ctx.phase = StreamPhase.FINALIZING

        if typing_fn:
            try:
                await typing_fn(chat_id, False)
            except Exception:
                pass

        ctx.phase = StreamPhase.DONE
        return ctx

    def cleanup(self, chat_id: str) -> StreamContext | None:
        return self._ctx.pop(chat_id, None)

    def is_active(self, chat_id: str) -> bool:
        ctx = self._ctx.get(chat_id)
        return ctx is not None and ctx.phase not in (
            StreamPhase.DONE,
            StreamPhase.ERROR,
        )

    def clear(self) -> None:
        self._ctx.clear()

    @property
    def active_count(self) -> int:
        return sum(1 for ctx in self._ctx.values() if ctx.phase not in (StreamPhase.DONE, StreamPhase.ERROR))


class UniversalChunkedSender:
    def __init__(
        self,
        send_fn: SendFn,
        chunk_size: int = 2000,
        prefix_format: str = "",
        rate_limit_delay_ms: int = 0,
        max_chunks: int = 10,
    ):
        self._send_fn = send_fn
        self._chunk_size = max(1, chunk_size)
        self._prefix_format = prefix_format
        self._rate_limit_delay_ms = rate_limit_delay_ms
        self._max_chunks = max_chunks
        self._sent_chunks: dict[str, int] = {}
        self._lock = asyncio.Lock()

    async def send_stream(
        self,
        chat_id: str,
        text: str,
        finished: bool = False,
    ) -> list[DeliveryResult]:
        results: list[DeliveryResult] = []

        async with self._lock:
            sent_count = self._sent_chunks.get(chat_id, 0)
            total_estimate = sent_count + 1 if not finished else sent_count + 1

            content = self._build_content(text, sent_count + 1, total_estimate)

        result = await self._send_fn(content)
        results.append(result)

        if result.success:
            async with self._lock:
                self._sent_chunks[chat_id] = sent_count + 1

        if finished:
            async with self._lock:
                self._sent_chunks.pop(chat_id, None)

            if self._prefix_format:
                summary = self._prefix_format.format_map({"total": sent_count + 1})
                if summary and summary != content:
                    summary_result = await self._send_fn(summary)
                    results.append(summary_result)

        return results

    async def send_text_chunked(
        self,
        chat_id: str,
        text: str,
        finished: bool = False,
    ) -> list[DeliveryResult]:
        results: list[DeliveryResult] = []
        paragraphs = text.split("\n\n")
        accumulated = ""
        chunk_idx = 0

        for para in paragraphs:
            accumulated += para + "\n\n"
            if len(accumulated) >= self._chunk_size:
                chunk_idx += 1
                if chunk_idx > self._max_chunks:
                    accumulated = accumulated[: self._chunk_size]
                    chunk_result = await self._send_fn(accumulated)
                    results.append(chunk_result)
                    accumulated = ""
                    continue

                content = accumulated.strip()
                if self._prefix_format:
                    content = f"{content}\n\n{self._prefix_format}"
                chunk_result = await self._send_fn(content)
                results.append(chunk_result)
                accumulated = ""

                if self._rate_limit_delay_ms > 0:
                    await asyncio.sleep(self._rate_limit_delay_ms / 1000.0)

        if accumulated.strip():
            chunk_idx += 1
            chunk_result = await self._send_fn(accumulated.strip())
            results.append(chunk_result)

        return results

    def _build_content(self, text: str, chunk_num: int, total: int) -> str:
        if not self._prefix_format:
            return text

        prefix = self._prefix_format.format_map({"chunk": str(chunk_num), "total": str(total)})

        return f"{prefix}\n{text}"


__all__ = [
    "SmartCoalesceBuffer",
    "StreamPhase",
    "StreamContext",
    "StreamStateMachine",
    "UniversalChunkedSender",
    "StreamingConfig",
    "STREAM_CONFIG_DEFAULTS",
    "SendFn",
    "TypingFn",
    "EditFn",
]