#!/home/thebears/envs/vector_search/bin/python

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



topic_subscribe = 'upload_to_vector_db'
def add_files_to_queue(paths, dry_run = False, force = False, do_upload = True, do_queue = False, use_bulk_import = 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('.oclip_embeds.npz'):
#                if decide_to_put_in_queue(new_path, force = force):
                    queued.add(new_path)

    if do_upload:
        upload_fn = bulk_import_vectors_to_db if use_bulk_import else bulk_insert_vectors_to_db
        results = upload_fn(list(tqdm(queued)))
        for x, ok in results.items():
            if not ok:
                print('Failed: ' + str(x))

    if do_queue:
        for x in tqdm(queued):
            try:
                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)')
    parser.add_argument("--bulk-import", action='store_true', dest='bulk_import',
                         help='EXPERIMENTAL/UNSAFE: use the Milvus server-side Bulk Import job API '
                              'instead of batched inserts. Known to crash this Milvus build when the '
                              'imported segment is loaded -- do not use against production collections.')
    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
    use_bulk_import = args.bulk_import
    if len(paths) == 0:
        paths = settings.dir_watch

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


# %%
# from common_code.vector_utils import upload_vectors_to_db
# ff = '/srv/ftp/hummingbird/2026/07/18/hummingbird_00_20260718144349.oclip_embeds.npz'
# upload_vectors_to_db(ff)
