from common_code.vector_utils import upload_vectors_to_db
from common_code import kwq, util, file_names, vector_utils
from multiprocessing import Pool
from tqdm import tqdm
import argparse
import os
import logging
logger = logging.getLogger(__name__)



topic_subscribe = 'upload_to_ch_db'
def add_files_to_queue(paths, dry_run = False, force = False, do_upload = True, do_queue = False):
    queued = set()
    for rt in paths:
        for root, dirs, files in os.walk(rt):
            for f in files:
                new_path = os.path.join(root, f)
                if new_path.endswith('detection.npz') and 'leopards' not in new_path.lower():
#                if decide_to_put_in_queue(new_path, force = force):
                    queued.add(new_path)

    for x in tqdm(queued):
        try:
            if do_upload:
                upload_vectors_to_db(x)
            if do_queue:
                kwq.publish(topic_subscribe, x, x)

        except Exception as e:
            print('Failed: '+str(x)+'with '+str(e))

        



if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="Watch and Fix Permissions And Push to Kafka Queue"
    )

    parser.add_argument("paths", nargs="*", help="Paths to monitor", default=())
    parser.add_argument("--dry-run", action='store_true', help='Dry Run')
    parser.add_argument("--force", action='store_true', help='Force adding (do not check if already done)')
    parser.add_argument("--skip_upload", action='store_true', help='Force adding (do not check if already done)')            
    parser.add_argument("--do_queue", action='store_true', help='Force adding (do not check if already done)')        
    args, _ = parser.parse_known_args()
    paths = args.paths
    force = args.force
    dry_run = args.dry_run
    do_queue = args.do_queue
    do_upload = not args.skip_upload
    if len(paths) == 0:
        paths = default_dir_watch

    add_files_to_queue(paths, dry_run = dry_run, force=force, do_upload = do_upload, do_queue = do_queue)

