#!/usr/bin/env python3
"""
Utilities for turning Markdown (e.g. PDF-to-Markdown output) into JSONL chunks for RAG.

Core steps
----------
1. Optionally strip Markdown image links.
2. Split on every header line (`#`, `##`, …) so each header starts a new chunk.
3. Merge "tiny" chunks (default < 500 chars) into the *following* chunk.
4. Write one JSONL object per chunk under the key `"text"`.
"""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import List, Literal, Pattern

import pandas as pd
from markdownify_rs import markdown_utils as rs_md


def coalesce_small_chunks(chunks: list[str], min_size: int = 125) -> list[str]:
    """Merge chunks shorter than *min_size* into the following chunk."""
    return rs_md.coalesce_small_chunks(chunks, min_size=min_size)


# ---------- regexes (still used by Python-local helpers) --------------------

_IMG = re.compile(r"!\[[^\]]*]\([^)]*\)")  # ![alt](url)
_MD_LINK_RE: Pattern[str] = re.compile(
    r"\[((?:[^\[\]]|\[[^\]]*\])*)\]\(((?:[^()]|\([^)]*\))*)\)",  # ① link text   ② URL
    re.IGNORECASE,
)


# ---------- Rust-backed helpers (thin wrappers) -----------------------------


def strip_data_uri_images(md: str) -> str:
    """Remove inline images or links whose URL starts with 'data:image'."""
    return rs_md.strip_data_uri_images(md)


def strip_html_and_contents(text: str) -> str:
    """Strip all HTML tags *and everything inside them*."""
    return rs_md.strip_html_and_contents(text)


def split_on_dividers(markdown: str) -> List[str]:
    """Split *markdown* into chunks separated by divider lines."""
    return rs_md.split_on_dividers(markdown)


def link_percentage(chunk: str) -> float:
    """Return the fraction of characters in *chunk* that belong to Markdown links."""
    return rs_md.link_percentage(chunk)


def strip_links_with_substring(text: str, substring: str) -> str:
    """Remove inline Markdown links whose URL contains *substring*, keeping link text."""
    return rs_md.strip_links_with_substring(text, substring)


def remove_large_tables(text: str, max_cells: int = 400) -> str:
    """Strip Markdown tables whose total cell count exceeds *max_cells*."""
    return rs_md.remove_large_tables(text, max_cells=max_cells)


# ---------- Rust-backed batch helpers ---------------------------------------

strip_links_with_substring_batch = rs_md.strip_links_with_substring_batch
remove_large_tables_batch = rs_md.remove_large_tables_batch
link_percentage_batch = rs_md.link_percentage_batch
filter_by_link_percentage = rs_md.filter_by_link_percentage
strip_html_and_contents_batch = rs_md.strip_html_and_contents_batch
remove_lines_with_substring_batch = rs_md.remove_lines_with_substring_batch
fix_newlines_batch = rs_md.fix_newlines_batch
text_pipeline_batch = rs_md.text_pipeline_batch


SplitHeuristic = Literal[
    "headings",
    "heading1",
    "heading2",
    "heading3",
    "heading4",
    "heading5",
    "heading6",
    "sections",
    "statute_sections",
    "lists",
    "links",
    "bold",
    "italic",
    "brute",
    "numbering",
]


def split_into_chunks(
    md: str,
    how: SplitHeuristic = "sections",
) -> List[str]:
    """Split so that each boundary (determined by *how*) begins a new chunk."""
    return rs_md.split_into_chunks(md, how=how)


# ---------- Python-local helpers (not ported to Rust) -----------------------


def strip_prelude(md: str, pattern: str, padding_lines: int = 0) -> str:
    """
    Remove all lines from the beginning until one matches 'pattern'.
    Add back 'padding_lines' removed.
    """
    lines = md.splitlines()
    start = -1
    for idx, line in enumerate(lines):
        if pattern in line:
            start = idx
            break

    start -= padding_lines

    return "\n".join(lines[start:])


def strip_postlude(md: str, pattern: str, padding_lines: int = 0):
    lines = md.splitlines()
    end = 0
    for idx, line in reversed(list(enumerate(lines))):
        if pattern in line:
            end = idx
            break

    end += padding_lines

    return "\n".join(lines[:end])


def remove_image_links(md: str) -> str:
    """Strip every Markdown image link."""
    return _IMG.sub("", md)


def remove_lines_with_substring(text: str, substring: str) -> str:
    """Remove lines containing *substring*."""
    return rs_md.remove_lines_with_substring(text, substring)


def fix_newlines(text: str) -> str:
    """Collapse runs of 3+ newlines to 2."""
    return rs_md.fix_newlines(text)


def extract_first_link_text(text: str) -> str:
    """Return the text inside the first Markdown link in *text*, or '' if none."""
    m = _MD_LINK_RE.search(text)
    return m.group(1) if m else ""


# ---------- title helpers ---------------------------------------------------


def _make_title(old: str | None, new: str | None):
    if not old and not new:
        return "[UNTITLED SECTION]"
    elif old and not new:
        return old
    elif new and not old:
        return new
    return f"{old} - {new}"


def title_from_first_heading(old: str | None, chunk: str) -> str:
    first = chunk.splitlines()[0].lstrip("# ").strip()
    return _make_title(old, first)


