from common_code.settings import get_logger, LogColorize
import logging
pfm = LogColorize.score_obj_det_embed
if not ('__file__' in vars() or '__file__' in globals()):
    __file__ = '/home/thebears/Source/pipelines/vision_v3/10_do_obj_det_and_clip/10_score_videos_objdet_and_clip.py'


logger = get_logger(__name__,'/var/log/ml_vision_logs/10_score_videos_objdet_and_clip', stdout=True, systemd=False, level = logging.INFO)

import sys
sys.path.insert(0,'/home/thebears/Source/task_runners/vision_v3')
from cuda_objdet_clip import infer
from common_code import kwq, util, file_names, vector_utils
from common_code.video_meta import FTPVideo
import os
import sys
import time
import json
import traceback


client_suffix = os.environ.get('CONSUMER_CLIENT_SUFFIX','')

# Call logger twice to reclaim
logger = get_logger(__name__,'/var/log/ml_vision_logs/10_score_videos_objdet_and_clip'+client_suffix, stdout=True, systemd=False, level = logging.INFO)
logger.info(pfm(f"Starting objdet_and_embedder scoring @ {__file__}"))


pref = ''
priority_queue = pref + kwq.TOPICS.enter_60_videos_embed_priority 
regular_queue = pref + kwq.TOPICS.exit_00_videos_to_score_detection
backfill_queue = pref + kwq.TOPICS.enter_60_videos_embed_backfill
invalid_file_topic = kwq.TOPICS.exit_60_videos_invalid
failed_embed_topic = kwq.TOPICS.exit_60_embedding_failed
success_embed_topic = kwq.TOPICS.exit_60_videos_embedded


group_id = "vision_v3"
client_id = "objdet_and_embedder"+client_suffix



queue_regular = kwq.create_consumer(group_id = group_id, client_id = client_id, auto_offset_reset = 'latest')
queue_regular.subscribe([regular_queue])


queue_priority = kwq.create_consumer(group_id = group_id, client_id = client_id, auto_offset_reset = 'latest')
queue_priority.subscribe([priority_queue])

queue_backfill = kwq.create_consumer(group_id = group_id, client_id = client_id, auto_offset_reset = 'earliest')
queue_backfill.subscribe([backfill_queue])



logging.debug(f'KAFKA: Created subscribers to {regular_queue} and {priority_queue}')


def check_file_validity(file_to_score):
    if not os.path.exists(file_to_score):
        return False, "file n/a"

    if os.stat(file_to_score).st_size < 100*1000:  #If less than 100kb don't score it
        return False, "file too small"

    if '_reduced' in file_to_score:
        return False, 'file reduced resolution'


    if '_trimmed' in file_to_score:
        return False, 'file is trimmed'

    if not file_to_score.endswith('.mp4'):
        return False, "Not mp4"

    if len(FTPVideo(file_to_score).get_frames_info()) == 0:
        return False, "Failed getting frame info"
    
    return True, ""



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


def perform_loop_kafka():
    score_this = None
    high_priority = False
    result_priority = queue_priority.poll(timeout=0.05)
    msg_value = None
    if result_priority is not None:
        score_this = get_score_file(result_priority)
        msg_value = result_priority.value()
        high_priority = True
        logging.info(f"KAFKA: high priority file {pfm(score_this)}")
        queue_priority.commit()
    else:
        result_regular = queue_regular.poll(timeout=0.05)
        if result_regular is not None:
            msg_value = result_regular.value()            
            score_this = get_score_file(result_regular)
            logging.info(f"KAFKA: regular priority file {pfm(score_this)}")
            queue_regular.commit()
        else:
            result_backfill = queue_backfill.poll(timeout=0.05)
            if result_backfill is not None:
                msg_value = result_backfill.value()            
                score_this = get_score_file(result_backfill)
                logging.info(f"KAFKA: backfill file {pfm(score_this)}")
                queue_backfill.commit()
            
    


    if score_this is None:
        return

    file_to_score = file_names.resolve_file_location(score_this)
    if file_to_score is None:
        return
    if file_to_score != score_this:
        logger.info(f"FILE_PATH_CHANGED: from {pfm(score_this)} to {pfm(file_to_score)}")
        
    file_is_valid, message = check_file_validity(file_to_score)

    if not file_is_valid:
        logger.info(f'INVALID_FILE: Skipping {pfm(file_to_score)}')
        kwq.publish(invalid_file_topic, file_to_score, message)
    else:
        logger.info(f'VALID_FILE: Scoring {pfm(file_to_score)}')
        try:
            det_scores, clip_scores = infer.score_video_cached(file_to_score, return_dict = True)
            vector_utils.upload_vectors_to_db( file_to_score, clip_scores['embeds'], clip_scores['frame_numbers'])
            kwq.publish(success_embed_topic, file_to_score) 
            logger.info(f'SCORING_SUCCESS: {pfm(file_to_score)}')
#            vector_utils.upload_to_doris( file_to_score, clip_scores['embeds'], clip_scores['frame_numbers'])
            logging.info(f'UPLOADED TO DB: {pfm(file_to_score)}')
        except Exception as e:
            traceback.print_exc()
            logging.error(f"SCORING_ERROR: {pfm(file_to_score)}: {e}")
            kwq.publish(failed_embed_topic, file_to_score, str(e))



while True:
    perform_loop_kafka()




