
import os
import sys
import time
from common_code import kwq
from common_code.util import get_cset_match
from common_code.file_names import resolve_file_location
import json
import sys
import logging
from common_code.settings import get_logger, LogColorize
sys.path.insert(0, "/home/thebears/Source/ml/sources_mltraining/model_runner/")
from video_reduce_resolution import convert_file_with_check
import numpy as np

pfm = LogColorize.file_modifications
logger = get_logger(__name__,'/var/log/ml_vision_logs/20_delete_files_and_create_other_videos', stdout=True, systemd=False, level = logging.INFO)

success_det_and_embed_topic = kwq.TOPICS.exit_60_videos_embedded

group_id = "vision_v3"
client_id = "file_converter"

if not ('__file__' in vars() or '__file__' in globals()):
    __file__ = '/home/thebears/Source/pipelines/vision_v3/20_delete_files_and_create_other_videos/20_delete_files_and_create_other_videos.py'


logger.info(pfm(f"Starting file modifier @ {__file__}"))


queue = kwq.create_consumer(group_id = group_id, client_id = client_id, auto_offset_reset = 'earliest')
queue.subscribe([success_det_and_embed_topic])

# %%
def get_ok_to_delete(file_path):
    max_look = 5
    c_path = file_path
    settings_default = {'class_threshold':10000, 'frames_with_dets_threshold': -10000, 'species_ignore': [], 'keep_reduced_mp4': False}

    depth_set = dict()
    for i in range(max_look):
        c_path  = os.path.abspath(os.path.join(c_path, '..'))
        test_path = os.path.join(c_path, 'settings.json')
        if os.path.exists(test_path):
            with open(test_path, 'r') as rr:
                depth_set[i] = json.load(rr)

    depth_set = dict(sorted(depth_set.items(), reverse=True))

    for depth, val in depth_set.items():
        settings_default.update(val)


    settings_default['species_ignore'] = set(settings_default['species_ignore'])
    return settings_default
# %%

def exec_file_remove_logic(cset, dry_run = False):
    did_purge = False
    if isinstance(cset, dict):
        file_path = cset['.mp4']
    else:
        file_path = cset
        cset = get_cset_match(file_path)
    
    settings = get_ok_to_delete(file_path)

    if settings['class_threshold'] > 1 or settings['class_threshold'] <= 0:
        print(f"THRESHOLD SET TO ABOVE 1, SKIPPING: {file_path}")
        return False


    if '.detection.npz' in cset and os.path.exists(cset['.detection.npz']):
        ff = np.load(cset['.detection.npz'], allow_pickle = True)
        det_results = dict(ff)

        del_idx = list()
        for idx, val in det_results['species_label_map'].item().items():
            if val in settings['species_ignore']:
                del_idx.append(idx)

        keep_mask = ~np.isin(det_results['final_labels'], del_idx)

        filt_scores = det_results['final_scores'][keep_mask]
        filt_frames = det_results['final_frames'][keep_mask]

        passes = filt_scores > settings['class_threshold']

        n_frames_above = len(np.unique(filt_frames[passes]))
        n_frames_scored = len(det_results['scored_frames'])

        frac_frames_scored = n_frames_above / n_frames_scored
        logger.info(f'Frac frames scored: {frac_frames_scored}')
        do_keep = frac_frames_scored > settings['frames_with_dets_threshold']


        keep_reduced_mp4 = settings['keep_reduced_mp4']
        if do_keep:
            cpath = os.path.splitext(cset['.mp4'])[0] + '.has_objs'
            with open(cpath,'w') as cc:
                pass
                did_purge = False
        else:
            did_purge = True

            for ff, cf in cset.items():
                if cf.endswith('npz') or cf.endswith('jpg'):
                    pass
                else:
                    if not dry_run and os.path.exists(cf):
                        if '_reduced' in cf:
                            logger.info(f"Skipping reduced_mp4 deletion {pfm(cf)}")
                            continue
                            
                        logger.info(f"Removing {pfm(cf)}")
                    else:
                        print(f'Skipping deleting {ff} because of dry run or path not exist')
    return did_purge

def get_score_file(f):
    return f.key().decode().strip('"')

from common_code.video_meta import FTPVideo


def perform_kafka_loop():
    try:
#    if True:
        result = queue.poll(timeout=0.05)
        if result is None:
            return
        
        queue.commit()
        file_to_score = resolve_file_location(get_score_file(result))
        if file_to_score is None:
            logger.info(f"Not found {pfm(file_to_score)}")
            return
        logger.info(f"Executing file operations {pfm(file_to_score)}")
        did_purge = exec_file_remove_logic(get_cset_match(file_to_score))
        logger.info(f"Did purge? {did_purge}   {get_cset_match(file_to_score)}")
        if not did_purge:
            did_convert = convert_file_with_check(file_to_score)
            FTPVideo(file_to_score).frames_info;
        
        logger.info(f"Finished file operations {pfm(file_to_score)}")        
    except Exception as e:
        logger.error(f"Failed: {pfm(file_to_score)}, {str(e)}") 


while True:
    perform_kafka_loop()



