import sys
import argparse
from . import chip, index, search


def main():
    parser = argparse.ArgumentParser(description="Earth Search CLI")
    subparsers = parser.add_subparsers(dest="command")

    chip_parser = subparsers.add_parser("chip")
    index_parser = subparsers.add_parser("index")
    search_parser= subparsers.add_parser("search")

    chip_parser.add_argument("--image_dir", type=str, required=True, help="Directory path of images to chip")
    chip_parser.add_argument("--chip_dir", type=str, required=True, help="Directory path to write chips to")
    chip_parser.add_argument("--window_size", type=int, required=False, default=512, help="Size of sliding window, e.g., 512")
    chip_parser.add_argument("--stride", type=float, required=False, default=0.0, help="Amount of overlap in x, y direction, e.g., 0.2")
    chip_parser.add_argument("--valid_exts", nargs="*", required=False, help="Image extensions to filter for")
    chip_parser.add_argument("--multiprocess", action="store_true", help="Use multiprocessing vs multithreading")

    index_parser.add_argument("--image_dir", type=str, required=True, help="Directory path of chipped images to index")
    index_parser.add_argument("--index_path", type=str, required=True, help="Path to save index")
    index_parser.add_argument("--indexed_images_path", type=str, required=True, help="Path to save file containing indexed image paths")
    index_parser.add_argument("--index_type", type=str, required=True, default="L2", help="Type of index to use (currently only using L2)")
    index_parser.add_argument("--model_type", type=str, required=True, default="dinov2_vits14_reg",help="Type of model to use")
    index_parser.add_argument("--device", type=str, required=False, default="cuda", help="Device to use for feature extraction (cuda, mps, cpu)")
    index_parser.add_argument("--batch_size", type=int, required=False, default=32, help="Number of images per batch")
    index_parser.add_argument("--overwrite_index", action="store_true", help="Overwrite existing index of same index_path")

    search_parser.add_argument("--query_image", type=str, required=True, help="Path to directory of images to index")
    search_parser.add_argument("--index_path", type=str, required=True, help="Path to saved index")
    search_parser.add_argument("--indexed_images_path", type=str, required=True, help="Path to file containing indexed image paths")
    search_parser.add_argument("--index_type", type=str, required=True, default="L2", help="Type of index to use (currently only using L2)")
    search_parser.add_argument("--top_k", type=int, required=True, default=10, help="top k-Nearest Neighbors to return")
    search_parser.add_argument("--model_type", type=str, required=True, default="dinov2_vits14_reg", help="Type of model to use")
    search_parser.add_argument("--device", type=str, required=False, default="cuda", help="Device to use for feature extraction (cuda, mps, cpu)")

    chip_parser.set_defaults(func=chip.chip)
    index_parser.set_defaults(func=index.index)
    search_parser.set_defaults(func=search.search)

    args = parser.parse_args()

    if args.command:
        func_args = {
            k: v for k, v in vars(args).items() if k not in ("command", "func")
        }
        try:
            args.func(**func_args)
        except NotImplementedError as e:
            print(e)
            sys.exit(-1)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()
