"""Extract highlights, typed notes, and ink strokes from reMarkable document bundles.

Uses rmscene to parse v6 .rm files for highlighted text (GlyphRange items),
typed text (Text items), and handwritten ink (Line items).
Uses PyMuPDF to render annotations onto the original PDF.
"""

import io
import json
import logging
import re
import zipfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from rmscene import read_tree, scene_items as si

log = logging.getLogger(__name__)

# Suppress rmscene "newer format" warning — benign, data is still extracted
logging.getLogger("rmscene.tagged_block_reader").setLevel(logging.ERROR)

# GlyphRange items on adjacent lines within this y-gap are merged
_MAX_LINE_GAP = 100.0

# Highlight annotation style
_HIGHLIGHT_COLOR = (1.0, 0.92, 0.3)  # soft yellow
_HIGHLIGHT_OPACITY = 0.35
_HIGHLIGHT_TRIM = 0.20  # shrink quads vertically by this fraction per side


def extract_original_pdf(zip_path: Path) -> Optional[bytes]:
    """Extract the original PDF from a reMarkable bundle with ink strokes.

    Returns PDF bytes with handwritten ink rendered (but no highlight
    annotations — those are added by Zotero's own annotation layer).
    Falls back to the raw PDF if ink rendering fails.
    """
    try:
        import pymupdf
    except ImportError:
        pass  # fall through to raw extraction

    try:
        with zipfile.ZipFile(zip_path, "r") as zf:
            pdf_names = [n for n in zf.namelist() if n.endswith(".pdf")]
            if not pdf_names:
                return None
            pdf_data = zf.read(pdf_names[0])

        # Render ink strokes onto the original PDF
        try:
            ink_by_page = extract_ink_strokes(zip_path)
            if ink_by_page:
                doc = pymupdf.open(stream=pdf_data, filetype="pdf")
                _render_ink_on_pdf(doc, ink_by_page)
                buf = io.BytesIO()
                doc.save(buf, garbage=3, deflate=True)
                doc.close()
                return buf.getvalue()
        except Exception:
            log.debug("Ink rendering on original PDF failed, using raw", exc_info=True)

        return pdf_data
    except Exception:
        log.warning("Failed to extract original PDF from %s", zip_path.name, exc_info=True)
        return None


def get_page_count(zip_path: Path) -> int:
    """Return the total page count from a reMarkable document bundle.

    Parses the .content file to count pages. Returns 0 on failure.
    """
    try:
        with zipfile.ZipFile(zip_path, "r") as zf:
            content_files = [n for n in zf.namelist() if n.endswith(".content")]
            if not content_files:
                return 0
            content_data = json.loads(zf.read(content_files[0]))
            page_ids = content_data.get("cPages", {}).get("pages", [])
            if not page_ids:
                page_ids = content_data.get("pages", [])
            return len(page_ids)
    except Exception:
        log.debug("Could not get page count from %s", zip_path, exc_info=True)
        return 0


def extract_highlights(zip_path: Path) -> Dict[int, List[str]]:
    """Extract highlighted text strings from a reMarkable document bundle.

    Returns a dict mapping page numbers (1-based) to lists of merged
    highlight passages for that page. Empty dict if no highlights found.
    """
    by_page: Dict[int, List[str]] = {}

    try:
        with zipfile.ZipFile(zip_path, "r") as zf:
            content_files = [n for n in zf.namelist() if n.endswith(".content")]
            if not content_files:
                return by_page

            content_data = json.loads(zf.read(content_files[0]))
            page_ids = content_data.get("cPages", {}).get("pages", [])
            if not page_ids:
                page_ids = content_data.get("pages", [])

            ordered_ids = []
            for page in page_ids:
                if isinstance(page, dict):
                    ordered_ids.append(page.get("id", ""))
                else:
                    ordered_ids.append(str(page))

            rm_files = [n for n in zf.namelist() if n.endswith(".rm")]

            # Build (page_index, rm_file) pairs and sort by page order
            indexed = []
            for rm_file in rm_files:
                stem = Path(rm_file).stem
                try:
                    idx = ordered_ids.index(stem)
                except ValueError:
                    continue
                indexed.append((idx, rm_file))
            indexed.sort()

            for page_idx, rm_file in indexed:
                try:
                    rm_data = zf.read(rm_file)
                    raw = _extract_raw_glyphs(rm_data)
                    merged = _merge_glyphs(raw)
                    if merged:
                        by_page[page_idx + 1] = merged  # 1-based page numbers
                except Exception:
                    log.warning("Failed to parse %s in %s", rm_file, zip_path, exc_info=True)

    except Exception:
        log.warning("Failed to read zip %s", zip_path, exc_info=True)

    by_page = _merge_cross_page(by_page)

    total = sum(len(v) for v in by_page.values())
    log.info("Extracted %d highlight(s) from %s", total, zip_path.name)
    return by_page


