"""Text chunking utilities for vector KB indexing."""

from __future__ import annotations

import re
from typing import Iterable, List, Sequence, Tuple

SECTION_PATTERN = re.compile(r"^(#{1,6})\s+(.+)$", re.MULTILINE)
TOKEN_PATTERN = re.compile(r"\w+|[^\w\s]", re.UNICODE)


def split_markdown_sections(text: str) -> List[Tuple[str, str]]:
    """Split markdown text into (section_title, section_text)."""
    sections: list[tuple[str, str]] = []
    last_title = "root"
    last_end = 0

    for match in SECTION_PATTERN.finditer(text):
        body = text[last_end:match.start()].strip()
        if body:
            sections.append((last_title, body))
        last_title = match.group(2).strip()
        last_end = match.end()

    tail = text[last_end:].strip()
    if tail:
        sections.append((last_title, tail))

    if not sections and text.strip():
        sections.append(("root", text.strip()))
    return sections


def _is_code_fence(line: str) -> bool:
    return line.strip().startswith("```")


def _is_table_line(line: str) -> bool:
    stripped = line.strip()
    return stripped.startswith("|") and "|" in stripped[1:]


def _iter_markdown_blocks(text: str) -> Iterable[str]:
    lines = text.splitlines()
    i = 0
    while i < len(lines):
        line = lines[i]

        if _is_code_fence(line):
            block = [line]
            i += 1
            while i < len(lines):
                block.append(lines[i])
                if _is_code_fence(lines[i]):
                    i += 1
                    break
                i += 1
            yield "\n".join(block).strip()
            continue

        if _is_table_line(line):
            block = [line]
            i += 1
            while i < len(lines) and _is_table_line(lines[i]):
                block.append(lines[i])
                i += 1
            yield "\n".join(block).strip()
            continue

        if line.strip().startswith("#"):
            yield line.strip()
            i += 1
            continue

        if not line.strip():
            i += 1
            continue

        paragraph = [line]
        i += 1
        while i < len(lines):
            current = lines[i]
            if not current.strip():
                i += 1
                break
            if _is_code_fence(current) or _is_table_line(current) or current.strip().startswith("#"):
                break
            paragraph.append(current)
            i += 1
        yield "\n".join(paragraph).strip()


def _split_long_plain_block(block: str, max_chars: int, overlap_chars: int) -> list[str]:
    if len(block) <= max_chars:
        return [block]
    result: list[str] = []
    start = 0
    while start < len(block):
        end = min(start + max_chars, len(block))
        result.append(block[start:end])
        if end >= len(block):
            break
        start = max(end - overlap_chars, start + 1)
    return result


def chunk_text_by_chars(
    text: str,
    *,
    max_chars: int = 800,
    min_chars: int = 120,
    overlap_chars: int = 80,
) -> List[str]:
    """Chunk markdown text by character size with basic structure protection."""
    blocks = [block for block in _iter_markdown_blocks(text) if block.strip()]
    chunks: list[str] = []
    current = ""

    for block in blocks:
        protected = block.startswith("```") or block.startswith("|")
        pieces = [block] if protected else _split_long_plain_block(block, max_chars, overlap_chars)

        for piece in pieces:
            candidate = f"{current}\n\n{piece}".strip() if current else piece
            if len(candidate) <= max_chars or protected:
                current = candidate
            else:
                if current:
                    chunks.append(current)
                current = piece

            if protected and len(current) > max_chars and current:
                chunks.append(current)
                current = ""

    if current:
        chunks.append(current)

    merged: list[str] = []
    for chunk in chunks:
        if merged and len(chunk) < min_chars and len(merged[-1]) + len(chunk) + 2 <= max_chars:
            merged[-1] = f"{merged[-1]}\n\n{chunk}"
        else:
            merged.append(chunk)
    return merged


def _token_count(text: str) -> int:
    return len(TOKEN_PATTERN.findall(text))


def chunk_text_by_tokens(
    text: str,
    *,
    max_tokens: int = 300,
    min_tokens: int = 60,
    overlap_tokens: int = 30,
) -> List[str]:
    """Chunk markdown text by token count with heading/code/table protection."""
    blocks = [block for block in _iter_markdown_blocks(text) if block.strip()]
    chunks: list[str] = []
    current_blocks: list[str] = []
    current_tokens = 0

    for block in blocks:
        block_tokens = _token_count(block)
        protected = block.startswith("```") or block.startswith("|")
        if block_tokens > max_tokens and not protected:
            if current_blocks:
                chunks.append("\n\n".join(current_blocks))
                current_blocks = []
                current_tokens = 0
            char_chunks = chunk_text_by_chars(block, max_chars=max_tokens * 4, min_chars=1, overlap_chars=overlap_tokens * 4)
            chunks.extend(char_chunks)
            continue

        if current_tokens + block_tokens <= max_tokens or protected:
            current_blocks.append(block)
            current_tokens += block_tokens
            if protected and current_tokens >= max_tokens:
                chunks.append("\n\n".join(current_blocks))
                current_blocks = []
                current_tokens = 0
        else:
            if current_blocks:
                chunks.append("\n\n".join(current_blocks))
            overlap_blocks: list[str] = []
            overlap_count = 0
            for old_block in reversed(current_blocks):
                old_tokens = _token_count(old_block)
                if overlap_count + old_tokens > overlap_tokens:
                    break
                overlap_blocks.insert(0, old_block)
                overlap_count += old_tokens
            current_blocks = overlap_blocks + [block]
            current_tokens = overlap_count + block_tokens

    if current_blocks:
        chunks.append("\n\n".join(current_blocks))

    merged: list[str] = []
    for chunk in chunks:
        token_len = _token_count(chunk)
        if merged and token_len < min_tokens:
            merged[-1] = f"{merged[-1]}\n\n{chunk}"
        else:
            merged.append(chunk)
    return [chunk for chunk in merged if chunk.strip()]


__all__ = [
    "chunk_text_by_chars",
    "chunk_text_by_tokens",
    "split_markdown_sections",
]
