"""Safe helpers for provider text streams.

Some OpenAI-compatible providers place reasoning in the normal ``content``
stream as ``<think>...</think>`` instead of using LangChain's dedicated
reasoning fields.  Those fragments must be separated from report/file content
while callers can still surface them in an explicitly marked live UI trace.
"""

from __future__ import annotations


class InlineThinkTagFilter:
    """Separate inline ``<think>`` blocks while preserving token streaming.

    Tags can be split across arbitrary provider chunks, so regex-per-chunk is
    not safe: it can leak a partial tag or the model's reasoning into the
    visible draft.  This small state machine holds only the suffix that might
    still form a tag and emits all confirmed answer text immediately.
    """

    _OPEN_TAG = "<think>"
    _CLOSE_TAG = "</think>"

    def __init__(self, *, allow_implicit_close: bool = False) -> None:
        self._inside_think = False
        self._pending = ""
        self._channel_decided = not allow_implicit_close

    @staticmethod
    def _possible_tag_prefix_size(content: str, tag: str) -> int:
        """Return the longest suffix which may become the next tag."""
        lowered = content.lower()
        for size in range(min(len(content), len(tag) - 1), 0, -1):
            if lowered[-size:] == tag[:size]:
                return size
        return 0

    def push_parts(self, chunk: str) -> tuple[str, str]:
        """Return ``(visible_text, thinking_text)`` for one provider chunk."""
        if not chunk:
            return "", ""

        remaining = self._pending + chunk
        self._pending = ""
        visible: list[str] = []
        thinking: list[str] = []

        if not self._channel_decided:
            lowered = remaining.lower()
            open_at = lowered.find(self._OPEN_TAG)
            close_at = lowered.find(self._CLOSE_TAG)
            if close_at >= 0 and (open_at < 0 or close_at < open_at):
                # Some templates inject the opening tag outside returned
                # content.  Keep the ambiguous prefix private until its close
                # marker arrives, even when that marker spans chunks.
                thinking.append(remaining[:close_at])
                remaining = remaining[close_at + len(self._CLOSE_TAG) :]
                self._channel_decided = True
            elif open_at >= 0:
                self._channel_decided = True
            else:
                self._pending = remaining
                return "", ""

        while remaining:
            tag = self._CLOSE_TAG if self._inside_think else self._OPEN_TAG
            marker_at = remaining.lower().find(tag)
            if marker_at >= 0:
                (thinking if self._inside_think else visible).append(remaining[:marker_at])
                remaining = remaining[marker_at + len(tag) :]
                self._inside_think = not self._inside_think
                continue

            suffix_size = self._possible_tag_prefix_size(remaining, tag)
            confirmed = remaining[:-suffix_size] if suffix_size else remaining
            (thinking if self._inside_think else visible).append(confirmed)
            self._pending = remaining[-suffix_size:] if suffix_size else ""
            break

        return "".join(visible), "".join(thinking)

    def push(self, chunk: str) -> str:
        """Compatibility helper returning only confirmed visible text."""
        return self.push_parts(chunk)[0]

    def confirm_visible_channel(self) -> str:
        """Stop close-only probing and return any buffered visible prelude.

        A provider that emits a dedicated reasoning field has already told us
        how to classify its trace.  Its ordinary ``content`` channel can then
        stream immediately even when the model supports thinking.
        """
        if self._channel_decided:
            return ""
        self._channel_decided = True
        pending = self._pending
        self._pending = ""
        return pending

    def finish_parts(self) -> tuple[str, str]:
        """Flush the remaining visible or thinking text at end of stream."""
        if not self._channel_decided:
            tail = self._pending
            self._pending = ""
            self._channel_decided = True
            return tail, ""
        if self._inside_think:
            thinking = self._pending
            self._pending = ""
            self._inside_think = False
            return "", thinking
        tail = self._pending
        self._pending = ""
        return tail, ""

    def finish(self) -> str:
        """Compatibility helper returning only the visible stream tail."""
        return self.finish_parts()[0]


def strip_inline_think_tags(content: str) -> str:
    """Return complete text with inline model-thinking blocks removed."""
    text_filter = InlineThinkTagFilter()
    return text_filter.push(content) + text_filter.finish()


__all__ = ["InlineThinkTagFilter", "strip_inline_think_tags"]