main.py

"""
MetaBoKG end-to-end pipeline.

Six stages can be invoked individually, chained, or all together:

  1. fetch   - Download citing PMIDs and their PDFs        [data_retriever/pmid.py]
  2. extract - Convert PDFs to markdown, find GNPS task    [data_retriever/find_massive_gnps.py]
               IDs / MassIVE / Zenodo accessions
  3. jobs    - Download GNPS job archives                  [data_retriever/job_download.py]
  4. map     - Generate TTL knowledge graphs from jobs     [mapping/script.py]
  5. load    - Bulk-load TTLs into Virtuoso                [mapping/load_to_virtuoso.py]
  6. cq      - Run competency queries                      [mapping/load_and_query_kg.py]

Examples:

    # full pipeline
    python main.py all

    # one stage, extra args forwarded to the underlying script
    python main.py fetch
    python main.py extract --workers 8
    python main.py jobs --gnps-version 1 --max-workers 8
    python main.py map --source both --rewrite
    python main.py load --reload
    python main.py cq --only CQ1 CQ3

    # chain a subset (in declared order); extras are forwarded to every stage
    python main.py map load cq
"""

from __future__ import annotations

import argparse
import sys
from typing import Callable


def _run_argv_main(module_main: Callable, stage: str, extra_args: list[str]) -> int:
    """Invoke an argparse-based ``main()`` with ``extra_args`` spliced into sys.argv."""
    saved = sys.argv
    sys.argv = [f"main.py:{stage}", *extra_args]
    try:
        rc = module_main()
        return rc if isinstance(rc, int) else 0
    finally:
        sys.argv = saved


def stage_fetch(extra_args: list[str]) -> int:
    if extra_args:
        print(
            f"[warn] 'fetch' takes no CLI args; ignoring {extra_args}. "
            "Edit SOURCE_PMID in data_retriever/pmid.py to change the seed PMID.",
            file=sys.stderr,
        )
    from data_retriever import pmid as mod
    mod.main()
    return 0


def stage_extract(extra_args: list[str]) -> int:
    from data_retriever import find_massive_gnps as mod
    return _run_argv_main(mod.main, "extract", extra_args)


def stage_jobs(extra_args: list[str]) -> int:
    from data_retriever import job_download as mod
    return _run_argv_main(mod.main, "jobs", extra_args)


def stage_map(extra_args: list[str]) -> int:
    from mapping import script as mod
    return _run_argv_main(mod.main, "map", extra_args)


def stage_load(extra_args: list[str]) -> int:
    from mapping import load_to_virtuoso as mod
    return _run_argv_main(mod.main, "load", extra_args)


def stage_cq(extra_args: list[str]) -> int:
    from mapping import load_and_query_kg as mod
    return _run_argv_main(mod.main, "cq", extra_args)


STAGES: dict[str, Callable[[list[str]], int]] = {
    "fetch":   stage_fetch,
    "extract": stage_extract,
    "jobs":    stage_jobs,
    "map":     stage_map,
    "load":    stage_load,
    "cq":      stage_cq,
}
STAGE_ORDER = ["fetch", "extract", "jobs", "map", "load", "cq"]


def parse_args() -> tuple[list[str], list[str]]:
    parser = argparse.ArgumentParser(
        prog="metabokg",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=__doc__,
    )
    parser.add_argument(
        "stages",
        nargs="+",
        metavar="STAGE",
        help=f"one or more of: {', '.join(STAGE_ORDER)} | all",
    )
    parser.add_argument(
        "rest",
        nargs=argparse.REMAINDER,
        help="extra arguments forwarded to the underlying script(s).",
    )
    args = parser.parse_args()

    extra = args.rest[1:] if args.rest and args.rest[0] == "--" else args.rest

    resolved: list[str] = []
    for s in args.stages:
        if s == "all":
            resolved.extend(STAGE_ORDER)
        elif s in STAGES:
            resolved.append(s)
        else:
            parser.error(
                f"unknown stage '{s}'. choose from {STAGE_ORDER + ['all']}"
            )
    return resolved, extra


def main() -> int:
    stages, extra = parse_args()
    if len(stages) > 1 and extra:
        print(
            f"[info] forwarding extra args {extra} to every chained stage",
            file=sys.stderr,
        )

    for stage in stages:
        bar = "=" * 70
        print(f"\n{bar}\n[{stage}] starting\n{bar}")
        rc = STAGES[stage](extra)
        if rc != 0:
            print(f"[{stage}] failed (exit {rc})", file=sys.stderr)
            return rc
        print(f"[{stage}] done")
    return 0


if __name__ == "__main__":
    sys.exit(main())