chunking/markdown_parser.py

from __future__ import annotations

import hashlib
import re
from pathlib import Path
from typing import Any

from src.chunking.markdown_loader import normalize_markdown
from src.chunking.page_detector import detect_page_markers, page_range_for_span
from src.chunking.schemas import MarkdownBlock, ParsedMarkdownDocument
from src.chunking.token_counter import TokenCounter


HEADING_RE = re.compile(r"^(#{1,6})\s+(.+?)\s*$")
PAGE_MARKER_RE = re.compile(r"^\s*<!--\s*page\s*:?\s*(\d+)\s*-->\s*$", re.IGNORECASE)
LIST_RE = re.compile(r"^\s*(?:[-*+]\s+|\d+[.)]\s+)")
TABLE_SEPARATOR_RE = re.compile(r"^\s*\|?\s*:?-{3,}:?\s*(?:\|\s*:?-{3,}:?\s*)+\|?\s*$")


def stable_doc_id(path: str | Path, metadata: dict[str, Any] | None = None) -> str:
    metadata = metadata or {}
    explicit = metadata.get("doc_id")
    if explicit:
        return _slug(str(explicit))
    return _slug(Path(path).stem)


def parse_markdown_file(
    path: str | Path,
    *,
    counter: TokenCounter | None = None,
) -> ParsedMarkdownDocument:
    path = Path(path)
    raw = path.read_text(encoding="utf-8", errors="ignore")
    return parse_markdown(
        raw,
        source_file=str(path),
        processed_markdown_path=str(path),
        counter=counter,
    )


def parse_markdown(
    text: str,
    *,
    source_file: str,
    processed_markdown_path: str,
    counter: TokenCounter | None = None,
) -> ParsedMarkdownDocument:
    counter = counter or TokenCounter()
    metadata, body = _split_frontmatter(text)
    body = normalize_markdown(body)
    content_hash = hashlib.sha256(body.encode("utf-8")).hexdigest()
    doc_id = stable_doc_id(processed_markdown_path, metadata)
    blocks = _parse_blocks(body, counter)
    title = _title_from_metadata_or_blocks(metadata, blocks, Path(source_file).stem)
    return ParsedMarkdownDocument(
        doc_id=doc_id,
        title=title,
        source_file=str(metadata.get("source_file") or source_file),
        processed_markdown_path=processed_markdown_path,
        content=body,
        content_hash=content_hash,
        metadata=metadata,
        blocks=blocks,
        token_count=counter.count(body),
        heading_count=sum(1 for block in blocks if block.block_type == "heading"),
    )


def _split_frontmatter(text: str) -> tuple[dict[str, Any], str]:
    normalized = text.replace("\r\n", "\n").replace("\r", "\n")
    if not normalized.startswith("---\n"):
        return {}, normalized
    end = normalized.find("\n---", 4)
    if end < 0:
        return {}, normalized
    raw_meta = normalized[4:end].strip()
    body_start = normalized.find("\n", end + 1)
    body = normalized[body_start + 1 :] if body_start >= 0 else ""
    return _parse_simple_yaml(raw_meta), body


def _parse_simple_yaml(raw: str) -> dict[str, Any]:
    metadata: dict[str, Any] = {}
    for line in raw.splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#") or ":" not in stripped:
            continue
        key, value = stripped.split(":", 1)
        key = key.strip()
        value = value.strip().strip("'\"")
        if value.lower() in {"true", "false"}:
            metadata[key] = value.lower() == "true"
        elif value.lower() in {"null", "none", "~"}:
            metadata[key] = None
        else:
            metadata[key] = value
    return metadata


