# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-19 for v1.3; prior component notices are in NOTICE.
"""``ax-doclayout`` command line: one page raster in, canonical JSON out."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from ax_doclayout.contract import PageImage
from ax_doclayout.errors import AXDocLayoutError, InvalidPageImageError
from ax_doclayout.model import BACKENDS, AXDocLayout

_PDF_MAGIC = b"%PDF-"


def _load_rgb(path: Path):
    """Decode a page raster; PDFs are rejected by content and by file name."""

    try:
        with path.open("rb") as stream:
            head = stream.read(len(_PDF_MAGIC))
    except OSError as error:
        raise InvalidPageImageError(
            f"cannot read page raster: {error.strerror}"
        ) from None
    if head.startswith(_PDF_MAGIC) or path.suffix.lower() == ".pdf":
        raise InvalidPageImageError("the CLI takes a decoded page raster, not a PDF")
    import numpy as np
    from PIL import Image, UnidentifiedImageError

    try:
        with Image.open(path) as image:
            return np.asarray(image.convert("RGB"), dtype=np.uint8)
    except (OSError, UnidentifiedImageError, Image.DecompressionBombError, ValueError) as error:
        raise InvalidPageImageError(
            f"cannot decode page raster: {type(error).__name__}"
        ) from None


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(prog="ax-doclayout")
    commands = parser.add_subparsers(dest="command", required=True)
    predict = commands.add_parser(
        "predict", help="detect layout regions on one page raster"
    )
    predict.add_argument("image", type=Path)
    predict.add_argument("--page-index", type=int, required=True)
    predict.add_argument("--model", required=True)
    predict.add_argument("--revision", required=True)
    predict.add_argument("--backend", choices=BACKENDS, required=True)
    predict.add_argument("--output", default="-")
    arguments = parser.parse_args(argv)

    try:
        page = PageImage(
            rgb=_load_rgb(arguments.image), page_index=arguments.page_index
        )
        model = AXDocLayout.from_pretrained(
            arguments.model, revision=arguments.revision, backend=arguments.backend
        )
        payload = json.dumps(model.predict(page).to_json(), indent=2)
    except AXDocLayoutError as error:
        print(f"ax-doclayout: {error}", file=sys.stderr)
        return 2
    if arguments.output == "-":
        print(payload)
    else:
        Path(arguments.output).write_text(payload + "\n", encoding="utf-8")
    return 0


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