pdf-script/preflight.py

#!/usr/bin/env python3
"""
preflight.py — Prerequisite checker for the PDF Accessibility Pipeline
=======================================================================

Called at the top of pipeline.py before any pipeline logic runs.
Checks system tools, Python packages, and required environment variables.
For missing config (API key, model path) it prompts interactively and
saves values to .env so the pipeline can continue without a restart.
"""

import importlib.util
import os
import platform
import shutil
import sys

# ---------------------------------------------------------------------------
# Pip package name → Python import name (where they differ)
# ---------------------------------------------------------------------------
_PIP_TO_IMPORT: dict = {
    "opencv-python": "cv2",
    "Pillow": "PIL",
    "python-dotenv": "dotenv",
}

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

_ROOT = os.path.dirname(os.path.abspath(__file__))


def _detect_os() -> str:
    s = platform.system()
    if s == "Darwin":
        return "mac"
    if s == "Linux":
        return "linux"
    return "other"


def _load_dot_env(root: str) -> None:
    env_path = os.path.join(root, ".env")
    if os.path.isfile(env_path):
        try:
            from dotenv import load_dotenv
            load_dotenv(env_path, override=True)
        except ImportError:
            pass  # dotenv not installed yet; will be caught in package check


def _reload_env(root: str) -> None:
    env_path = os.path.join(root, ".env")
    from dotenv import load_dotenv
    load_dotenv(env_path, override=True)


def _check_system_tool(tool: str) -> bool:
    return shutil.which(tool) is not None


