ForcePilot/backend/package/yuxi/channel/streaming/block_chunker.py

from __future__ import annotations

import asyncio
import logging
import re
import time
from collections.abc import AsyncIterator, Callable
from dataclasses import dataclass
from enum import StrEnum

from yuxi.channel.streaming.fences import (
    FenceSpan,
    find_fence_span_at,
    is_safe_fence_break,
    parse_fence_spans,
)
from yuxi.channel.streaming.models import (
    BlockReplyCoalescing,
    StreamingConfig,
)

logger = logging.getLogger(__name__)

_CODE_FENCE_RE = re.compile(r"^```[a-zA-Z]*\s*$")
_SENTENCE_BREAK_RE = re.compile(r"[.!?](?=\s|$)")


class BreakPreference(StrEnum):
    PARAGRAPH = "paragraph"
    NEWLINE = "newline"
    SENTENCE = "sentence"
    WHITESPACE = "whitespace"


@dataclass
class StreamingBlock:
    text: str
    index: int
    is_final: bool = False


# ── helpers ────────────────────────────────────────────────────────────


def _skip_leading_newlines(value: str, start: int = 0) -> int:
    i = start
    while i < len(value) and value[i] == "\n":
        i += 1
    return i


def _strip_leading_newlines(value: str) -> str:
    start = _skip_leading_newlines(value)
    return value[start:] if start > 0 else value


def _find_safe_sentence_break(text: str, fence_spans: list[FenceSpan], min_chars: int, offset: int = 0) -> int:
    for match in _SENTENCE_BREAK_RE.finditer(text):
        at = match.start()
        if at < min_chars:
            continue
        candidate = at + 1
        if is_safe_fence_break(fence_spans, offset + candidate):
            return candidate
    return -1


def _find_safe_paragraph_break(
    text: str,
    fence_spans: list[FenceSpan],
    min_chars: int,
    reverse: bool,
    offset: int = 0,
) -> int:
    sep = "\n\n"
    paragraph_idx = text.rfind(sep) if reverse else text.find(sep)
    while (paragraph_idx >= min_chars) if reverse else (paragraph_idx != -1):
        candidate = paragraph_idx + len(sep)
        if candidate >= min_chars and candidate < len(text) and is_safe_fence_break(fence_spans, offset + candidate):
            return candidate
        paragraph_idx = text.rfind(sep, 0, paragraph_idx) if reverse else text.find(sep, paragraph_idx + len(sep))
    return -1


def _find_safe_newline_break(
    text: str,
    fence_spans: list[FenceSpan],
    min_chars: int,
    reverse: bool,
    offset: int = 0,
) -> int:
    char = "\n"
    newline_idx = text.rfind(char) if reverse else text.find(char)
    while (newline_idx >= min_chars) if reverse else (newline_idx != -1):
        if newline_idx >= min_chars and is_safe_fence_break(fence_spans, offset + newline_idx):
            return newline_idx
        newline_idx = text.rfind(char, 0, newline_idx) if reverse else text.find(char, newline_idx + 1)
    return -1


def _find_next_paragraph_break(
    buffer: str,
    fence_spans: list[FenceSpan],
    start_index: int = 0,
    min_chars_from_start: int = 1,
) -> tuple[int, int] | None:
    if start_index < 0:
        return None
    re_paragraph = re.compile(r"\n[\t ]*\n+")
    for match in re_paragraph.finditer(buffer):
        index = match.start()
        if index < start_index:
            continue
        if index - start_index < min_chars_from_start:
            continue
        if not is_safe_fence_break(fence_spans, index):
            continue
        return (index, len(match.group()))
    return None


def _find_fence_close_line_start(buffer: str, fence: FenceSpan, offset: int = 0) -> int:
    relative_fence_end = min(len(buffer), max(0, fence.end - offset))
    if relative_fence_end <= 0:
        return -1
    last_newline = buffer.rfind("\n", 0, relative_fence_end)
    return last_newline + 1 if last_newline >= 0 else -1


