# chunk_page.py: the section-citation chunker (E) as one standalone file.
# Header sections, tiny sections merged, long sections split with LangChain's recursive splitter.
import re
import tiktoken
from langchain_text_splitters import Language, RecursiveCharacterTextSplitter

ENC = tiktoken.get_encoding("cl100k_base")
CAP, OVERLAP, MERGE_BELOW = 400, 60, 80
FENCE = re.compile(r"^\s*(```|~~~)")
HEAD = re.compile(r"^(#{1,3})\s+(.+?)\s*#*\s*$")
LINK = re.compile(r"\[([^\]]*)\]\([^)]*\)")
SPLITTER = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    encoding_name="cl100k_base", chunk_size=CAP, chunk_overlap=OVERLAP, is_separator_regex=True,
    separators=RecursiveCharacterTextSplitter.get_separators_for_language(Language.MARKDOWN))


def tokens(text):
    return len(ENC.encode(text, disallowed_special=()))


def clean_heading(h):
    """Drop markdown links from a heading (FastAPI adds a [¶](...) permalink to every heading)."""
    return LINK.sub(r"\1", h).replace("¶", "").strip()


def chunk_page(md, url, title=""):
    # 1. sections at # to ### headings, ignoring lines inside code fences; keep the heading path
    secs, stack, in_fence, pos, start, path = [], {}, False, 0, 0, ""
    for line in md.splitlines(keepends=True):
        in_fence ^= bool(FENCE.match(line))
        m = None if in_fence else HEAD.match(line.rstrip("\n"))
        if m:
            if pos > start:
                secs.append([path, start, pos])
            lvl = len(m.group(1))
            stack = {k: v for k, v in stack.items() if k < lvl}
            stack[lvl] = clean_heading(m.group(2))
            path, start = " > ".join(stack[k] for k in sorted(stack)), pos
        pos += len(line)
    if pos > start:
        secs.append([path, start, pos])
    # 2. merge sections under MERGE_BELOW tokens: into the previous one, or forward if the previous one is tiny
    merged = []
    for s in secs:
        small = tokens(md[s[1]:s[2]]) < MERGE_BELOW
        prev_small = bool(merged) and tokens(md[merged[-1][1]:merged[-1][2]]) < MERGE_BELOW
        if merged and small and tokens(md[merged[-1][1]:s[2]]) <= CAP:
            merged[-1][2] = s[2]
        elif prev_small:
            merged[-1] = [s[0], merged[-1][1], s[2]]
        else:
            merged.append(s)
    # 3. cap: sections over CAP tokens go through the recursive splitter
    out = []
    for path, s, e in merged:
        body = md[s:e]
        for piece in ([body] if tokens(body) <= CAP else SPLITTER.split_text(body)):
            if piece.strip():
                out.append({"url": url, "section_path": path or title, "body": piece, "tokens": tokens(piece)})
    return out