def _merge_cross_page(by_page: Dict[int, List[str]]) -> Dict[int, List[str]]:
    """Merge highlight passages that span page breaks.

    If the last passage on page N ends without terminal punctuation and the
    first passage on page N+1 starts lowercase, they are likely the same
    sentence split across pages. Merge them into page N.
    """
    if len(by_page) < 2:
        return by_page

    pages = sorted(by_page.keys())
    for i in range(len(pages) - 1):
        page_n = pages[i]
        page_next = pages[i + 1]

        if page_next != page_n + 1:
            continue

        # Pages may have been deleted by a prior merge in this loop
        if page_n not in by_page or page_next not in by_page:
            continue

        passages_n = by_page[page_n]
        passages_next = by_page[page_next]

        if not passages_n or not passages_next:
            continue

        last = passages_n[-1]
        first = passages_next[0]

        ends_mid = not last.rstrip().endswith((".", "!", "?", ":", '"'))
        starts_lower = first[:1].islower()

        if ends_mid and starts_lower:
            passages_n[-1] = last.rstrip() + " " + first.lstrip()
            passages_next.pop(0)
            if not passages_next:
                del by_page[page_next]

    return by_page


def _recover_pdf_text(page_text: str, search_text: str) -> Optional[str]:
    """Find search_text in page_text via whitespace/hyphen-normalized matching.

    reMarkable OCR often strips spaces and hyphens at line breaks, producing
    concatenated text like ``proofs,standard`` or ``math`` (from ``math-\\n
    ematics``).  This function normalises both texts, locates the match, then
    returns the original PDF substring with line-break hyphens rejoined and
    remaining newlines collapsed to spaces so that PyMuPDF ``search_for`` can
    find it.
    """
    norm_chars: List[str] = []
    norm_to_orig: List[int] = []
    for i, ch in enumerate(page_text):
        if not ch.isspace() and ch not in ("-", "\u00ad"):
            norm_to_orig.append(i)
            norm_chars.append(ch.lower())
    page_norm = "".join(norm_chars)

    search_norm = re.sub(r"[\s\-\u00ad]+", "", search_text.lower())
    if not search_norm:
        return None

    pos = page_norm.find(search_norm)
    if pos < 0:
        return None

    orig_start = norm_to_orig[pos]
    orig_end = norm_to_orig[pos + len(search_norm) - 1] + 1
    raw = page_text[orig_start:orig_end]
    # Rejoin hyphenated words, collapse remaining newlines to spaces
    return raw.replace("-\n", "").replace("\n", " ")


def _search_highlight_positions(
    doc: Any,
    highlights_by_page: Dict[int, List[Tuple[str, Optional[float]]]],
) -> List[Dict[str, Any]]:
    """Search for highlight text in a PyMuPDF document and return positions.

    Shared helper used by both render_annotated_pdf() and
    extract_zotero_highlights().  Returns a list of dicts with keys:
    text, page_index (0-based), quads (list of pymupdf.Quad), page_height.
    """
    import pymupdf

    _RM_HEIGHT = 1872.0
    _RM_TO_PDF_SCALE = 0.70

    results: List[Dict[str, Any]] = []

    for page_idx, glyph_list in highlights_by_page.items():
        if page_idx >= len(doc):
            continue
        page = doc[page_idx]
        page_h = page.rect.height

        highlighted: List = []  # list of pymupdf.Rect

        for text, rm_y in glyph_list:
            quads = page.search_for(text, quads=True)
            if not quads:
                # Fallback: recover actual PDF text via normalized matching
                recovered = _recover_pdf_text(page.get_text("text"), text)
                if recovered:
                    quads = page.search_for(recovered, quads=True)
            if not quads:
                continue

            groups = _group_quads(quads, page_h)

            if len(groups) > 1 and rm_y is not None:
                expected_frac = (rm_y / _RM_HEIGHT) * _RM_TO_PDF_SCALE
                groups = [min(
                    groups,
                    key=lambda g: abs(g[0].ul.y / page_h - expected_frac),
                )]

            selected = groups[0]

            new_quads = []
            for q in selected:
                r = q.rect
                center = pymupdf.Point(
                    (r.x0 + r.x1) / 2, (r.y0 + r.y1) / 2,
                )
                if any(hr.contains(center) for hr in highlighted):
                    continue
                new_quads.append(q)
                highlighted.append(r)

            if new_quads:
                results.append({
                    "text": text,
                    "page_index": page_idx,
                    "quads": new_quads,
                    "page_height": page_h,
                })

    return results