def _break_pref_order(pref: BreakPreference) -> list[BreakPreference]:
    order = [
        BreakPreference.PARAGRAPH,
        BreakPreference.NEWLINE,
        BreakPreference.SENTENCE,
        BreakPreference.WHITESPACE,
    ]
    idx = order.index(pref) if pref in order else 1
    return order[idx:] + order[:idx]


def _find_break_in_window(
    window: str,
    fence_spans: list[FenceSpan],
    min_chars: int,
    break_preference: BreakPreference,
    reverse: bool,
    offset: int = 0,
) -> int:
    for pref in _break_pref_order(break_preference):
        if pref == BreakPreference.PARAGRAPH:
            idx = _find_safe_paragraph_break(window, fence_spans, min_chars, reverse=reverse, offset=offset)
            if idx != -1:
                return idx

        if pref in (BreakPreference.PARAGRAPH, BreakPreference.NEWLINE):
            idx = _find_safe_newline_break(window, fence_spans, min_chars, reverse=reverse, offset=offset)
            if idx != -1:
                return idx

        if pref != BreakPreference.NEWLINE:
            idx = _find_safe_sentence_break(window, fence_spans, min_chars, offset=offset)
            if idx != -1:
                return idx

        if pref == BreakPreference.WHITESPACE:
            search_range = range(len(window) - 1, min_chars - 1, -1) if reverse else range(min_chars, len(window))
            for i in search_range:
                if window[i].isspace() and is_safe_fence_break(fence_spans, offset + i):
                    return i

    return -1


# ── BlockChunker ───────────────────────────────────────────────────────


class BlockChunker:
    def __init__(self, config: StreamingConfig | None = None):
        cfg = config or StreamingConfig()
        self._min_chars = cfg.block_streaming_chunk_min_chars
        self._max_chars = cfg.block_streaming_chunk_max_chars
        self._break_preference = BreakPreference(cfg.block_streaming_chunk_break_preference)
        self._coalesce_idle_ms = cfg.block_streaming_coalesce_idle_ms

    def chunk(self, text: str) -> list[StreamingBlock]:
        blocks: list[StreamingBlock] = []
        if len(text) <= self._max_chars:
            blocks.append(StreamingBlock(text=text, index=0))
            return blocks

        fence_spans = parse_fence_spans(text)
        start = 0
        idx = 0
        while start < len(text):
            end = min(start + self._max_chars, len(text))

            if end < len(text):
                end = self._find_break_point(text, start, end, fence_spans)

            blocks.append(StreamingBlock(text=text[start:end], index=idx))
            start = end
            idx += 1

        return blocks

    async def chunk_from_stream(self, stream: AsyncIterator[str]) -> AsyncIterator[StreamingBlock]:
        buffer = ""
        block_index = 0

        async for chunk in stream:
            buffer += chunk

            if len(buffer) >= self._max_chars:
                blocks = self.chunk(buffer)
                buffer = ""
                for b in blocks:
                    b.index = block_index
                    block_index += 1
                    yield b

        if buffer:
            yield StreamingBlock(text=buffer, index=block_index)

    def _find_break_point(self, text: str, start: int, end: int, fence_spans: list[FenceSpan]) -> int:
        if end - start < self._min_chars:
            return end

        search_end = start + self._min_chars

        safe_end = end
        for search_start in range(end, search_end, -1):
            segment = text[search_end:search_start]
            if self._is_safe_segment(segment):
                safe_end = search_start
                break

        search_window = text[search_end:safe_end]
        if not search_window:
            return safe_end

        idx = _find_break_in_window(
            search_window,
            fence_spans,
            min_chars=0,
            break_preference=self._break_preference,
            reverse=True,
            offset=search_end,
        )
        if idx != -1:
            return search_end + idx

        return safe_end

    @staticmethod
    def _is_safe_segment(segment: str) -> bool:
        lines = segment.split("\n")
        fence_count = sum(1 for line in lines if _CODE_FENCE_RE.match(line))
        return fence_count % 2 == 0


