from __future__ import annotations

from typing import Literal

NATURAL_BOUNDARIES = [
    "\n\n",
    "\n",
    "。",
    "！",
    "？",
    "；",
    "：",
    ". ",
    "! ",
    "? ",
]


class WeChatChunker:
    def chunk_text(self, text: str, limit: int = 2048) -> list[str]:
        if len(text) <= limit:
            return [text]

        chunks: list[str] = []
        remaining = text
        while len(remaining) > limit:
            cut = self._find_boundary(remaining, limit)
            if cut <= 0:
                cut = limit
            chunks.append(remaining[:cut])
            remaining = remaining[cut:]
        if remaining:
            chunks.append(remaining)
        return chunks

    def chunk_markdown(self, text: str, limit: int = 2048) -> list[str]:
        if len(text) <= limit:
            return [text]

        chunks: list[str] = []
        remaining = text
        code_block = False

        while len(remaining) > limit:
            if "```" in remaining[:limit]:
                fence_idx = remaining.find("```")
                if fence_idx >= 0:
                    code_block = not code_block
                    if code_block:
                        close_idx = remaining.find("```", fence_idx + 3)
                        block_end = close_idx + 3 if close_idx >= 0 else limit
                        if block_end <= limit:
                            cut = block_end
                        else:
                            cut = limit
                    else:
                        cut = self._find_boundary(remaining, limit)
                    chunks.append(remaining[:cut])
                    remaining = remaining[cut:]
                    continue

            cut = self._find_boundary(remaining, limit)
            if cut <= 0:
                cut = limit
            chunks.append(remaining[:cut])
            remaining = remaining[cut:]

        if remaining:
            chunks.append(remaining)
        return chunks

    def chunk(self, text: str, limit: int, mode: Literal["text", "markdown"] = "text") -> list[str]:
        if mode == "markdown":
            return self.chunk_markdown(text, limit)
        return self.chunk_text(text, limit)

    @staticmethod
    def _find_boundary(text: str, limit: int) -> int:
        best = -1
        for boundary in NATURAL_BOUNDARIES:
            idx = text.rfind(boundary, 0, limit)
            if idx > best:
                best = idx
        return best + 1 if best >= 0 else 0