def render_annotated_pdf(zip_path: Path, output_path: Path) -> bool:
    """Render highlight annotations onto the original PDF using PyMuPDF.

    Extracts the original PDF and highlight text from the zip bundle,
    searches for each text line on the correct page, uses the RM
    y-coordinate to pick the right match when there are duplicates,
    and deduplicates overlapping regions.

    Returns True on success, False on failure.
    """
    try:
        import pymupdf
    except ImportError:
        log.warning("pymupdf not installed, cannot render annotated PDF")
        return False

    try:
        highlights_by_page = _extract_highlights_by_page(zip_path)
        if not highlights_by_page:
            log.info("No highlights to render for %s", zip_path.name)

        with zipfile.ZipFile(zip_path, "r") as zf:
            pdf_names = [n for n in zf.namelist() if n.endswith(".pdf")]
            if not pdf_names:
                log.warning("No PDF found in %s", zip_path)
                return False
            pdf_data = zf.read(pdf_names[0])

        doc = pymupdf.open(stream=pdf_data, filetype="pdf")
        positions = _search_highlight_positions(doc, highlights_by_page)

        for pos in positions:
            page = doc[pos["page_index"]]
            slimmed = [_slim_quad(pymupdf, q, _HIGHLIGHT_TRIM)
                       for q in pos["quads"]]
            annot = page.add_highlight_annot(slimmed)
            annot.set_colors(stroke=_HIGHLIGHT_COLOR)
            annot.set_opacity(_HIGHLIGHT_OPACITY)
            annot.update()

        # Render handwritten ink strokes
        ink_count = 0
        try:
            ink_by_page = extract_ink_strokes(zip_path)
            if ink_by_page:
                ink_count = _render_ink_on_pdf(doc, ink_by_page)
        except Exception:
            log.debug("Ink stroke rendering failed, continuing", exc_info=True)

        doc.save(str(output_path), garbage=3, deflate=True)
        doc.close()
        log.info(
            "Rendered annotated PDF with %d highlight(s) and %d ink stroke(s): %s",
            len(positions), ink_count, output_path,
        )
        return True

    except Exception:
        log.warning("Failed to render annotated PDF for %s", zip_path.name, exc_info=True)
        return False


def extract_zotero_highlights(
    zip_path: Path,
    pdf_bytes: Optional[bytes] = None,
) -> List[Dict[str, Any]]:
    """Extract highlight positions suitable for Zotero annotation API.

    Returns list of dicts with keys: text, page_index (0-based),
    page_label (1-based str), rects (list of [x0, y0, x1, y1] in PDF
    bottom-left coordinates), sort_index, color.
    """
    try:
        import pymupdf
    except ImportError:
        log.warning("pymupdf not installed, cannot extract Zotero highlights")
        return []

    try:
        highlights_by_page = _extract_highlights_by_page(zip_path)
        if not highlights_by_page:
            return []

        if pdf_bytes is None:
            with zipfile.ZipFile(zip_path, "r") as zf:
                pdf_names = [n for n in zf.namelist() if n.endswith(".pdf")]
                if not pdf_names:
                    return []
                pdf_bytes = zf.read(pdf_names[0])

        doc = pymupdf.open(stream=pdf_bytes, filetype="pdf")
        positions = _search_highlight_positions(doc, highlights_by_page)

        results: List[Dict[str, Any]] = []
        for pos in positions:
            page_idx = pos["page_index"]
            page_h = pos["page_height"]
            page = doc[page_idx]
            page_text = page.get_text("text")

            # Convert quads from PyMuPDF (top-left origin) to PDF (bottom-left)
            rects = []
            for q in pos["quads"]:
                r = q.rect
                rects.append([
                    round(r.x0, 3),
                    round(page_h - r.y1, 3),
                    round(r.x1, 3),
                    round(page_h - r.y0, 3),
                ])

            # Compute sort_index: PPPPP|CCCCCC|TTTTT
            char_offset = page_text.find(pos["text"])
            top_y = min(q.rect.y0 for q in pos["quads"])
            sort_index = (
                f"{page_idx:05d}"
                f"|{max(char_offset, 0):06d}"
                f"|{int(page_h - top_y):05d}"
            )

            results.append({
                "text": pos["text"],
                "page_index": page_idx,
                "page_label": str(page_idx + 1),
                "rects": rects,
                "sort_index": sort_index,
                "color": "#ffd400",
            })

        doc.close()
        log.info("Extracted %d Zotero highlight position(s)", len(results))
        return results

    except Exception:
        log.warning("Failed to extract Zotero highlights", exc_info=True)
        return []


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _group_quads(quads, page_h: float):
    """Group consecutive quads that belong to the same match.

    PyMuPDF returns one quad per text-run/line for a single match.
    Consecutive quads whose vertical distance is small (< 3% of page
    height) are treated as part of the same match.
    """
    if not quads:
        return []
    threshold = page_h * 0.03
    groups: list = [[quads[0]]]
    for q in quads[1:]:
        prev = groups[-1][-1]
        if abs(q.ul.y - prev.ul.y) < threshold or q.ul.y - prev.lr.y < threshold:
            groups[-1].append(q)
        else:
            groups.append([q])
    return groups