# ── EmbeddedBlockChunker ───────────────────────────────────────────────


class _BreakResult:
    def __init__(self, index: int, fence_split: dict | None = None):
        self.index = index
        self.fence_split = fence_split


class EmbeddedBlockChunker:
    def __init__(self, config: StreamingConfig | None = None):
        cfg = config or StreamingConfig()
        self._min_chars = cfg.block_streaming_chunk_min_chars
        self._max_chars = cfg.block_streaming_chunk_max_chars
        self._break_preference = BreakPreference(cfg.block_streaming_chunk_break_preference)
        self._flush_on_paragraph = cfg.block_streaming_flush_on_paragraph
        self._coalesce_min = cfg.block_streaming_coalesce_min_chars or self._min_chars
        self._coalesce_max = cfg.block_streaming_coalesce_max_chars or self._max_chars * 2
        self._coalesce_idle_ms = cfg.block_streaming_coalesce_idle_ms

        self._buffer = ""
        self._block_index = 0
        self._reopen_fence: FenceSpan | None = None

        self._coalesce_buffer: list[StreamingBlock] = []
        self._last_append_ts = 0.0

    def append(self, text: str) -> None:
        if text:
            self._buffer += text

    def reset(self) -> None:
        self._buffer = ""
        self._reopen_fence = None

    @property
    def buffered_text(self) -> str:
        return self._buffer

    def has_buffered(self) -> bool:
        return len(self._buffer) > 0

    def drain(self, force: bool) -> list[str]:
        if len(self._buffer) < self._min_chars and not force:
            return []

        if force and len(self._buffer) <= self._max_chars:
            stripped = self._buffer.strip()
            result = [stripped] if stripped else []
            self._buffer = ""
            self._reopen_fence = None
            return result

        source = self._buffer
        fence_spans = parse_fence_spans(source)
        start = 0
        reopen_fence: FenceSpan | None = self._reopen_fence
        chunks: list[str] = []

        while start < len(source):
            reopen_prefix = f"{reopen_fence.open_line}\n" if reopen_fence and reopen_fence.open_line else ""
            remaining_length = len(reopen_prefix) + (len(source) - start)

            if not force and remaining_length < self._min_chars:
                break

            if self._flush_on_paragraph and not force:
                para_break = _find_next_paragraph_break(source, fence_spans, start, self._min_chars)
                para_limit = max(1, self._max_chars - len(reopen_prefix))
                if para_break and para_break[0] - start <= para_limit:
                    chunk = f"{reopen_prefix}{source[start : para_break[0]]}"
                    if chunk.strip():
                        chunks.append(chunk)
                    start = _skip_leading_newlines(source, para_break[0] + para_break[1])
                    reopen_fence = None
                    continue
                if remaining_length < self._max_chars:
                    break

            view = source[start:]

            if force and remaining_length <= self._max_chars:
                break_result = self._pick_soft_break(view, fence_spans, min_chars=1, offset=start)
            else:
                break_result = self._pick_break(view, fence_spans, min_chars=1 if force else None, offset=start)

            if break_result.index <= 0:
                if force:
                    chunk_text = f"{reopen_prefix}{source[start:]}"
                    chunks.append(chunk_text)
                    start = len(source)
                    reopen_fence = None
                break

            absolute_idx = start + break_result.index
            raw_chunk = f"{reopen_prefix}{source[start:absolute_idx]}"

            fence_split = break_result.fence_split
            if fence_split:
                suffix = "\n" if raw_chunk.endswith("\n") else ""
                raw_chunk = f"{raw_chunk}{suffix}{fence_split['close_line']}\n"

            if raw_chunk.strip():
                chunks.append(raw_chunk)

            if fence_split:
                start = absolute_idx
                reopen_fence = fence_split["fence"]
            else:
                next_start = (
                    absolute_idx + 1 if absolute_idx < len(source) and source[absolute_idx].isspace() else absolute_idx
                )
                start = _skip_leading_newlines(source, next_start)
                reopen_fence = None

            next_prefix_len = len(reopen_fence.open_line) + 1 if reopen_fence and reopen_fence.open_line else 0
            next_length = next_prefix_len + (len(source) - start)
            if next_length < self._min_chars and not force:
                break
            if next_length < self._max_chars and not force and not self._flush_on_paragraph:
                break

        if reopen_fence and reopen_fence.open_line:
            self._buffer = f"{reopen_fence.open_line}\n{_strip_leading_newlines(source[start:])}"
        else:
            self._buffer = _strip_leading_newlines(source[start:])

        self._reopen_fence = reopen_fence
        return chunks

    def drain_streaming_blocks(self, force: bool) -> list[StreamingBlock]:
        texts = self.drain(force)
        blocks: list[StreamingBlock] = []
        for t in texts:
            block = StreamingBlock(text=t, index=self._block_index)
            self._block_index += 1
            blocks.append(block)
        return blocks

    def feed(self, text_delta: str) -> list[StreamingBlock]:
        self.append(text_delta)
        return self.drain_streaming_blocks(force=False)

    async def feed_async(self, text_delta: str) -> list[StreamingBlock]:
        blocks = self.feed(text_delta)
        if not blocks:
            return []

        now = time.monotonic()
        idle_ms = int((now - self._last_append_ts) * 1000) if self._last_append_ts else 0

        if self._last_append_ts and idle_ms >= self._coalesce_idle_ms:
            flushed = list(self._coalesce_buffer)
            self._coalesce_buffer = []
            self._last_append_ts = now
            return flushed + blocks

        self._coalesce_buffer.extend(blocks)
        self._last_append_ts = now

        coalesced_len = sum(len(b.text) for b in self._coalesce_buffer)
        if coalesced_len >= self._coalesce_min:
            flushed = list(self._coalesce_buffer)
            self._coalesce_buffer = []
            return flushed

        if coalesced_len >= self._coalesce_max:
            flushed = list(self._coalesce_buffer)
            self._coalesce_buffer = []
            return flushed

        return []

    def flush(self) -> list[StreamingBlock]:
        output: list[StreamingBlock] = []

        remaining = list(self._coalesce_buffer)
        self._coalesce_buffer = []

        if self._buffer:
            output.append(StreamingBlock(text=self._buffer, index=self._block_index))
            self._buffer = ""
            self._block_index += 1

        return remaining + output

    def _pick_soft_break(
        self,
        buffer: str,
        fence_spans: list[FenceSpan],
        min_chars: int | None = None,
        offset: int = 0,
    ) -> _BreakResult:
        min_c = max(1, min_chars if min_chars is not None else self._min_chars)
        if len(buffer) < min_c:
            return _BreakResult(index=-1)

        idx = _find_break_in_window(
            buffer,
            fence_spans,
            min_chars=min_c,
            break_preference=self._break_preference,
            reverse=False,
            offset=offset,
        )
        if idx != -1:
            return _BreakResult(index=idx)

        return _BreakResult(index=-1)

    def _pick_break(
        self,
        buffer: str,
        fence_spans: list[FenceSpan],
        min_chars: int | None = None,
        offset: int = 0,
    ) -> _BreakResult:
        min_c = max(1, min_chars if min_chars is not None else self._min_chars)
        max_c = max(min_c, self._max_chars)
        if len(buffer) < min_c:
            return _BreakResult(index=-1)

        window = buffer[: min(max_c, len(buffer))]

        idx = _find_break_in_window(
            window,
            fence_spans,
            min_chars=min_c,
            break_preference=self._break_preference,
            reverse=True,
            offset=offset,
        )
        if idx != -1:
            return _BreakResult(index=idx)

        if self._break_preference == BreakPreference.NEWLINE and len(buffer) < max_c:
            return _BreakResult(index=-1)

        for i in range(len(window) - 1, min_c - 1, -1):
            if window[i].isspace() and is_safe_fence_break(fence_spans, offset + i):
                return _BreakResult(index=i)

        if len(buffer) >= max_c:
            if is_safe_fence_break(fence_spans, offset + max_c):
                return _BreakResult(index=max_c)
            fence = find_fence_span_at(fence_spans, offset + max_c)
            if fence:
                close_line_start = _find_fence_close_line_start(buffer, fence, offset)
                if min_c <= close_line_start < max_c:
                    return _BreakResult(
                        index=close_line_start,
                        fence_split={
                            "close_line": f"{fence.indent}{fence.marker}",
                            "fence": fence,
                        },
                    )
                return _BreakResult(
                    index=max_c,
                    fence_split={
                        "close_line": f"{fence.indent}{fence.marker}",
                        "fence": fence,
                    },
                )
            return _BreakResult(index=max_c)

        return _BreakResult(index=-1)


