
from CommonCode import kwq
import time
import json
import logging
import os
from CommonCode.settings import get_logger, LogColorize
from kafka import TopicPartition
from kafka.structs import OffsetAndMetadata
pfm = LogColorize.score_obj_det_orin 
logger = get_logger(__name__,'/var/log/ml_vision_logs/01_score_obj_det_orin', stdout=True, systemd=False)

os.system("sudo /usr/bin/systemctl restart --now systemd-journal-upload.service")

logger.info(pfm(f"Starting wait_for_new_messages.py on orin for scoring object detection"))
input_topic = kwq.TOPICS.videos_to_score_detection
producer = kwq.producer
topic_produce = kwq.TOPICS.videos_scored_detection

client_id = 'obj_detector_orin_3'
group_id = client_id
# %%


import json


logger.debug("Starting Kafka Consumer")

from deepstream_obj_det import run_inference, target_width, target_height
import os

os.environ.pop("DISPLAY",None)

def run_inference_for_file(file_path):
    start_time = time.time()

    end_time = time.time()

    pre_path, _ = os.path.splitext(file_path)

    det_path = pre_path + '.json.orin'
    emb_path = pre_path + '.oclip.orin'
    if os.path.exists(det_path) and os.path.exists(emb_path):
        return "Already scored"

    if not os.path.exists(file_path):
        return "Movie does not exist"

# %%

    cmd = f'/usr/bin/ffprobe -v error -select_streams v:0 -count_packets -show_entries stream=nb_read_packets -of csv=p=0 {file_path}'#/srv/ftp/railing/2025/02/28/railing_00_20250228115800.mp4
    try:
        frames = int(os.popen(cmd).read().strip())
    except:
        frames = 0

    logger.info(f"TOTAL_FRAMES: {frames}")
    if frames < 30:
        logger.info(f"TOTAL_FRAMES_SKIPPING: {file_path}")
        return "FAILED, NOT ENOUGH FRAMES"

    while True:
        try:
            with open(det_path,'w') as ff:
                out = ff.write(' '*100)
        except OSError as e:
            logger.error(f"NO_SPACE :{det_path}:{e}")                
        else:
            break
        
        time.sleep(5)

    try:
        os.remove(det_path)
    except:
        pass
    logger.info(f"ENOUGH SPACE, STARTING INFERENCE")
    detector_results, embedder_results = run_inference(file_path)
    obj_det_dict = dict()
    obj_det_dict['meta'] = {'model_version':'orin_v1'}
    obj_det_dict['scoring'] = {'start_time':start_time, 'end_time': end_time}
    obj_det_dict['json'] ={'path':det_path}
    obj_det_dict['video'] = {'path':file_path, 'target_w': target_width, 'target_h': target_height}




    by_frame_num = dict()
    for idx, sc in enumerate(detector_results):
        c_res = dict()
        c_frame = sc['frame_number']
        c_res['score'] = sc['score']
        c_res['L'] = sc['left']
        c_res['T'] = sc['top']
        c_res['W'] = sc['width']
        c_res['H'] = sc['height']
        c_res['name'] = sc['label']
        c_res['idx'] = sc['class_id']
        if c_frame not in by_frame_num:
            by_frame_num[c_frame] = list()

        by_frame_num[c_frame].append(c_res)

    obj_det_dict['scores'] = [{'frame':key, 'detections':val} for key,val in by_frame_num.items()]


    with open(det_path,'w') as ff:
        json.dump(obj_det_dict, ff, indent=4)


    emb_dict = dict()
    emb_dict['meta'] = {'model_version':'ViT-L-16-SigLIP2-512','host':'orin'}
    emb_dict['scoring'] = {'start_time':start_time, 'end_time': end_time}
    emb_dict['json'] ={'path':det_path}
    emb_dict['video'] = {'path':file_path, 'target_w': target_width, 'target_h': target_height}
    emb_dict['scores'] = list()


    for c_score in embedder_results:
        fr_num = c_score['frame_number']
        vect = c_score['vector']
        emb_dict['scores'].append({'score':vect, 'frame':fr_num})

    with open(emb_path,'w') as ff:
        json.dump(emb_dict, ff, indent=4)

    return "Success"

# %%
consumer = kwq.create_consumer(input_topic, group_id = group_id, client_id = client_id)
#consumer.subscribe(input_topic)

c_part = TopicPartition(input_topic, 0)
consumer.assign([c_part])


c_committed = consumer.committed(c_part)
logger.info(f"KAFKA_POSITION_IS: {str(consumer.position(c_part))}")

if c_committed is None:
    logger.info(f"KAFKA_POSITION_NOT_COMMITTED")
else:
    logger.info(f"KAFKA_POSITION_COMMITTED_IS: {c_committed}")
    consumer.seek(c_part, c_committed)
logger.info("START POLLING")

#while True:
#    out = consumer.poll(timeout_ms=5000 , update_offsets = False)
#    msgs = list()
#    logger.info(f"KAFKA_POSITION_COMMITTED_IS: {str(consumer.committed(c_part))}")    
#    logger.info(f"KAFKA_POSITION_IS: {str(consumer.position(c_part))}")
#    
#    for k, v in out.items():
#        msgs.extend(v)
#    for message in msgs:
for message in consumer:
        logger.info(f"KAFKA_POSITION_COMMITTED_IS: {str(consumer.committed(c_part))}")    
        logger.info(f"KAFKA_POSITION_IS: {str(consumer.position(c_part))}")
        
        logger.info(f"MSG_RECEIVED :{message}")        
        logger.info(f"INFERENCE_START: {pfm(message.key)}")
        result = run_inference_for_file(message.key)
        logger.info(f"INFERENCE_DONE:{pfm(result)} {message.key}")
        oandm = OffsetAndMetadata(message.offset,'')
        consumer.commit({c_part:oandm})
        producer.send(topic_produce, value=message.value, key=message.key)   