def _slim_quad(pymupdf, q, trim_frac: float):
    """Shrink a quad vertically by *trim_frac* on each side."""
    h = q.ll.y - q.ul.y
    trim = h * trim_frac
    return pymupdf.Quad(
        pymupdf.Point(q.ul.x, q.ul.y + trim),
        pymupdf.Point(q.ur.x, q.ur.y + trim),
        pymupdf.Point(q.ll.x, q.ll.y - trim),
        pymupdf.Point(q.lr.x, q.lr.y - trim),
    )


def _extract_highlights_by_page(
    zip_path: Path,
) -> Dict[int, List[Tuple[str, float | None]]]:
    """Return {page_index: [(line_text, rm_y), ...]} from the zip bundle.

    Each entry is a single GlyphRange line with its RM y-coordinate
    for position-aware matching in the PDF.
    """
    result: Dict[int, List[Tuple[str, float | None]]] = {}

    try:
        with zipfile.ZipFile(zip_path, "r") as zf:
            content_files = [n for n in zf.namelist() if n.endswith(".content")]
            if not content_files:
                return result

            content_data = json.loads(zf.read(content_files[0]))
            page_ids = content_data.get("cPages", {}).get("pages", [])
            if not page_ids:
                page_ids = content_data.get("pages", [])

            ordered_ids = []
            for page in page_ids:
                if isinstance(page, dict):
                    ordered_ids.append(page.get("id", ""))
                else:
                    ordered_ids.append(str(page))

            rm_files = [n for n in zf.namelist() if n.endswith(".rm")]

            for rm_file in rm_files:
                stem = Path(rm_file).stem
                try:
                    page_idx = ordered_ids.index(stem)
                except ValueError:
                    continue

                try:
                    rm_data = zf.read(rm_file)
                    raw_glyphs = _extract_raw_glyphs(rm_data)
                    if raw_glyphs:
                        result[page_idx] = [
                            (text, y) for text, y, _color in raw_glyphs
                        ]
                except Exception:
                    log.warning(
                        "Failed to parse %s in %s",
                        rm_file, zip_path, exc_info=True,
                    )

    except Exception:
        log.warning("Failed to read zip %s", zip_path, exc_info=True)

    return result


def _join_dedup(parts: List[str]) -> str:
    """Join text parts, removing overlapping words at boundaries.

    When two GlyphRange items overlap (e.g. one ends with 'species' and the
    next starts with 'species or'), the duplicate words are removed.
    """
    if not parts:
        return ""
    result = parts[0]
    for part in parts[1:]:
        result_words = result.split()
        part_words = part.split()
        # Check for 1-3 word overlap at boundary
        overlap = 0
        for n in range(min(3, len(result_words), len(part_words)), 0, -1):
            if result_words[-n:] == part_words[:n]:
                overlap = n
                break
        if overlap:
            result += " " + " ".join(part_words[overlap:])
        else:
            result += " " + part
    return result