def title_from_first_bold(old: str | None, chunk: str) -> str:
    try:
        first = chunk.split("**")[1]
        return _make_title(old, first.strip())
    except IndexError:
        return title_from_first_sentence(old, chunk)


def title_from_first_italic(old: str | None, chunk: str) -> str:
    try:
        first = chunk.split("*")[1]
        return _make_title(old, first.strip())
    except IndexError:
        return title_from_first_sentence(old, chunk)


def title_from_first_sentence(old: str | None, chunk: str) -> str:
    first = chunk.strip().splitlines()[0].split(". ")[0].strip()
    return _make_title(old, first)


def title_from_first_line(old: str | None, chunk: str) -> str:
    line = chunk.strip().splitlines()[0].strip("# ")
    if len(line) < 100 and len(line) > 10:
        return _make_title(old, line)
    else:
        return title_from_first_sentence(old, chunk)


def title_from_first_link(old: str | None, chunk: str) -> str:
    """
    If the chunk contains a Markdown link, append its link-text
    (the part inside the square brackets) to the title.
    """
    m = _MD_LINK_RE.search(chunk)
    if m:
        text = m.group(1).strip()
        return _make_title(old, text)
    return _make_title(old, None)


def title_from_numbered_section(old: str | None, chunk: str) -> str:
    """Extract title from numbered section like '1.01 Definitions.'"""
    first_line = chunk.strip().splitlines()[0]
    # Take up to the first period or em-dash, or the whole line if short
    for sep in [".—", "—", ". "]:
        if sep in first_line:
            title = first_line.split(sep)[0] + ("." if sep == ". " else "")
            return _make_title(old, title.strip())
    return _make_title(old, first_line[:80].strip())


def title_from_statute_section(old: str | None, chunk: str) -> str:
    first_line = chunk.strip().splitlines()[0].strip()
    if len(first_line) > 200:
        first_line = first_line[:200].rstrip() + "…"
    return _make_title(old, first_line.rstrip("."))


# ---------- cascading split ----------------------------------------

HOW_TO_TITLE_FN = {
    "headings": title_from_first_heading,
    "heading1": title_from_first_heading,
    "heading2": title_from_first_heading,
    "heading3": title_from_first_heading,
    "heading4": title_from_first_heading,
    "heading5": title_from_first_heading,
    "heading6": title_from_first_heading,
    "sections": title_from_first_line,
    "statute_sections": title_from_statute_section,
    "lists": title_from_first_line,
    "links": title_from_first_link,
    "bold": title_from_first_bold,
    "italic": title_from_first_italic,
    "brute": title_from_first_sentence,
    "numbering": title_from_numbered_section,
}


def cascading_split(
    df: pd.DataFrame,
    how: list[SplitHeuristic],
    max_length: int = 8_000,
    coalesce_min: int = 225,
):
    assert "content" in df.columns and "title" in df.columns, "df missing title/content"
    working = df.copy()
    for step in how:
        # batch split+coalesce via Rust
        working["content"] = rs_md.cascading_split_text_batch(
            working["content"].tolist(),
            how=[step],
            max_length=max_length,
            coalesce_min=coalesce_min,
        )
        working["num_subsections"] = working["content"].apply(len)
        working = working.explode("content", ignore_index=True)

        # update titles if a rule was provided and we actually split
        title_fn = HOW_TO_TITLE_FN[step]
        working["title"] = working.apply(
            lambda row: row["title"]
            if row["num_subsections"] == 1
            else title_fn(row["title"], row["content"]),
            axis=1,
        )

    working["content_length"] = working["content"].apply(len)
    return working


# ---------- pipeline & file helpers ----------------------------------------


def process(md: str, *, remove_images: bool = True, min_size: int = 500) -> List[str]:
    """Run the full pipeline and return the final chunks."""
    if remove_images:
        md = remove_image_links(md)
    chunks = split_into_chunks(md)
    return coalesce_small_chunks(chunks, min_size=min_size)


def markdown_to_jsonl(
    src: Path,
    dst: Path,
    *,
    remove_images: bool = True,
    min_size: int = 500,
) -> None:
    """Convert *src* Markdown file to JSONL at *dst*."""
    chunks = process(
        src.read_text(encoding="utf-8"), remove_images=remove_images, min_size=min_size
    )
    with dst.open("w", encoding="utf-8") as fh:
        for c in chunks:
            json.dump({"text": c}, fh, ensure_ascii=False)
            fh.write("\n")


# ---------- CLI ------------------------------------------------------------


def _cli() -> None:
    ap = argparse.ArgumentParser(description="Ingest Markdown and output JSONL chunks.")
    ap.add_argument("input_md", type=Path, help="Source Markdown file")
    ap.add_argument("output_jsonl", type=Path, help="Destination JSONL file")
    ap.add_argument(
        "--keep-images", action="store_true", help="Do NOT strip image links"
    )
    ap.add_argument(
        "--min-size", type=int, default=500, help="Minimum chunk size in characters"
    )
    args = ap.parse_args()

    markdown_to_jsonl(
        args.input_md,
        args.output_jsonl,
        remove_images=not args.keep_images,
        min_size=args.min_size,
    )


if __name__ == "__main__":
    _cli()