def _check_python_packages(req_path: str) -> list:
    """Return list of pip package names that are not importable."""
    if not os.path.isfile(req_path):
        return []
    missing = []
    with open(req_path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            # Strip version specifiers
            pkg = line.split("==")[0].split(">=")[0].split("<=")[0].split("!=")[0].strip()
            import_name = _PIP_TO_IMPORT.get(
                pkg,
                pkg.lower().replace("-", "_"),
            )
            if importlib.util.find_spec(import_name) is None:
                missing.append(pkg)
    return missing


def _write_env_vars(root: str, values: dict) -> None:
    """Upsert key=value pairs into .env (create if absent)."""
    env_path = os.path.join(root, ".env")
    existing_lines = []
    if os.path.isfile(env_path):
        with open(env_path, encoding="utf-8") as f:
            existing_lines = f.readlines()

    for key, val in values.items():
        key_found = False
        for i, line in enumerate(existing_lines):
            stripped = line.split("=")[0].strip() if "=" in line else ""
            if stripped == key:
                existing_lines[i] = f"{key}={val}\n"
                key_found = True
                break
        if not key_found:
            existing_lines.append(f"{key}={val}\n")

    try:
        with open(env_path, "w", encoding="utf-8") as f:
            f.writelines(existing_lines)
    except PermissionError:
        print(f"\n[preflight] ERROR: Cannot write to {env_path} — permission denied.")
        print("            Manually add the following lines to your .env file:")
        for key, val in values.items():
            print(f"            {key}={val}")
        sys.exit(1)


def _prompt_once(prompt: str) -> str:
    """Print prompt and read a line from stdin. Strip whitespace."""
    try:
        return input(prompt).strip()
    except (EOFError, KeyboardInterrupt):
        print()
        sys.exit(1)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def run_preflight_checks(
    *,
    skip: bool = False,
    need_python_part: bool = True,
    need_java_part: bool = True,
) -> None:
    """
    Run all prerequisite checks before the pipeline starts.

    Args:
        skip:             If True, skip all checks (--no-preflight flag).
        need_python_part: If True, check tesseract, poppler, Python packages,
                          OPENAI_API_KEY, and YOLO_MODEL_PATH.
        need_java_part:   If True, check java and mvn.
    """
    if skip:
        print("[preflight] Skipped (--no-preflight).")
        return

    print("[preflight] Checking prerequisites...\n")

    os_type = _detect_os()
    _load_dot_env(_ROOT)

    # ------------------------------------------------------------------
    # Phase 1 — System tools
    # ------------------------------------------------------------------

    # tool_name → (mac install command, linux install command)
    tools_to_check: list[tuple] = []
    if need_python_part:
        tools_to_check += [
            ("tesseract",
             "brew install tesseract",
             "sudo apt install tesseract-ocr"),
            ("pdftoppm",
             "brew install poppler",
             "sudo apt install poppler-utils"),
        ]
    if need_java_part:
        tools_to_check += [
            ("java",
             "brew install openjdk@11  # or download from https://adoptium.net/",
             "sudo apt install default-jdk"),
            ("mvn",
             "brew install maven",
             "sudo apt install maven"),
        ]

    missing_tools: list[tuple] = []
    for tool, mac_cmd, linux_cmd in tools_to_check:
        ok = _check_system_tool(tool)
        status = "OK " if ok else "MISSING"
        label = tool if tool != "pdftoppm" else "pdftoppm (poppler)"
        print(f"  [{status}]  {label}")
        if not ok:
            missing_tools.append((tool, mac_cmd, linux_cmd))

    # ------------------------------------------------------------------
    # Phase 2 — Python packages
    # ------------------------------------------------------------------

    missing_pkgs: list = []
    if need_python_part:
        req_path = os.path.join(_ROOT, "requirements.txt")
        missing_pkgs = _check_python_packages(req_path)
        status = "OK " if not missing_pkgs else "MISSING"
        print(f"  [{status}]  Python packages")

    # ------------------------------------------------------------------
    # Early exit if any system tools or packages are missing
    # ------------------------------------------------------------------

    if missing_tools or missing_pkgs:
        print()

        if missing_tools:
            print("[preflight] Missing required system tools:\n")
            for tool, mac_cmd, linux_cmd in missing_tools:
                label = tool if tool != "pdftoppm" else "pdftoppm (poppler)"
                print(f"  {label}:")
                if os_type == "mac":
                    print(f"    macOS:  {mac_cmd}")
                elif os_type == "linux":
                    print(f"    Linux:  {linux_cmd}")
                else:
                    print(f"    macOS:  {mac_cmd}")
                    print(f"    Linux:  {linux_cmd}")
                print()

        if missing_pkgs:
            print("[preflight] Missing Python packages:\n")
            print(f"  {', '.join(missing_pkgs)}\n")
            print("  Run:  pip install -r requirements.txt\n")

        print("[preflight] Fix the issues above, then re-run pipeline.py.")
        sys.exit(1)

    # ------------------------------------------------------------------
    # Phase 3 — Environment variables
    # ------------------------------------------------------------------

    missing_env: list = []
    if need_python_part:
        api_key = os.environ.get("OPENAI_API_KEY", "").strip()
        yolo_path = os.environ.get("YOLO_MODEL_PATH", "").strip()

        if not api_key:
            print("  [MISSING] OPENAI_API_KEY  (will prompt)")
            missing_env.append("OPENAI_API_KEY")
        else:
            print("  [OK ]  OPENAI_API_KEY")

        if not yolo_path or not os.path.isfile(yolo_path):
            reason = "path not set" if not yolo_path else f"file not found: {yolo_path}"
            print(f"  [MISSING] YOLO_MODEL_PATH  ({reason}, will prompt)")
            missing_env.append("YOLO_MODEL_PATH")
        else:
            print("  [OK ]  YOLO_MODEL_PATH")

    if not missing_env:
        print(f"\n[preflight] All prerequisites satisfied. Starting pipeline...\n")
        return

    print()
    is_interactive = sys.stdin.isatty()

    # Non-interactive (CI/CD): print export hints and exit
    if not is_interactive:
        print("[preflight] Running in non-interactive mode (no TTY).")
        print("            Set the following environment variables and re-run:\n")
        if "OPENAI_API_KEY" in missing_env:
            print('            export OPENAI_API_KEY="sk-..."')
        if "YOLO_MODEL_PATH" in missing_env:
            print('            export YOLO_MODEL_PATH="/absolute/path/to/yolov11x_best.pt"')
        print()
        sys.exit(1)

    # Interactive: prompt for each missing var
    print("[preflight] Some required configuration is missing. Let's set it up.\n")
    collected: dict = {}

    if "OPENAI_API_KEY" in missing_env:
        print("  OPENAI_API_KEY")
        print("  Get your key from: https://platform.openai.com/api-keys\n")
        val = _prompt_once("  Enter your OpenAI API key (sk-...): ")
        if not val:
            val = _prompt_once("  (cannot be empty) Enter your OpenAI API key: ")
        if not val:
            print("\n[preflight] OPENAI_API_KEY is required. Exiting.")
            sys.exit(1)
        collected["OPENAI_API_KEY"] = val
        os.environ["OPENAI_API_KEY"] = val
        print()

    if "YOLO_MODEL_PATH" in missing_env:
        print("  YOLO_MODEL_PATH")
        print("  Download the model weights from:")
        print("    https://github.com/moured/YOLOv11-Document-Layout-Analysis/releases")
        print("    File: yolov11x_best.pt\n")
        path = _prompt_once("  Enter the absolute path to the downloaded .pt file: ")
        if not os.path.isfile(path):
            path = _prompt_once(f"  File not found at '{path}'. Enter path again: ")
        if not os.path.isfile(path):
            print(f"\n[preflight] YOLO model not found at '{path}'. Exiting.")
            sys.exit(1)
        collected["YOLO_MODEL_PATH"] = path
        os.environ["YOLO_MODEL_PATH"] = path
        print()

    # Save to .env and reload
    _write_env_vars(_ROOT, collected)
    _reload_env(_ROOT)
    print("[preflight] Saved to .env.\n")
    print("[preflight] All prerequisites satisfied. Starting pipeline...\n")