def _clean_highlight_text(text: str) -> str:
    """Clean up OCR artifacts from reMarkable highlight text."""
    # Strip superscript citation markers: (p1), (p2), (1), (23), etc.
    text = re.sub(r'\(p?\d+\)', '', text)

    # Clean punctuation artifacts from citation removal
    text = re.sub(r',\s*,', ',', text)       # ",," → ","
    text = re.sub(r';\s*;', ';', text)       # ";;" → ";"
    text = re.sub(r',(\s*and\b)', r'\1', text)  # ", and" after removed citation

    # Insert space after sentence-ending punctuation followed by a letter
    text = re.sub(r'([.;!?])([A-Za-z])', r'\1 \2', text)

    # Insert space after comma followed directly by a letter
    text = re.sub(r',([A-Za-z])', r', \1', text)

    # Fix broken line-wrap joins: lowercase followed by uppercase mid-sentence
    # e.g. "operationsWe" → "operations We" (but preserve acronyms like "GenAI")
    text = re.sub(r'([a-z]{2})([A-Z][a-z])', r'\1 \2', text)

    # Collapse multiple spaces
    text = re.sub(r' {2,}', ' ', text)

    return text.strip()


def _extract_raw_glyphs(
    rm_data: bytes,
) -> List[Tuple[str, float | None, int]]:
    """Parse a .rm file and return [(text, y, color), ...] per GlyphRange."""
    raw_glyphs: List[Tuple[str, float | None, int]] = []
    try:
        tree = read_tree(io.BytesIO(rm_data))
        for item in tree.walk():
            if isinstance(item, si.GlyphRange) and item.text:
                text = item.text.strip()
                if not text:
                    continue
                y = item.rectangles[0].y if item.rectangles else None
                raw_glyphs.append((text, y, item.color))
    except Exception:
        log.debug("Could not parse .rm data as v6 scene tree", exc_info=True)
    return raw_glyphs


def _merge_glyphs(
    raw_glyphs: List[Tuple[str, float | None, int]],
) -> List[str]:
    """Merge consecutive GlyphRange items on adjacent lines into passages.

    Sorts glyphs by y-coordinate first so that items from different columns
    (which have distinct y-ranges) are grouped correctly instead of
    being interleaved.
    """
    if not raw_glyphs:
        return []

    # Sort by y so items from the same column are adjacent
    sorted_glyphs = sorted(raw_glyphs, key=lambda g: g[1] if g[1] is not None else 0)

    passages: List[str] = []
    current_parts: List[str] = [sorted_glyphs[0][0]]
    prev_y = sorted_glyphs[0][1]
    prev_color = sorted_glyphs[0][2]

    for text, y, color in sorted_glyphs[1:]:
        same_passage = (
            color == prev_color
            and prev_y is not None
            and y is not None
            and abs(y - prev_y) < _MAX_LINE_GAP
        )
        if same_passage:
            current_parts.append(text)
        else:
            passages.append(_clean_highlight_text(_join_dedup(current_parts)))
            current_parts = [text]
        prev_y = y
        prev_color = color

    passages.append(_clean_highlight_text(_join_dedup(current_parts)))
    return passages


# ---------------------------------------------------------------------------
# Page-ID helpers (shared by highlight, ink, and typed-note extraction)
# ---------------------------------------------------------------------------


def _load_rm_pages(zip_path: Path) -> List[Tuple[int, bytes]]:
    """Load .rm page data from a reMarkable bundle.

    Returns a list of (page_index, rm_bytes) tuples sorted by page order.
    """
    result: List[Tuple[int, bytes]] = []
    try:
        with zipfile.ZipFile(zip_path, "r") as zf:
            content_files = [n for n in zf.namelist() if n.endswith(".content")]
            if not content_files:
                return result

            content_data = json.loads(zf.read(content_files[0]))
            page_ids = content_data.get("cPages", {}).get("pages", [])
            if not page_ids:
                page_ids = content_data.get("pages", [])

            ordered_ids = []
            for page in page_ids:
                if isinstance(page, dict):
                    ordered_ids.append(page.get("id", ""))
                else:
                    ordered_ids.append(str(page))

            rm_files = [n for n in zf.namelist() if n.endswith(".rm")]

            indexed = []
            for rm_file in rm_files:
                stem = Path(rm_file).stem
                try:
                    idx = ordered_ids.index(stem)
                except ValueError:
                    continue
                indexed.append((idx, rm_file))
            indexed.sort()

            for page_idx, rm_file in indexed:
                try:
                    result.append((page_idx, zf.read(rm_file)))
                except Exception:
                    log.warning("Failed to read %s in %s", rm_file, zip_path)
    except Exception:
        log.warning("Failed to read zip %s", zip_path, exc_info=True)
    return result


