#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11"
# dependencies = ["markdown-it-py", "pyyaml"]
# ///
"""Structure-aware markdown chunker for the wiki skill.

Splits a document along its H1-H3 heading hierarchy, then merges small siblings and
splits oversized sections along block boundaries. Every chunk carries a breadcrumb
(`path > title > section > subsection`) which goes into the embedded text as well as
the metadata, so a vector represents a passage in context rather than in isolation.

markdown-it-py supplies the AST (it knows the CommonMark edge cases); the chunking
policy below is ours. Block boundaries come from token line maps, so the emitted text
is the original markdown — code fences and tables stay byte-for-byte intact.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field

import yaml
from markdown_it import MarkdownIt  # ty: ignore[unresolved-import]

CHUNKER_VERSION = "1"

MERGE_BELOW = 200
SPLIT_ABOVE = 800
OVERLAP = 64
CHARS_PER_TOKEN = 4

SECTION_LEVELS = (1, 2, 3)
BREADCRUMB_SEP = " > "

_FRONTMATTER_RE = re.compile(r"\A---[ \t]*\r?\n(.*?)\r?\n---[ \t]*(?:\r?\n|\Z)", re.DOTALL)

_md = MarkdownIt("commonmark")


@dataclass(frozen=True)
class Chunk:
    breadcrumb: str
    text: str


@dataclass(frozen=True)
class ParsedFile:
    title: str | None = None
    tags: list[str] = field(default_factory=list)
    headings: list[str] = field(default_factory=list)
    chunks: list[Chunk] = field(default_factory=list)


@dataclass(frozen=True)
class _Heading:
    start: int
    end: int
    level: int
    title: str


@dataclass(frozen=True)
class _Section:
    breadcrumb: tuple[str, ...]
    text: str

    @property
    def parent(self) -> tuple[str, ...]:
        return self.breadcrumb[:-1]


def token_estimate(text: str) -> int:
    """Approximate token count. Sizing does not need a real tokenizer (plan: 4 chars/token)."""
    return max(1, len(text) // CHARS_PER_TOKEN)


def parse_markdown(text: str, path: str) -> ParsedFile:
    """Chunk one markdown document. `path` is the source-relative path, the breadcrumb root."""
    frontmatter, body = _split_frontmatter(text)
    frontmatter_title = _frontmatter_title(frontmatter)
    tags = _frontmatter_tags(frontmatter)

    root: tuple[str, ...] = (path,) if frontmatter_title is None else (path, frontmatter_title)
    sections, headings = _split_sections(body, root)
    merged = _merge_small_siblings(sections)
    chunks = [chunk for section in merged for chunk in _split_oversized(section)]

    # The catalog title falls back to the first heading, because notes carry their title as
    # `# H1` far more often than as frontmatter. It deliberately does NOT feed the breadcrumb
    # root: the H1 already reaches the breadcrumb through the heading stack, and changing the
    # root would rewrite every chunk's text and force a full reindex.
    title = frontmatter_title or (headings[0] if headings else None)
    return ParsedFile(title=title, tags=tags, headings=headings, chunks=chunks)


# ---------------------------------------------------------------------------
# frontmatter
# ---------------------------------------------------------------------------


def _split_frontmatter(text: str) -> tuple[dict, str]:
    """Peel off a leading YAML frontmatter block. Malformed frontmatter stays body text."""
    match = _FRONTMATTER_RE.match(text)
    if not match:
        return {}, text
    try:
        data = yaml.safe_load(match.group(1))
    except yaml.YAMLError:
        return {}, text
    if not isinstance(data, dict):
        return {}, text
    return data, text[match.end() :]


def _frontmatter_title(frontmatter: dict) -> str | None:
    title = frontmatter.get("title")
    if isinstance(title, str) and title.strip():
        return title.strip()
    return None


def _frontmatter_tags(frontmatter: dict) -> list[str]:
    raw = frontmatter.get("tags")
    if isinstance(raw, str):
        values = raw.split(",")
    elif isinstance(raw, list):
        values = [str(item) for item in raw]
    else:
        return []
    return [tag.strip() for tag in values if tag.strip()]


# ---------------------------------------------------------------------------
# sectioning
# ---------------------------------------------------------------------------


def _split_sections(body: str, root: tuple[str, ...]) -> tuple[list[_Section], list[str]]:
    """Cut the body at H1-H3 headings. H4+ stay inside their parent section."""
    lines = body.split("\n")
    tokens = _md.parse(body)

    starts: list[_Heading] = []
    for index, token in enumerate(tokens):
        if token.type != "heading_open" or token.map is None:
            continue
        level = int(token.tag[1:])
        if level not in SECTION_LEVELS:
            continue
        inline = tokens[index + 1] if index + 1 < len(tokens) else None
        title = inline.content.strip() if inline is not None else ""
        starts.append(_Heading(start=token.map[0], end=token.map[1], level=level, title=title))

    headings = [heading.title for heading in starts]
    boundaries = [heading.start for heading in starts] + [len(lines)]

    sections: list[_Section] = []
    preamble = "\n".join(lines[: boundaries[0]]).strip()
    if preamble:
        sections.append(_Section(root, preamble))

    stack: list[tuple[int, str]] = []
    for position, heading in enumerate(starts):
        while stack and stack[-1][0] >= heading.level:
            stack.pop()
        stack.append((heading.level, heading.title))
        section_end = boundaries[position + 1]
        # A heading with no prose of its own would embed as a bare title; the heading
        # still reaches the index through its children's breadcrumbs, so drop it.
        if not "\n".join(lines[heading.end : section_end]).strip():
            continue
        text = "\n".join(lines[heading.start : section_end]).strip()
        sections.append(_Section(_dedupe(root + tuple(title for _, title in stack)), text))
    return sections, headings


def _dedupe(parts: tuple[str, ...]) -> tuple[str, ...]:
    """Drop consecutive repeats so a frontmatter title matching the H1 shows up once."""
    out: list[str] = []
    for part in parts:
        if not out or out[-1] != part:
            out.append(part)
    return tuple(out)


def _merge_small_siblings(sections: list[_Section]) -> list[_Section]:
    """Glue a too-small section onto the following sibling under the same parent."""
    merged: list[_Section] = []
    for section in sections:
        if not merged:
            merged.append(section)
            continue
        previous = merged[-1]
        combined = f"{previous.text}\n\n{section.text}"
        if (
            token_estimate(previous.text) < MERGE_BELOW
            and previous.parent == section.parent
            and token_estimate(combined) <= SPLIT_ABOVE
        ):
            merged[-1] = _Section(previous.breadcrumb, combined)
        else:
            merged.append(section)
    return merged


# ---------------------------------------------------------------------------
# splitting
# ---------------------------------------------------------------------------


def _split_oversized(section: _Section) -> list[Chunk]:
    """Break a section over SPLIT_ABOVE into block-aligned pieces with OVERLAP carry-over."""
    breadcrumb = BREADCRUMB_SEP.join(section.breadcrumb)
    if token_estimate(section.text) <= SPLIT_ABOVE:
        return [Chunk(breadcrumb, _embed_text(breadcrumb, section.text))]

    blocks = _top_level_blocks(section.text)
    chunks: list[Chunk] = []
    current: list[str] = []
    for block in blocks:
        candidate = [*current, block]
        if current and token_estimate("\n\n".join(candidate)) > SPLIT_ABOVE:
            chunks.append(Chunk(breadcrumb, _embed_text(breadcrumb, "\n\n".join(current))))
            current = [*_overlap_tail(current), block]
        else:
            current = candidate
    if current:
        chunks.append(Chunk(breadcrumb, _embed_text(breadcrumb, "\n\n".join(current))))
    return chunks


def _top_level_blocks(text: str) -> list[str]:
    """Top-level markdown blocks, taken from token line maps so fences stay whole."""
    lines = text.split("\n")
    tokens = _md.parse(text)
    starts = sorted({token.map[0] for token in tokens if token.level == 0 and token.map})
    if not starts:
        return [text]
    bounds = [*starts, len(lines)]
    blocks = ["\n".join(lines[bounds[i] : bounds[i + 1]]).strip() for i in range(len(starts))]
    return [block for block in blocks if block]


def _overlap_tail(blocks: list[str]) -> list[str]:
    """Trailing whole blocks of the emitted chunk, up to OVERLAP tokens."""
    tail: list[str] = []
    budget = OVERLAP
    for block in reversed(blocks):
        cost = token_estimate(block)
        if cost > budget:
            break
        tail.insert(0, block)
        budget -= cost
    return tail


def _embed_text(breadcrumb: str, text: str) -> str:
    """The stored chunk text carries its breadcrumb, so the vector sees the context."""
    return f"{breadcrumb}\n\n{text}"