#!/usr/bin/env python3
"""
anydoc_vlm.py — document -> Markdown + vision description pipeline.

This is the conversion engine behind the `dsh-anydoc-markdown` DeepSeek Harness
host plugin. It turns a document (Word, PowerPoint, Excel, OpenDocument, RTF,
EPUB, CSV or PDF) into clean GitHub-Flavored Markdown using the Rust
`firecrawl-anydoc` crate, then replaces every embedded image in the output with
a text description produced by an OpenAI-compatible vision model (VLM).

Design notes
------------
* anydoc is a pure local, dependency-free Rust converter. It exposes
  ``anydoc.to_document(data)`` which returns the full document model plus the
  embedded binary assets (images, MIME types, source part / alt text markers) on
  ``Document.assets``. That is the seam we use — exactly the ``llm_client`` +
  ``llm_model`` pattern MarkItDown pioneered, implemented for anydoc here.
* Markdown cannot embed bytes, so anydoc renders an embedded image as its alt
  text (a bare paragraph) and drops it entirely when the alt text is empty.
  This wrapper therefore:
    * walks ``Document.blocks`` to recover every image inline *in document
      order* together with its asset, and
    * substitutes the rendered alt-text paragraph (or, for alt-less images that
      anydoc dropped, an appended note) with a description from the VLM.
* Images are NEVER all sent in one API request — that degrades vision
  performance. They are split into batches of ``max_images_per_request``
  (default 10) and each batch is sent as its own request. Batches larger than
  the default are configurable.
* The VLM client is OpenAI-compatible (``/chat/completions`` with
  ``image_url``/base64 parts). When no endpoint / model / key is configured the
  wrapper falls back to a deterministic metadata description so it always
  yields output (and is testable offline).

Everything reads bytes from a single file path and writes the finished Markdown
to stdout (or a file / JSON meta object).
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import re
import sys
import time
import urllib.error
import urllib.request
from typing import Any, Callable, Dict, List, Optional, Tuple

try:
    import anydoc
except Exception as exc:  # pragma: no cover - import guard
    sys.stderr.write(f"anydoc_vlm: cannot import anydoc: {exc}\n")
    sys.stderr.write("Install it with: pip install firecrawl-anydoc\n")
    raise

# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #

DEFAULT_MAX_IMAGES_PER_REQUEST = 10
DEFAULT_DESCRIPTION_LIMIT = 1200
DEFAULT_TIMEOUT_SECONDS = 120
DEFAULT_IMAGE_PARAGRAPH = "![{desc}]({uri})"
# Vision calls are retried on a transient failure (HTTP error, timeout, or an
# unparseable model response) so a single flaky request does not degrade a whole
# batch to the offline description. This mirrors the fail-closed philosophy but
# gives the endpoint a chance to recover first.
DEFAULT_VISION_ATTEMPTS = 3
DEFAULT_RETRY_BACKOFF_SECONDS = 1.0


# --------------------------------------------------------------------------- #
# Small helpers
# --------------------------------------------------------------------------- #

def _data_uri(media_type: str, data: bytes) -> str:
    """Base64 data URI for a vision request."""
    return f"data:{media_type};base64,{base64.b64encode(data).decode('ascii')}"


def _metadata_description(image: Dict[str, Any], reason: str = "no vision model was configured") -> str:
    """Deterministic, offline description derived from the asset metadata.

    Used when no VLM is configured (or the VLM call fails) so the pipeline still
    turns an image into *some* readable text. The description names the media
    type, its source part, byte size and any alt text the author supplied. The
    ``reason`` is surfaced in the note so the reader can tell a genuinely
    unconfigured vision model apart from a request that simply failed.
    """
    media = image.get("media_type") or "image"
    part = image.get("origin_part")
    size = len(image.get("data") or b"")
    alt = (image.get("alt") or "").strip()
    bits = [f"an embedded {media}"]
    if part:
        bits.append(f"from `{part}`")
    bits.append(f"{size} bytes")
    if alt:
        bits.append(f"author alt text: {alt}")
    return ", ".join(bits) + f". [Offline metadata description - {reason}]"


def _extract_json_array(text: str) -> Optional[List[Any]]:
    """Best-effort JSON-array extraction, tolerating prose/fences around it.

    The strict path is ``json.loads`` on the whole body. If that fails, we look
    for a balanced ``[...]`` region (or a fenced ```json ... ``` block) and try to
    parse just that, so a response like ``Sure, here are the descriptions: [..]``
    still yields an array instead of degrading the batch to offline.
    """
    if not isinstance(text, str):
        return None
    s = text.strip()
    if not s:
        return None
    # Direct parse first.
    try:
        parsed = json.loads(s)
        if isinstance(parsed, list):
            return parsed
    except (json.JSONDecodeError, ValueError):
        pass
    # Fenced ```json ... ``` block.
    fence = re.search(r"```(?:json)?\s*(\[.*?\])\s*```", s, re.S)
    candidate = fence.group(1) if fence else None
    if candidate is None:
        # Fall back to the outermost [...] region (first '[' to last ']').
        start = s.find("[")
        end = s.rfind("]")
        if start != -1 and end > start:
            candidate = s[start:end + 1]
    if candidate is None:
        return None
    try:
        parsed = json.loads(candidate)
        if isinstance(parsed, list):
            return parsed
    except (json.JSONDecodeError, ValueError):
        return None
    return None


def _cap(text: str, limit: int) -> str:
    text = (text or "").strip()
    if len(text) <= limit:
        return text
    return text[:limit].rstrip() + "…"


# --------------------------------------------------------------------------- #
# Image extraction from the anydoc document model
# --------------------------------------------------------------------------- #

def extract_images(document: Any) -> List[Dict[str, Any]]:
    """Walk ``document.blocks`` and return every embedded image inline in order.

    Each returned mapping contains the resolved asset bytes plus the metadata a
    vision model call (or the offline fallback) needs:
    ``{asset_id, alt, media_type, origin_part, data, url}``.

    URL-kind images (``source.kind == 'url'``) are returned too but their
    ``data``/``asset_id`` are ``None`` — they already render as ordinary
    Markdown images, so callers should leave them untouched unless they want to
    describe an external image by fetching it.
    """
    images: List[Dict[str, Any]] = []
    assets: Optional[List[Any]] = getattr(document, "assets", None) or []

    def walk_inline(inline: Any) -> None:
        kind = getattr(inline, "kind", None)
        if kind != "image":
            return
        source = getattr(inline, "source", None)
        if source is None:
            return
        src_kind = getattr(source, "kind", None)
        alt = (getattr(inline, "alt", None) or "").strip()
        if src_kind == "url":
            images.append({
                "asset_id": None,
                "alt": alt,
                "media_type": None,
                "origin_part": None,
                "data": None,
                "url": getattr(source, "url", None),
            })
            return
        asset_id = getattr(source, "asset_id", None)
        asset = assets[asset_id] if asset_id is not None and asset_id < len(assets) else None
        images.append({
            "asset_id": asset_id,
            "alt": alt,
            "media_type": getattr(asset, "media_type", None) if asset else None,
            "origin_part": getattr(asset, "origin_part", None) if asset else None,
            "data": getattr(asset, "data", None) if asset else None,
            "url": None,
        })

    def walk_block(block: Any) -> None:
        content = getattr(block, "content", None)
        if isinstance(content, list):
            for inline in content:
                walk_inline(inline)
        nested = getattr(block, "blocks", None)
        if isinstance(nested, list):
            for sub in nested:
                walk_block(sub)
        table = getattr(block, "table", None)
        if table is not None:
            for _row in getattr(table, "rows", []) or []:
                for cell in getattr(_row, "cells", []) or []:
                    walk_block(cell)

    for block in getattr(document, "blocks", None) or []:
        walk_block(block)

    return [img for img in images if img.get("data") is not None or img.get("url") is not None]


# --------------------------------------------------------------------------- #
# PDF image extraction (pdfplumber)
# --------------------------------------------------------------------------- #
#
# A PDF has no anydoc document model, so `to_document()` raises
# `UnsupportedError` and the generic `extract_images()` path yields nothing. A
# born-digital PDF still carries real embedded raster images (JPEG / raw
# bitmaps), so we recover them with an auxiliary PDF library and re-encode each
# one to PNG (or pass the JPEG bytes through) so the vision client can send it as
# an `image_url` data URI. These images have no per-image alt text, so they land
# in the described reference list the same way an alt-less docx/pptx image does.

def _pdf_stream_pixel_size(image: Dict[str, Any]) -> tuple:
    """Return (width, height) in pixels for a pdfplumber image object.

    pdfplumber's ``image['width']/['height']`` are the *display* dimensions in
    PDF points (floats); the true pixel size lives on the stream's ``Width`` /
    ``Height`` (or ``Columns`` / ``Rows``) entries.
    """
    stream = image.get("stream")
    if stream is None:
        return None, None
    attrs = getattr(stream, "attrs", None) or {}
    try:
        width = int(attrs.get("Width") or attrs.get("Columns"))
        height = int(attrs.get("Height") or attrs.get("Rows"))
    except (TypeError, ValueError):
        return None, None
    return width, height


def _pdf_image_bytes(image: Dict[str, Any]) -> tuple:
    """Return ``(encoded_bytes, media_type)`` for one pdfplumber image object.

    Handles the image stream encodings seen in real PDFs:

    * JPEG streams (``DCTDecode``): ``stream.get_data()`` is already the encoded
      JPEG bytes, so they are validated and passed through unchanged.
    * Raw bitmaps (``FlateDecode``) with a plain ``DeviceRGB``/``DeviceGray``/
      ``DeviceCMYK`` colours space, ``ICCBased``, or an ``Indexed`` palette: the
      raw samples are reshaped with Pillow and re-encoded to PNG.
    """
    import io as _io

    from PIL import Image

    stream = image.get("stream")
    if stream is None:
        raise ValueError("pdf image has no stream")
    width, height = _pdf_stream_pixel_size(image)
    if not (width and height):
        raise ValueError("pdf image has no usable pixel dimensions")

    def _name(value) -> str:
        return getattr(value, "name", None) or value

    filter_list = stream.get_filters() or []
    filter_names = [
        _name(f[0]) if isinstance(f, (list, tuple)) and f else _name(f)
        for f in filter_list
    ]
    data = stream.get_data()

    # JPEG: get_data() returns the encoded JPEG bytes (header FF D8 FF ..).
    if any("DCTDecode" in name for name in filter_names):
        # Validate it decodes, then pass the original bytes through unchanged.
        Image.open(_io.BytesIO(data)).load()
        return data, "image/jpeg"

    # Colorspace is a 1-element list wrapping the spec (e.g. ['DeviceRGB'] or
    # [['Indexed','DeviceRGB',<hival>,<lookup>]]).
    spec = image.get("colorspace")
    if isinstance(spec, list) and len(spec) == 1:
        spec = spec[0]
    if isinstance(spec, (list, tuple)):
        components = [_name(x) for x in spec]
        first = components[0] if components else ""
    else:
        first = components = [_name(spec)] if spec is not None else [""]
    first = components[0] if components else ""

    def _from_rgb():
        return Image.frombytes("RGB", (width, height), data[: width * height * 3])

    if first == "Indexed":
        lookup = spec[3]
        palette = lookup.get_data() if hasattr(lookup, "get_data") else bytes(lookup)
        pim = Image.frombytes("P", (width, height), data[: width * height])
        pim.putpalette(palette[:768].ljust(768, b"\x00"))
        pim = pim.convert("RGB")
    elif first == "ICCBased":
        icc = spec[1]
        n_components = (getattr(icc, "attrs", None) or {}).get("N")
        if n_components == 4:
            pim = Image.frombytes("CMYK", (width, height), data[: width * height * 4]).convert("RGB")
        else:
            pim = _from_rgb()
    elif first == "DeviceCMYK":
        pim = Image.frombytes("CMYK", (width, height), data[: width * height * 4]).convert("RGB")
    elif first in ("DeviceGray", "CalGray"):
        pim = Image.frombytes("L", (width, height), data[: width * height])
    elif first in ("DeviceRGB", "CalRGB", ""):
        pim = _from_rgb()
    else:
        raise ValueError(f"unhandled pdf colours space {first!r}")

    buf = _io.BytesIO()
    pim.convert("RGB").save(buf, format="PNG")
    return buf.getvalue(), "image/png"


def extract_pdf_images(data: bytes) -> List[Dict[str, Any]]:
    """Extract embedded images from PDF ``data`` via pdfplumber.

    Returns the same image-mapping shape as :func:`extract_images`: ``{asset_id,
    alt, media_type, origin_part, data, url}``. ``alt`` is empty (PDFs carry no
    per-image alt text) so these entries are appended to the described reference
    list rather than substituted inline. ``data`` is re-encoded to PNG (or the
    original JPEG bytes), so the vision client can send it directly.

    Fails closed: if pdfplumber/PIL is unavailable or the PDF cannot be parsed,
    returns ``[]`` and the conversion still produces the (text) Markdown.
    """
    import io as _io

    try:
        import pdfplumber
    except Exception as exc:  # pragma: no cover - optional dependency
        sys.stderr.write(f"anydoc_vlm: pdfplumber not available for PDF image extraction: {exc}\n")
        return []

    images: List[Dict[str, Any]] = []
    try:
        pdf = pdfplumber.open(_io.BytesIO(data))
    except Exception as exc:  # pragma: no cover
        sys.stderr.write(f"anydoc_vlm: could not open PDF for image extraction: {exc}\n")
        return []

    asset_id = 0
    for page in pdf.pages:
        page_number = getattr(page, "page_number", None)
        for image in page.images:
            try:
                encoded, media_type = _pdf_image_bytes(image)
            except Exception as exc:  # pragma: no cover - fail closed per image
                sys.stderr.write(f"anydoc_vlm: skipped a PDF image on page {page_number}: {exc}\n")
                continue
            images.append({
                "asset_id": asset_id,
                "alt": "",
                "media_type": media_type,
                "origin_part": f"page {page_number}" if page_number else None,
                "data": encoded,
                "url": None,
            })
            asset_id += 1
    try:
        pdf.close()
    except Exception:  # pragma: no cover
        pass
    return images


# --------------------------------------------------------------------------- #
# Vision client (OpenAI-compatible) with batching
# --------------------------------------------------------------------------- #

class VisionClient:
    """Batch-capable OpenAI-compatible vision client.

    ``describe_many`` splits the images into groups of ``max_per_request`` and
    issues one chat-completions request per group, so no single request carries
    more than a configurable number of images (default 10). Order is preserved.
    """

    def __init__(
        self,
        endpoint: Optional[str],
        model: Optional[str],
        api_key: Optional[str],
        max_per_request: int,
        timeout: float,
        *,
        describe_batch: Optional[Callable[[List[Dict[str, Any]]], List[str]]] = None,
        attempts: int = DEFAULT_VISION_ATTEMPTS,
        retry_backoff: float = DEFAULT_RETRY_BACKOFF_SECONDS,
    ) -> None:
        self.endpoint = (endpoint or "").rstrip("/") or None
        self.model = model or None
        self.api_key = api_key or None
        self.max_per_request = max(1, int(max_per_request))
        self.timeout = timeout
        self._describe_batch_override = describe_batch
        self.attempts = max(1, int(attempts))
        self.retry_backoff = max(0.0, float(retry_backoff))
        self.requests = 0
        self.batch_sizes: List[int] = []
        self.retries: int = 0

    @property
    def configured(self) -> bool:
        return bool(self.endpoint and self.model)

    def describe_many(self, images: List[Dict[str, Any]]) -> List[str]:
        """Return one description per image, batched by ``max_per_request``."""
        results: List[str] = []
        for start in range(0, len(images), self.max_per_request):
            batch = images[start:start + self.max_per_request]
            self.requests += 1
            self.batch_sizes.append(len(batch))
            results.extend(self._describe_batch(batch))
        return results

    def _describe_batch(self, batch: List[Dict[str, Any]]) -> List[str]:
        if self._describe_batch_override is not None:
            return self._describe_batch_override(batch)
        if not self.configured:
            return [_metadata_description(img) for img in batch]
        return self._openai_batch(batch)

    def _openai_batch(self, batch: List[Dict[str, Any]]) -> List[str]:
        content: List[Dict[str, Any]] = [{
            "type": "text",
            "text": (
                f"You are a vision model. There are {len(batch)} image(s) below, "
                "labelled IMAGE 1..N in order. For EACH image produce exactly one "
                "concise description of its visible content. Respond with a single "
                "JSON array of exactly N strings, one per image, in the same order. "
                "No other text."
            ),
        }]
        for idx, img in enumerate(batch, start=1):
            if img.get("data") is None:
                # URL image: reference it directly if we can.
                url = img.get("url")
                if not url:
                    content.append({"type": "text", "text": f"IMAGE {idx}: <no bytes>"})
                    continue
                content.append({
                    "type": "text",
                    "text": f"IMAGE {idx}: {url}",
                })
                continue
            media = img.get("media_type") or "image/png"
            content.append({
                "type": "image_url",
                "image_url": {"url": _data_uri(media, img["data"])},
            })

        body: Dict[str, Any] = {
            "model": self.model,
            "messages": [{"role": "user", "content": content}],
            "temperature": 0,
            "max_tokens": 4096,
        }
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        url = f"{self.endpoint}/chat/completions" if not self.endpoint.endswith("chat/completions") \
            else self.endpoint
        request = urllib.request.Request(
            url, data=json.dumps(body).encode("utf-8"), headers=headers, method="POST"
        )

        # Retry on a transient failure (HTTP error / timeout / unparseable
        # response) so one flaky request does not degrade a whole batch.
        last_reason = "no response"
        for attempt in range(1, self.attempts + 1):
            payload = None
            try:
                with urllib.request.urlopen(request, timeout=self.timeout) as response:
                    payload = json.loads(response.read().decode("utf-8"))
            except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, json.JSONDecodeError) as exc:
                last_reason = f"request failed: {exc}"
                if attempt < self.attempts:
                    time.sleep(self.retry_backoff * attempt)
                    self.retries += 1
                    continue
                sys.stderr.write(f"anydoc_vlm: vision request failed after {self.attempts} attempts: {exc}\n")
                return [_metadata_description(img, "vision request failed after retries") for img in batch]

            # Extract assistant text (handles both string and array content).
            assistant = (payload.get("choices") or [{}])[0].get("message") or {}
            text = assistant.get("content")
            if isinstance(text, list):
                text = "".join(
                    part.get("text", "") for part in text if isinstance(part, dict)
                )
            parsed = self._parse_descriptions(text, batch)
            if parsed is not None:
                return parsed
            last_reason = "vision response was not a JSON array"
            if attempt < self.attempts:
                time.sleep(self.retry_backoff * attempt)
                self.retries += 1
                continue

        sys.stderr.write(
            f"anydoc_vlm: vision batch gave no usable descriptions after {self.attempts} attempts ({last_reason}); "
            f"falling back to offline metadata\n"
        )
        return [_metadata_description(img, "vision unavailable after retries") for img in batch]

    def _parse_descriptions(self, text: Optional[str], batch: List[Dict[str, Any]]) -> Optional[List[str]]:
        """Return a clean JSON-array parse, or ``None`` if it cannot be salvaged.

        Requires exactly one string per image (``len == len(batch)``) so a
        truncated or length-mismatched body is rejected and retried rather than
        silently misaligned. Tolerates prose / fenced code around the array.
        """
        parsed = _extract_json_array(text) if isinstance(text, str) else None
        if parsed is None or len(parsed) != len(batch):
            return None
        return [_cap(str(item), DEFAULT_DESCRIPTION_LIMIT) for item in parsed]


# --------------------------------------------------------------------------- #
# Substitution
# --------------------------------------------------------------------------- #

def _image_markdown(image: Dict[str, Any], description: str) -> str:
    """Render one embedded image as Markdown with the description as alt text.

    The URI is normally ``asset://<id>`` (a pointer to the original bytes). When
    the image was persisted to disk as a PNG (``saved_file`` set — see
    :func:`_save_image_png`), the reference points at that file instead, so a
    Markdown reader resolves the actual PNG sitting next to the ``.md`` file.
    URL-kind images are left as the ordinary external Markdown image.
    """
    if image.get("url"):
        return f"![{_cap(description, DEFAULT_DESCRIPTION_LIMIT)}]({image['url']})"
    saved = image.get("saved_file")
    if saved:
        return f"![{_cap(description, DEFAULT_DESCRIPTION_LIMIT)}]({saved})"
    uri = f"asset://{image['asset_id']}"
    return DEFAULT_IMAGE_PARAGRAPH.format(desc=_cap(description, DEFAULT_DESCRIPTION_LIMIT), uri=uri)


def _save_image_png(image: Dict[str, Any], save_dir: str, index: int) -> Optional[str]:
    """Re-encode one embedded image to a PNG file and write it into ``save_dir``.

    Returns the relative file name (e.g. ``image_0001.png``) when the image was
    written, or ``None`` when the bytes cannot be decoded (fail-closed: the
    caller keeps the ``asset://`` reference so nothing is silently lost). The
    image is embedded as a raster (``data`` not ``None``); URL-kind images are
    never persisted and return ``None``.

    Transparent images (RGBA/LA/PA) are composited onto white so the PNG is a
    plain RGB bitmap rather than black-backed transparent PNG.
    """
    import io as _io

    from PIL import Image

    data = image.get("data")
    if not data:
        return None
    try:
        img = Image.open(_io.BytesIO(data))
        img.load()
    except Exception as exc:  # pragma: no cover - fail closed per image
        sys.stderr.write(f"anydoc_vlm: couldn't decode image {index} to save as PNG: {exc}\n")
        return None
    try:
        if img.mode in ("RGBA", "LA", "PA"):
            background = Image.new("RGB", img.size, (255, 255, 255))
            alpha = img.getchannel("A") if img.mode == "RGBA" else img.convert("RGBA").getchannel("A")
            background.paste(img.convert("RGBA"), mask=alpha)
            img = background
        else:
            img = img.convert("RGB")
    except Exception as exc:  # pragma: no cover - fail closed per image
        sys.stderr.write(f"anydoc_vlm: couldn't convert image {index} to RGB for saving: {exc}\n")
        return None
    filename = f"image_{index:04d}.png"
    target = os.path.join(save_dir, filename)
    try:
        img.save(target, format="PNG")
    except Exception as exc:  # pragma: no cover - fail closed per image
        sys.stderr.write(f"anydoc_vlm: couldn't write {target}: {exc}\n")
        return None
    return filename


def save_images_as_png(
    images: List[Dict[str, Any]], save_dir: str
) -> List[str]:
    """Persist every embeddable (has ``data``) image to ``save_dir`` as a PNG.

    Mutates each such image in place to set ``saved_file`` (the relative file
    name) so :func:`_image_markdown` can reference the on-disk PNG. Returns the
    relative file names written, in document order. Fails closed per image: a
    non-decodable image is skipped and keeps its ``asset://`` reference.
    """
    os.makedirs(save_dir, exist_ok=True)
    saved: List[str] = []
    index = 0
    for image in images:
        if image.get("data") is None:
            continue  # URL-kind image: already a Markdown image, nothing to save
        index += 1
        filename = _save_image_png(image, save_dir, index)
        if filename:
            image["saved_file"] = filename
            saved.append(filename)
    return saved


def substitute(markdown: str, entries: List[Dict[str, Any]]) -> str:
    """Replace each image's rendered alt-text paragraph with its description.

    anydoc renders an embedded image (when it has alt text) as a bare paragraph
    containing exactly that alt text. We walk the images in document order and
    replace the next standalone-occurence of each alt text after the previous
    image position, so repeated alt texts and interleaved prose stay aligned.
    Images anydoc dropped (empty alt) are appended as a described reference list,
    so no image content is lost.

    ``entries`` is the list returned by :func:`extract_images` extended with
    ``desc`` (the description). URL-kind images are already Markdown images and
    are left untouched.
    """
    import re

    edits: List[tuple] = []
    dropped: List[Dict[str, Any]] = []
    cursor = 0

    for img in entries:
        if img.get("data") is None and img.get("url"):
            continue  # already a Markdown image; nothing to substitute
        alt = (img.get("alt") or "").strip()
        if not alt:
            dropped.append(img)
            continue
        pattern = re.compile(r"(?m)^[ \t]*" + re.escape(alt) + r"[ \t]*$")
        match = pattern.search(markdown, cursor)
        if match is None:
            dropped.append(img)
            continue
        description = _image_markdown(img, img.get("desc") or _metadata_description(img))
        edits.append((match.start(), match.end(), description))
        cursor = match.end()

    for start, end, description in reversed(edits):
        markdown = markdown[:start] + description + markdown[end:]

    if dropped:
        lines = ["", "> **Embedded images** (anydoc dropped these from the body; described here):", ""]
        for img in dropped:
            uri = _image_markdown(img, img.get("desc") or _metadata_description(img))
            lines.append(f"- {uri}")
        markdown = markdown.rstrip() + "\n" + "\n".join(lines) + "\n"

    return markdown


# --------------------------------------------------------------------------- #
# Conversion engine
# --------------------------------------------------------------------------- #

def convert(
    path: str,
    *,
    max_images_per_request: int = DEFAULT_MAX_IMAGES_PER_REQUEST,
    vision: Optional[VisionClient] = None,
    ocr: str = "reject",
    format_name: Optional[str] = None,
    save_images_dir: Optional[str] = None,
) -> Dict[str, Any]:
    """Convert ``path`` to Markdown, describing every embedded image via VLM.

    When ``save_images_dir`` is given, every embedded raster image is also
    persisted to that directory as a PNG file and the Markdown image references
    are rewritten to point at the on-disk PNGs (relative file names). This lets
    the plugin keep the extracted images next to the saved ``.md`` file.

    Returns ``{"markdown": str, "meta": dict}`` where ``meta`` records the
    conversion and vision statistics (assets, images, requests, batch sizes,
    saved-image file names).
    """
    with open(path, "rb") as handle:
        data = handle.read()

    # 1) Base Markdown from the bytes (required). `ocr='hosted'` routes scanned
    #    PDF pages to Firecrawl Parse; the local Rust crate does no OCR itself.
    try:
        markdown = anydoc.to_markdown_bytes(data, format_name, ocr=ocr) if format_name \
            else anydoc.to_markdown_bytes(data, ocr=ocr)
    except (anydoc.NeedsOcrError, anydoc.EncryptedError, anydoc.UnsupportedError,
            anydoc.MalformedError, anydoc.MissingPartError, anydoc.ResourceLimitError,
            anydoc.HostedError) as exc:
        raise type(exc)(f"anydoc conversion of {os.path.basename(path)} failed: {exc}") from exc

    # 2) Embedded images for the vision pass — best-effort. Formats with a document
    #    model (docx/pptx/xlsx/odt/…) are recovered from `Document.assets` via
    #    `to_document`. PDFs have no such model (`to_document` raises
    #    `UnsupportedError`); for those we recover the images directly from the PDF
    #    bytes with `pdfplumber` so they are described too.
    images: List[Dict[str, Any]] = []
    asset_count = 0
    try:
        document = anydoc.to_document(data, format_name) if format_name else anydoc.to_document(data)
        images = extract_images(document)
        asset_count = len(getattr(document, "assets", None) or [])
    except anydoc.ConvertError as exc:
        # PDFs (and other "direct-to-markdown" formats) have no document model, so
        # `to_document` raises here. Rather than silently skipping imagery, recover
        # the embedded raster images from a PDF with an auxiliary library so they
        # can still be described by the vision model (fail-closed when unavailable).
        is_pdf = path.lower().endswith(".pdf") or (format_name or "").lower() == "pdf"
        if is_pdf:
            try:
                images = extract_pdf_images(data)
                asset_count = sum(1 for img in images if img.get("data") is not None)
            except Exception as pdf_exc:  # pragma: no cover - fail closed
                sys.stderr.write(f"anydoc_vlm: PDF image extraction failed: {pdf_exc}\n")
        else:
            sys.stderr.write(f"anydoc_vlm: no document model for {os.path.basename(path)} "
                             f"({type(exc).__name__}); skipping embedded-image description\n")

    embeddable = [img for img in images if img.get("data") is not None]

    if embeddable and vision is not None:
        descriptions = vision.describe_many(embeddable)
        desc_by_asset = {}
        for img, desc in zip(embeddable, descriptions):
            desc_by_asset[(img.get("asset_id"), img.get("origin_part"), img.get("alt"))] = desc
        for img in images:
            img["desc"] = desc_by_asset.get(
                (img.get("asset_id"), img.get("origin_part"), img.get("alt"))
            )
    else:
        for img in images:
            img["desc"] = _metadata_description(img)

    # Persist every embeddable image as a PNG beside the Markdown (when asked to).
    # Done after the vision pass so the description is unrelated to the saved
    # bytes, and before `substitute` so the image references point at the files.
    saved_images: List[str] = []
    if save_images_dir:
        saved_images = save_images_as_png(images, save_images_dir)

    final_markdown = substitute(markdown, images)

    meta = {
        "path": path,
        "assets": asset_count,
        "images": len(images),
        "images_described": sum(1 for img in images if img.get("desc")),
        "image_dir": save_images_dir,
        "images_saved": len(saved_images),
        "saved_images": saved_images,
        "vision_configured": bool(vision and vision.configured) if vision else False,
        "vision_requests": getattr(vision, "requests", 0) if vision else 0,
        "vision_retries": getattr(vision, "retries", 0) if vision else 0,
        "vision_batch_sizes": getattr(vision, "batch_sizes", []) if vision else [],
        "max_images_per_request": max_images_per_request,
        "dropped_images": [img.get("asset_id") for img in images if not (img.get("alt") or "").strip() and img.get("data") is not None],
    }
    return {"markdown": final_markdown, "meta": meta}


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

def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="anydoc_vlm.py",
        description="Convert a document to Markdown via firecrawl-anydoc, describing embedded images with a VLM.",
    )
    parser.add_argument("path", help="Path to the document to convert.")
    parser.add_argument("--format", dest="format_name", default=None,
                        help="Explicit format name (e.g. 'csv'); only needed for signature-less formats.")
    parser.add_argument("--ocr", dest="ocr", default="reject", choices=["reject", "hosted"],
                        help="'hosted' sends scanned PDF pages to Firecrawl Parse.")
    parser.add_argument("--max-images", dest="max_images", type=int, default=DEFAULT_MAX_IMAGES_PER_REQUEST,
                        help=f"Max images per vision request (default {DEFAULT_MAX_IMAGES_PER_REQUEST}).")
    parser.add_argument("--vision-endpoint", dest="vision_endpoint", default=os.environ.get("ANYDOC_VISION_ENDPOINT", ""),
                        help="OpenAI-compatible vision endpoint (e.g. https://api.openai.com/v1).")
    parser.add_argument("--vision-model", dest="vision_model", default=os.environ.get("ANYDOC_VISION_MODEL", ""),
                        help="Vision model id (e.g. gpt-4o).")
    parser.add_argument("--api-key", dest="api_key", default=os.environ.get("OPENAI_API_KEY", ""),
                        help="API key (or set OPENAI_API_KEY).")
    parser.add_argument("--timeout", dest="timeout", type=float, default=DEFAULT_TIMEOUT_SECONDS,
                        help=f"Vision request timeout seconds (default {DEFAULT_TIMEOUT_SECONDS}).")
    parser.add_argument("--mock-vision", dest="mock_vision", action="store_true",
                        help="Use the offline metadata description without calling a VLM (for testing).")
    parser.add_argument("--save-images-dir", dest="save_images_dir", default=None,
                        help="Directory to write every extracted embedded image as a PNG file; the "
                             "Markdown image references are rewritten to point at the saved files.")
    parser.add_argument("--json", dest="emit_json", action="store_true",
                        help="Emit a JSON envelope {\"markdown\", \"meta\"} on stdout instead of plain Markdown.")
    parser.add_argument("--out", dest="out_path", default=None,
                        help="Write the Markdown to this file.")
    return parser


def main(argv: Optional[List[str]] = None) -> int:
    args = _build_parser().parse_args(argv)
    if not os.path.isfile(args.path):
        sys.stderr.write(f"anydoc_vlm: no such file: {args.path}\n")
        return 2

    # Always build the batch-capable client. When no endpoint/model is
    # configured (default), describe_batch falls back to the offline metadata
    # description WITHOUT any HTTP call — but the batching, request counting and
    # order preservation are still exercised. `--mock-vision` forces the offline
    # description even when an endpoint/model is configured (for testing).
    vision = VisionClient(
        endpoint=args.vision_endpoint,
        model=args.vision_model,
        api_key=args.api_key,
        max_per_request=args.max_images,
        timeout=args.timeout,
        describe_batch=(lambda batch: [_metadata_description(img) for img in batch]) if args.mock_vision else None,
    )

    try:
        result = convert(
            args.path,
            max_images_per_request=args.max_images,
            vision=vision,
            ocr=args.ocr,
            format_name=args.format_name,
            save_images_dir=args.save_images_dir,
        )
    except Exception as exc:  # anydoc errors and file IO
        sys.stderr.write(f"anydoc_vlm: {exc}\n")
        return 1

    if args.out_path:
        with open(args.out_path, "w", encoding="utf-8") as handle:
            handle.write(result["markdown"])
    if args.emit_json:
        sys.stdout.write(json.dumps({
            "markdown": result["markdown"],
            "meta": result["meta"],
        }, indent=2) + "\n")
    else:
        sys.stdout.write(result["markdown"])
        if not result["markdown"].endswith("\n"):
            sys.stdout.write("\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())