import argparse

from table_parsing.table_parsig import annotate_tables_in_pdf
from box_detection.box_detection import annotate_boxes_in_pdf
from layout_detection.layout_detection import annotate_layout_in_pdf


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("pdf_path")
    parser.add_argument("page_index", type=int)
    parser.add_argument("--object", choices=["table", "box", "layout"], default="table")

    args = parser.parse_args()

    return args


if __name__ == "__main__":
    args = parse_args()
    if args.object == "table":
        annotate_tables_in_pdf(args.pdf_path, page_index=args.page_index)
    elif args.object == "box":
        annotate_boxes_in_pdf(args.pdf_path, page_index=args.page_index)
    elif args.object == "layout":
        annotate_layout_in_pdf(args.pdf_path, page_index=args.page_index)