# ---------------------------------------------------------------------------
# Typed notes extraction (keyboard-typed text on reMarkable)
# ---------------------------------------------------------------------------


def extract_typed_notes(zip_path: Path) -> Dict[int, str]:
    """Extract keyboard-typed text notes per page from a reMarkable bundle.

    Returns a dict mapping page index (0-based) to the typed text on that page.
    Only pages with non-empty text are included.
    """
    from rmscene.text import TextDocument

    result: Dict[int, str] = {}
    for page_idx, rm_data in _load_rm_pages(zip_path):
        try:
            tree = read_tree(io.BytesIO(rm_data))
            if tree.root_text is None:
                continue
            doc = TextDocument.from_scene_item(tree.root_text)
            paragraphs = []
            for para in doc.contents:
                style = para.style.value if para.style else si.ParagraphStyle.PLAIN
                text = str(para).strip()
                if not text:
                    continue
                if style == si.ParagraphStyle.HEADING:
                    paragraphs.append(f"### {text}")
                elif style == si.ParagraphStyle.BOLD:
                    paragraphs.append(f"**{text}**")
                elif style in (si.ParagraphStyle.BULLET, si.ParagraphStyle.BULLET2):
                    paragraphs.append(f"- {text}")
                elif style == si.ParagraphStyle.CHECKBOX:
                    paragraphs.append(f"- [ ] {text}")
                elif style == si.ParagraphStyle.CHECKBOX_CHECKED:
                    paragraphs.append(f"- [x] {text}")
                else:
                    paragraphs.append(text)
            if paragraphs:
                result[page_idx] = "\n".join(paragraphs)
        except Exception:
            log.debug("Could not parse typed notes from page %d of %s",
                      page_idx, zip_path.name, exc_info=True)
    if result:
        log.info("Extracted typed notes from %d page(s) of %s",
                 len(result), zip_path.name)
    return result


# ---------------------------------------------------------------------------
# Ink stroke extraction and rendering
# ---------------------------------------------------------------------------

# Eraser tools to exclude from ink extraction
_ERASER_TOOLS = {si.Pen.ERASER, si.Pen.ERASER_AREA}

# reMarkable page dimensions in device units (classic RM1/RM2)
_RM_WIDTH = 1404.0
_RM_HEIGHT = 1872.0

# reMarkable Paper Pro renders at 227 DPI; stroke coordinates are in
# this native-DPI space (not the classic 1404×1872 viewport).
_RM_PRO_DPI = 227.0
_PDF_DPI = 72.0

# Map reMarkable pen colors to RGB tuples
_PEN_COLOR_MAP = {
    si.PenColor.BLACK: (0, 0, 0),
    si.PenColor.GRAY: (0.5, 0.5, 0.5),
    si.PenColor.WHITE: (1, 1, 1),
    si.PenColor.YELLOW: (0.9, 0.8, 0),
    si.PenColor.GREEN: (0, 0.6, 0),
    si.PenColor.PINK: (0.9, 0.2, 0.5),
    si.PenColor.BLUE: (0, 0.2, 0.8),
    si.PenColor.RED: (0.8, 0, 0),
    si.PenColor.GREEN_2: (0, 0.6, 0),
    si.PenColor.CYAN: (0, 0.6, 0.7),
    si.PenColor.MAGENTA: (0.7, 0, 0.7),
    si.PenColor.YELLOW_2: (0.9, 0.8, 0),
}


def extract_ink_strokes(zip_path: Path) -> Dict[int, List[si.Line]]:
    """Extract handwritten ink strokes (non-highlighter, non-eraser) per page.

    Returns a dict mapping page index (0-based) to a list of Line items.
    Only pages with ink strokes are included.
    """
    result: Dict[int, List[si.Line]] = {}
    for page_idx, rm_data in _load_rm_pages(zip_path):
        try:
            tree = read_tree(io.BytesIO(rm_data))
            strokes = []
            for item in tree.walk():
                if not isinstance(item, si.Line):
                    continue
                if si.Pen.is_highlighter(item.tool):
                    continue
                if item.tool in _ERASER_TOOLS:
                    continue
                if len(item.points) < 2:
                    continue
                strokes.append(item)
            if strokes:
                result[page_idx] = strokes
        except Exception:
            log.debug("Could not parse ink from page %d of %s",
                      page_idx, zip_path.name, exc_info=True)
    if result:
        total = sum(len(v) for v in result.values())
        log.info("Extracted %d ink stroke(s) from %d page(s) of %s",
                 total, len(result), zip_path.name)
    return result


