WCX/scripts/process_pending_ocr.py

#!/usr/bin/env python3

"""
Process all movies waiting for OCR in the WCX SQLite index.

The script queries the movie table for records whose OCR status is "pending"
and invokes check_ocr.py once for each movie ID.

check_ocr.py performs the actual OCR request, parsing, validation, and database
update. This batch script only coordinates those individual executions.

Processing continues if one movie fails. A final summary reports how many
movies completed successfully and how many failed or require manual review.

By default, all pending movies are processed. Use --limit to restrict the
number of movies during testing.

Usage:
    process_pending_ocr.py
    process_pending_ocr.py --limit 1
    process_pending_ocr.py --database /path/to/wcx-test.db

Requirements:
    GOOGLE_VISION_API_KEY must be available in the environment.
"""

import argparse
import os
import sqlite3
import subprocess
import sys
from pathlib import Path


DEFAULT_DATABASE_FILE = Path(
    "/storage/disk1/WCX/database/wcx.db"
)
CHECK_OCR_SCRIPT = Path(
    "/storage/disk1/WCX/scripts/check_ocr.py"
)


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Process movies whose OCR status is pending."
    )

    parser.add_argument(
        "--database",
        type=Path,
        default=DEFAULT_DATABASE_FILE,
        help=f"SQLite database; default: {DEFAULT_DATABASE_FILE}",
    )

    parser.add_argument(
        "--limit",
        type=int,
        help="Process at most N pending movies.",
    )

    return parser.parse_args()


def get_pending_movies(
    database_file: Path,
    limit: int | None,
) -> list[tuple[str, str]]:
    sql = """
        SELECT id, name
        FROM movie
        WHERE ocr_status = 'pending'
        ORDER BY published, id
    """

    parameters: tuple[int, ...] = ()

    if limit is not None:
        sql += "\nLIMIT ?"
        parameters = (limit,)

    with sqlite3.connect(database_file) as connection:
        rows = connection.execute(sql, parameters).fetchall()

    return [
        (str(movie_id), str(name))
        for movie_id, name in rows
    ]


def get_ocr_status(
    database_file: Path,
    movie_id: str,
) -> tuple[str | None, str | None]:
    with sqlite3.connect(database_file) as connection:
        row = connection.execute(
            """
            SELECT ocr_status, ocr_error
            FROM movie
            WHERE id = ?
            """,
            (movie_id,),
        ).fetchone()

    if row is None:
        return None, "Movie disappeared from the database."

    return row[0], row[1]


def main() -> None:
    args = parse_arguments()

    if args.limit is not None and args.limit <= 0:
        raise ValueError("--limit must be greater than zero.")

    if not args.database.is_file():
        raise FileNotFoundError(
            f"Database file does not exist: {args.database}"
        )

    if not CHECK_OCR_SCRIPT.is_file():
        raise FileNotFoundError(
            f"OCR processing script does not exist: {CHECK_OCR_SCRIPT}"
        )

    if not os.environ.get("GOOGLE_VISION_API_KEY"):
        raise ValueError(
            "GOOGLE_VISION_API_KEY is not set in the environment."
        )

    pending_movies = get_pending_movies(
        database_file=args.database,
        limit=args.limit,
    )

    if not pending_movies:
        print("No movies are waiting for OCR.")
        return

    print(f"Movies waiting for OCR: {len(pending_movies)}")
    print()

    completed = 0
    manual_review = 0
    failed = 0
    unexpected = 0

    for position, (movie_id, name) in enumerate(
        pending_movies,
        start=1,
    ):
        print(
            f"[{position}/{len(pending_movies)}] "
            f"{movie_id} - {name}"
        )

        result = subprocess.run(
            [
                str(CHECK_OCR_SCRIPT),
                movie_id,
                "--database",
                str(args.database),
            ],
            text=True,
            capture_output=True,
        )

        if result.stdout.strip():
            print(result.stdout.strip())

        if result.stderr.strip():
            print(result.stderr.strip(), file=sys.stderr)

        status, error_message = get_ocr_status(
            database_file=args.database,
            movie_id=movie_id,
        )

        if status == "completed":
            completed += 1
            print("Result: completed")
        elif status == "manual_review":
            manual_review += 1
            print("Result: manual_review")
        elif status == "failed":
            failed += 1
            print("Result: failed")
        else:
            unexpected += 1
            print(f"Result: unexpected status {status!r}")

        if error_message:
            print(f"Error: {error_message}")

        print()

    print("OCR batch summary")
    print("-----------------")
    print(f"Selected:       {len(pending_movies)}")
    print(f"Completed:      {completed}")
    print(f"Manual review:  {manual_review}")
    print(f"Failed:         {failed}")
    print(f"Unexpected:     {unexpected}")

    if failed or unexpected:
        sys.exit(1)


if __name__ == "__main__":
    try:
        main()
    except (
        FileNotFoundError,
        ValueError,
        sqlite3.Error,
    ) as error:
        print(f"Error: {error}", file=sys.stderr)
        sys.exit(1)