# ── BlockReplyCoalescer ────────────────────────────────────────────────


class BlockReplyCoalescer:
    def __init__(
        self,
        config: BlockReplyCoalescing,
        on_flush: Callable[[str], None],
        is_stopped: Callable[[], bool] | None = None,
    ):
        cfg = config
        self._min_chars = max(1, cfg.min_chars)
        self._max_chars = max(self._min_chars, cfg.max_chars)
        self._idle_ms = max(0, cfg.idle_ms)
        self._joiner = cfg.joiner
        self._flush_on_enqueue = cfg.flush_on_enqueue
        self._on_flush = on_flush
        self._is_stopped = is_stopped or (lambda: False)

        self._buffer_text: str = ""
        self._idle_timer: asyncio.Task[None] | None = None

    def enqueue(self, text: str) -> None:
        if self._is_stopped():
            return
        if not text:
            return

        if self._flush_on_enqueue:
            if self._buffer_text:
                self._on_flush(self._buffer_text)
            self._buffer_text = text
            self._on_flush(self._buffer_text)
            self._buffer_text = ""
            return

        next_text = f"{self._buffer_text}{self._joiner}{text}" if self._buffer_text else text

        if len(next_text) > self._max_chars:
            if self._buffer_text:
                self._on_flush(self._buffer_text)
            if len(text) >= self._max_chars:
                self._on_flush(text)
                self._buffer_text = ""
            else:
                self._buffer_text = text
                self._schedule_idle_flush()
            return

        self._buffer_text = next_text
        if len(self._buffer_text) >= self._max_chars:
            self.flush()
            return
        self._schedule_idle_flush()

    def flush(self) -> None:
        self._cancel_timer()
        if self._is_stopped() or not self._buffer_text:
            return
        text = self._buffer_text
        self._buffer_text = ""
        self._on_flush(text)

    def has_buffered(self) -> bool:
        return bool(self._buffer_text)

    def stop(self) -> None:
        self._cancel_timer()
        self._buffer_text = ""

    def _schedule_idle_flush(self) -> None:
        if self._idle_ms <= 0:
            return
        self._cancel_timer()

        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            return

        async def _idle_flush():
            try:
                await asyncio.sleep(self._idle_ms / 1000)
                self.flush()
            except Exception:
                logger.exception("BlockReplyCoalescer idle flush failed")

        self._idle_timer = loop.create_task(_idle_flush())
        self._idle_timer.add_done_callback(self._on_timer_done)

    def _cancel_timer(self) -> None:
        if self._idle_timer is not None and not self._idle_timer.done():
            self._idle_timer.cancel()
        self._idle_timer = None

    def _on_timer_done(self, task: asyncio.Task) -> None:
        if self._idle_timer is task:
            self._idle_timer = None
        try:
            task.result()
        except Exception:
            logger.exception("BlockReplyCoalescer timer task failed")