def _is_paper_pro_coords(ink_by_page: Dict[int, List[si.Line]]) -> bool:
    """Detect whether stroke coordinates use Paper Pro native-DPI space.

    Paper Pro coordinates are centered at x=0 with the PDF scaled at
    ~227 DPI, producing stroke x-values far below 0 and y-values well
    above 1872.  Classic RM coordinates stay within 0..1404 / 0..1872
    (with small margins for edge strokes).
    """
    for lines in ink_by_page.values():
        for line in lines:
            for p in line.points:
                if p.x < -200 or p.y > 2100:
                    return True
    return False


def _rm_to_pdf_mapping(
    pdf_w: float, pdf_h: float,
    paper_pro: bool = False,
) -> tuple[float, float, float]:
    """Compute reMarkable → PDF coordinate mapping.

    Returns (rm_scale, x_offset, y_offset) where rm_scale is RM units
    per PDF point and offsets are the RM coordinates of the PDF origin.

    Classic RM (1404×1872 viewport, bestFit):
        pdf_coord = (rm_coord - offset) / rm_scale

    Paper Pro (227 DPI native coordinates, PDF centered at x≈0):
        Same formula with scale=227/72 and PDF horizontally centered.
    """
    if paper_pro:
        rm_scale = _RM_PRO_DPI / _PDF_DPI
        x_off = -(pdf_w * rm_scale) / 2
        y_off = 0.0
    else:
        rm_scale = min(_RM_WIDTH / pdf_w, _RM_HEIGHT / pdf_h)
        rendered_w = pdf_w * rm_scale
        rendered_h = pdf_h * rm_scale
        x_off = (_RM_WIDTH - rendered_w) / 2
        y_off = (_RM_HEIGHT - rendered_h) / 2
    return rm_scale, x_off, y_off


def _render_ink_on_pdf(doc, ink_by_page: Dict[int, List[si.Line]]) -> int:
    """Draw ink strokes onto PDF pages using PyMuPDF's drawing API.

    Detects the coordinate system (classic RM vs Paper Pro) from the
    stroke data, then maps RM coordinates to PDF page coordinates.
    Only renders strokes within the rendered PDF area.
    Returns the number of strokes rendered.
    """
    import pymupdf

    paper_pro = _is_paper_pro_coords(ink_by_page)
    if paper_pro:
        log.debug("Detected Paper Pro coordinate system")

    total_rendered = 0
    for page_idx, lines in ink_by_page.items():
        if page_idx >= len(doc):
            continue
        page = doc[page_idx]
        pdf_w = page.rect.width
        pdf_h = page.rect.height
        rm_scale, x_off, y_off = _rm_to_pdf_mapping(pdf_w, pdf_h, paper_pro)
        rendered_w = pdf_w * rm_scale
        rendered_h = pdf_h * rm_scale

        shape = page.new_shape()
        for line in lines:
            # Skip strokes entirely outside the rendered PDF area
            xs = [p.x for p in line.points]
            ys = [p.y for p in line.points]
            if max(xs) < x_off or min(xs) > x_off + rendered_w:
                continue
            if max(ys) < y_off or min(ys) > y_off + rendered_h:
                continue
            # Map RM coordinates to PDF coordinates, clamping to page
            points = [
                pymupdf.Point(
                    max(0, min((p.x - x_off) / rm_scale, pdf_w)),
                    max(0, min((p.y - y_off) / rm_scale, pdf_h)),
                )
                for p in line.points
            ]
            color = _PEN_COLOR_MAP.get(line.color, (0, 0, 0))
            width = max(0.8, line.thickness_scale / rm_scale)
            shape.draw_polyline(points)
            shape.finish(color=color, width=width, closePath=False)
            total_rendered += 1
        shape.commit()
    return total_rendered


# ---------------------------------------------------------------------------
# Handwritten OCR via Apple Vision (macOS only)
# ---------------------------------------------------------------------------