def _parse_blocks(text: str, counter: TokenCounter) -> list[MarkdownBlock]:
    markers = detect_page_markers(text)
    lines = text.splitlines(keepends=True)
    offsets: list[int] = []
    cursor = 0
    for line in lines:
        offsets.append(cursor)
        cursor += len(line)

    blocks: list[MarkdownBlock] = []
    heading_stack: list[tuple[int, str]] = []
    i = 0
    while i < len(lines):
        line = lines[i]
        stripped = line.strip()
        if not stripped:
            i += 1
            continue

        start = offsets[i]
        heading_match = HEADING_RE.match(line.rstrip("\n"))
        if heading_match:
            level = len(heading_match.group(1))
            heading_text = heading_match.group(2).strip().rstrip("#").strip()
            heading_stack = [
                item for item in heading_stack if item[0] < level
            ] + [(level, heading_text)]
            end = start + len(line)
            section_path = [item[1] for item in heading_stack]
            blocks.append(
                _make_block(
                    "heading",
                    line,
                    section_path,
                    start,
                    end,
                    markers,
                    counter,
                    heading_level=level,
                    heading_text=heading_text,
                )
            )
            i += 1
            continue

        if PAGE_MARKER_RE.match(stripped):
            end = start + len(line)
            blocks.append(
                _make_block(
                    "page_marker",
                    line,
                    [item[1] for item in heading_stack],
                    start,
                    end,
                    markers,
                    counter,
                )
            )
            i += 1
            continue

        if stripped.startswith("```") or stripped.startswith("~~~"):
            fence = stripped[:3]
            j = i + 1
            while j < len(lines) and not lines[j].strip().startswith(fence):
                j += 1
            if j < len(lines):
                j += 1
            blocks.append(_block_from_range("code", lines, offsets, i, j, heading_stack, markers, counter))
            i = j
            continue

        if _is_table_start(lines, i):
            j = i + 2
            while j < len(lines) and _looks_like_table_row(lines[j]):
                j += 1
            blocks.append(_block_from_range("table", lines, offsets, i, j, heading_stack, markers, counter))
            i = j
            continue

        if LIST_RE.match(line):
            j = i + 1
            while j < len(lines):
                candidate = lines[j]
                if not candidate.strip():
                    if j + 1 < len(lines) and LIST_RE.match(lines[j + 1]):
                        j += 1
                        continue
                    break
                if _is_special_start(candidate) and not candidate.startswith((" ", "\t")):
                    break
                if LIST_RE.match(candidate) or candidate.startswith((" ", "\t")):
                    j += 1
                    continue
                break
            blocks.append(_block_from_range("list", lines, offsets, i, j, heading_stack, markers, counter))
            i = j
            continue

        j = i + 1
        while j < len(lines):
            candidate = lines[j]
            if not candidate.strip() or _is_special_start(candidate):
                break
            j += 1
        blocks.append(_block_from_range("paragraph", lines, offsets, i, j, heading_stack, markers, counter))
        i = j

    return blocks


def _block_from_range(
    block_type: str,
    lines: list[str],
    offsets: list[int],
    start_index: int,
    end_index: int,
    heading_stack: list[tuple[int, str]],
    markers,
    counter: TokenCounter,
) -> MarkdownBlock:
    start = offsets[start_index]
    end = offsets[end_index - 1] + len(lines[end_index - 1])
    content = "".join(lines[start_index:end_index])
    return _make_block(
        block_type,
        content,
        [item[1] for item in heading_stack],
        start,
        end,
        markers,
        counter,
    )


def _make_block(
    block_type: str,
    content: str,
    section_path: list[str],
    char_start: int,
    char_end: int,
    markers,
    counter: TokenCounter,
    *,
    heading_level: int | None = None,
    heading_text: str | None = None,
) -> MarkdownBlock:
    page_start, page_end = page_range_for_span(markers, char_start, char_end)
    return MarkdownBlock(
        block_type=block_type,
        content=content.strip(),
        section_path=list(section_path),
        char_start=char_start,
        char_end=char_end,
        token_count=counter.count(content),
        page_start=page_start,
        page_end=page_end,
        heading_level=heading_level,
        heading_text=heading_text,
    )


def _is_special_start(line: str) -> bool:
    stripped = line.strip()
    return (
        bool(HEADING_RE.match(line.rstrip("\n")))
        or bool(PAGE_MARKER_RE.match(stripped))
        or stripped.startswith("```")
        or stripped.startswith("~~~")
        or LIST_RE.match(line) is not None
        or _looks_like_table_row(line)
    )


def _is_table_start(lines: list[str], index: int) -> bool:
    if index + 1 >= len(lines):
        return False
    return _looks_like_table_row(lines[index]) and TABLE_SEPARATOR_RE.match(lines[index + 1].strip()) is not None


def _looks_like_table_row(line: str) -> bool:
    stripped = line.strip()
    return stripped.startswith("|") and stripped.endswith("|") and stripped.count("|") >= 2


def _title_from_metadata_or_blocks(
    metadata: dict[str, Any],
    blocks: list[MarkdownBlock],
    fallback: str,
) -> str:
    for key in ("title", "document_title", "name"):
        if metadata.get(key):
            return str(metadata[key])
    for block in blocks:
        if block.block_type == "heading" and block.heading_text:
            return block.heading_text
    return fallback


def _slug(value: str) -> str:
    slug = re.sub(r"[^\w]+", "_", value.lower(), flags=re.UNICODE).strip("_")
    return slug or hashlib.sha256(value.encode("utf-8")).hexdigest()[:12]