# chunkers.py: the five chunkers compared in the post.
# A fixed, B recursive, C header_only, D header_capped (with heading path), E = D without the prefix.
import hashlib, json, re, sys
import tiktoken
from langchain_text_splitters import Language, MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter

PARAMS = {
    "encoding": "cl100k_base",
    "cap_tokens": 400,        # max body tokens per chunk for A, B, D, E
    "overlap_tokens": 60,     # 15% overlap for A, B, and for forced splits in D, E
    "merge_below_tokens": 80, # D, E: sections smaller than this merge into the previous section
    "header_levels": 3,       # split on #, ##, ###
    "prefix_format": "{path joined by ' > '}\\n\\n",  # D only; page title first, then headings
    "embed_token_limit": 8000,  # texts longer than this are truncated before embedding (model limit 8191)
}
ENC = tiktoken.get_encoding(PARAMS["encoding"])
FENCE = re.compile(r"^\s*(```|~~~)")
HEAD = re.compile(r"^(#{1,3})\s+(.+?)\s*#*\s*$")


LINK = re.compile(r"\[([^\]]*)\]\([^)]*\)")


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


def ntok(s):
    return len(ENC.encode(s, disallowed_special=()))


def rec(chunk, md, url, title, path, start, end, body, text):
    return {"url": url, "title": title, "section_path": path, "char_start": start, "char_end": end,
            "body": body, "text": text, "tokens": ntok(text)}


def fixed(page):
    md = page["markdown"]
    toks = ENC.encode(md, disallowed_special=())
    _, offs = ENC.decode_with_offsets(toks)
    cap, step = PARAMS["cap_tokens"], PARAMS["cap_tokens"] - PARAMS["overlap_tokens"]
    out, i = [], 0
    while i < len(toks):
        j = min(i + cap, len(toks))
        s, e = offs[i], (offs[j] if j < len(toks) else len(md))
        body = md[s:e]
        if body.strip():
            out.append(rec(None, md, page["url"], page["title"], "", s, e, body, body))
        if j == len(toks):
            break
        i += step
    return out


def _recursive_splitter():
    return RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        encoding_name=PARAMS["encoding"], chunk_size=PARAMS["cap_tokens"], chunk_overlap=PARAMS["overlap_tokens"],
        separators=RecursiveCharacterTextSplitter.get_separators_for_language(Language.MARKDOWN),
        is_separator_regex=True,  # the markdown separators are regexes; without this flag they are escaped and the heading separators never match
        add_start_index=False)


def _split_text(text, offset):
    """Recursive split of text; yields (start, end, body) in page coordinates."""
    docs = _recursive_splitter().create_documents([text])
    prev_start = -1
    for d in docs:
        # LangChain's start_index is unreliable when text repeats, so place each chunk by searching forward from the previous start
        s = text.find(d.page_content, prev_start + 1)
        if s < 0:
            s = max(prev_start + 1, 0)
        prev_start = s
        yield offset + s, offset + s + len(d.page_content), d.page_content


def recursive(page):
    md = page["markdown"]
    return [rec(None, md, page["url"], page["title"], "", s, e, b, b) for s, e, b in _split_text(md, 0) if b.strip()]


def header_only(page):
    md = page["markdown"]
    sp = MarkdownHeaderTextSplitter([("#", "h1"), ("##", "h2"), ("###", "h3")], strip_headers=False)
    docs = [d for d in sp.split_text(md) if d.page_content.strip()]
    starts, cursor = [], 0
    for d in docs:  # locate each chunk in the page: its first line, searching forward
        first = d.page_content.split("\n", 1)[0].strip()
        p = md.find(first, cursor)
        p = cursor if p < 0 else p
        starts.append(p)
        cursor = p + 1
    out = []
    for k, d in enumerate(docs):
        e = starts[k + 1] if k + 1 < len(docs) else len(md)
        path = " > ".join(clean_heading(d.metadata[h]) for h in ("h1", "h2", "h3") if h in d.metadata)
        out.append(rec(None, md, page["url"], page["title"], path, starts[k], e, d.page_content, d.page_content))
    return out


def sections(md):
    """Split on #..### headings outside fenced code blocks. Returns [(path, start, end)]."""
    secs, stack, in_fence, pos = [], {}, False, 0
    cur_start, cur_path = 0, ""
    for line in md.splitlines(keepends=True):
        if FENCE.match(line):
            in_fence = not in_fence
        m = None if in_fence else HEAD.match(line.rstrip("\n"))
        if m:
            if pos > cur_start:
                secs.append([cur_path, cur_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))
            cur_path = " > ".join(stack[k] for k in sorted(stack))
            cur_start = pos
        pos += len(line)
    if pos > cur_start:
        secs.append([cur_path, cur_start, pos])
    return secs


def header_capped(page, prefix=True):
    md, title = page["markdown"], page["title"]
    secs = sections(md)
    # merge tiny sections into the previous one; a tiny preamble (for example a "Skip to content" link) merges forward
    merged = []
    for s in secs:
        small = ntok(md[s[1]:s[2]]) < PARAMS["merge_below_tokens"]
        prev_small = bool(merged) and ntok(md[merged[-1][1]:merged[-1][2]]) < PARAMS["merge_below_tokens"]
        if merged and small and ntok(md[merged[-1][1]:s[2]]) <= PARAMS["cap_tokens"]:
            merged[-1][2] = s[2]
        elif prev_small:
            merged[-1] = [s[0], merged[-1][1], s[2]]  # absorb into the tiny previous section, keep this heading path
        else:
            merged.append(list(s))
    out = []
    for path, s, e in merged:
        body_all = md[s:e]
        pieces = [(s, e, body_all)] if ntok(body_all) <= PARAMS["cap_tokens"] else list(_split_text(body_all, s))
        label = path or title
        pre = f"{label}\n\n" if prefix else ""
        for ps, pe, b in pieces:
            if b.strip():
                out.append(rec(None, md, page["url"], title, label, ps, pe, b, pre + b))
    return out


CHUNKERS = {
    "A_fixed": fixed,
    "B_recursive": recursive,
    "C_header_only": header_only,
    "D_header_capped_prefix": lambda p: header_capped(p, True),
    "E_header_capped_noprefix": lambda p: header_capped(p, False),
}

if __name__ == "__main__":
    if sys.argv[1:] == ["--freeze"]:
        blob = json.dumps({"params": PARAMS, "langchain_text_splitters": __import__("importlib.metadata", fromlist=["x"]).version("langchain-text-splitters"),
                           "src_sha256": hashlib.sha256(open(__file__, "rb").read()).hexdigest()}, indent=2)
        from datetime import datetime, timezone
        d = {"frozen_at_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"), **json.loads(blob),
             "params_sha256": hashlib.sha256(json.dumps(PARAMS, sort_keys=True).encode()).hexdigest()}
        open("params.json", "w").write(json.dumps(d, indent=2))
        print(json.dumps(d, indent=2))