def _render_strokes_to_image(lines: List[si.Line]):
    """Render ink strokes to a PIL Image for OCR.

    Auto-sizes the canvas to fit all strokes (handles margin notes and
    scrolled pages where coordinates extend beyond the default RM viewport).
    Returns a greyscale PIL Image with black strokes on white background,
    or None if Pillow is not installed.
    """
    try:
        from PIL import Image, ImageDraw
    except ImportError:
        log.debug("Pillow not installed, cannot render strokes to image")
        return None

    # Compute bounding box of all strokes
    all_x: list[float] = []
    all_y: list[float] = []
    for line in lines:
        for p in line.points:
            all_x.append(p.x)
            all_y.append(p.y)
    if not all_x:
        return None

    x_min, x_max = min(all_x), max(all_x)
    y_min, y_max = min(all_y), max(all_y)

    # Add padding and compute canvas size
    pad = 40
    canvas_w = int(x_max - x_min) + 2 * pad
    canvas_h = int(y_max - y_min) + 2 * pad
    # Ensure minimum size for OCR quality
    canvas_w = max(canvas_w, 400)
    canvas_h = max(canvas_h, 200)

    img = Image.new("L", (canvas_w, canvas_h), 255)
    draw = ImageDraw.Draw(img)
    for line in lines:
        if len(line.points) < 2:
            continue
        # Shift coordinates so all strokes fit on canvas
        coords = [(p.x - x_min + pad, p.y - y_min + pad) for p in line.points]
        stroke_width = max(2, int(line.thickness_scale * 2))
        draw.line(coords, fill=0, width=stroke_width)
    return img


def _ocr_image_vision(image) -> str:
    """OCR a PIL Image using macOS Vision framework.

    Returns recognized text, or empty string if Vision is unavailable.
    """
    try:
        import Vision
        from Quartz import (
            CGImageCreate,
            CGDataProviderCreateWithData,
            CGColorSpaceCreateDeviceGray,
        )
    except ImportError:
        log.debug("pyobjc Vision framework not available")
        return ""

    # Convert PIL greyscale image to CGImage
    width, height = image.size
    raw_data = image.tobytes()
    provider = CGDataProviderCreateWithData(None, raw_data, len(raw_data), None)
    cg_image = CGImageCreate(
        width, height,
        8, 8,  # bits per component, bits per pixel
        width,  # bytes per row
        CGColorSpaceCreateDeviceGray(),
        0,  # bitmap info
        provider, None, False, 0,
    )
    if cg_image is None:
        log.debug("Failed to create CGImage for OCR")
        return ""

    # Create and run handwriting recognition request
    request = Vision.VNRecognizeTextRequest.alloc().init()
    request.setRecognitionLevel_(Vision.VNRequestTextRecognitionLevelAccurate)
    # Enable handwriting recognition alongside printed text
    request.setUsesLanguageCorrection_(True)

    handler = Vision.VNImageRequestHandler.alloc().initWithCGImage_options_(
        cg_image, None,
    )
    success = handler.performRequests_error_([request], None)
    if not success[0]:
        log.debug("Vision OCR request failed")
        return ""

    results = request.results()
    if not results:
        return ""

    # Collect recognized text lines sorted by vertical position (top to bottom)
    text_lines = []
    for observation in results:
        candidate = observation.topCandidates_(1)
        if candidate:
            text = candidate[0].string()
            # Vision y-coordinate is bottom-up, so 1-y for top-to-bottom sort
            bbox = observation.boundingBox()
            y_pos = 1 - bbox.origin.y - bbox.size.height
            text_lines.append((y_pos, text))

    text_lines.sort(key=lambda t: t[0])
    return "\n".join(line for _, line in text_lines)


def ocr_handwritten_notes(zip_path: Path) -> Dict[int, str]:
    """Extract and OCR handwritten notes from a reMarkable bundle.

    Renders ink strokes to images and uses Apple Vision for recognition.
    Returns a dict mapping page index (0-based) to recognized text.
    Falls back gracefully if Pillow or Vision are not available.
    """
    ink_by_page = extract_ink_strokes(zip_path)
    if not ink_by_page:
        return {}

    results: Dict[int, str] = {}
    for page_idx, lines in ink_by_page.items():
        if not lines:
            continue
        img = _render_strokes_to_image(lines)
        if img is None:
            return {}  # Pillow not available, stop trying
        text = _ocr_image_vision(img)
        if text.strip():
            results[page_idx] = text.strip()

    if results:
        log.info("OCR'd handwritten notes from %d page(s) of %s",
                 len(results), zip_path.name)
    return results
