import cv2
import json
import os

def draw_video(video_path,labels_file,video_save):

    cap = cv2.VideoCapture(video_path)
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    size = (width,height)
    video_basename = os.path.basename(video_path)

    labels_pptsm_list = ["Nodding",'Not_Playing_Mobile','Not_sleet','Playing_Mobile','slppe']

    with open(labels_file, "r") as json_file:
        data_dict = json.load(json_file)
    labels_list = data_dict['big_dict']

    video_save_file = video_save
    os.makedirs(video_save_file, exist_ok=True)
    video_save_path = os.path.join(video_save_file, video_basename)
    videoWriter =cv2.VideoWriter(video_save_path,cv2.VideoWriter_fourcc('X','V','I','D'),fps,size)

    count_fps = 0
    while cap.isOpened():
        success, frame = cap.read()   
        print('count_fps:',count_fps)
        if not success:
            print(video_path,"Ignoring empty camera frame.")
                    # print('video_fps:',video_fps,'count_fps:',count_fps)
            break

        re_anno_list = get_bbox_list(count_fps=count_fps,bbox_dict_list=labels_list)

        for re_dic in re_anno_list:
            
            re_txt = re_dic[0]
            re_bbox = re_dic[1]
          
            cv2.rectangle(frame, (int(re_bbox[0]), int(re_bbox[1])),(int(re_bbox[2]), int(re_bbox[3])), (0, 255, 255), 2)
            cv2.putText(frame, re_txt, (int(re_bbox[0]) - 10, int(re_bbox[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)

            if re_txt != 'person':
                re_txt_1 = labels_pptsm_list[int(re_txt)]
                cv2.putText(frame, re_txt_1, (int(50), (int(50) + int(re_txt)*30)), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 0, 255), 3)
            
        videoWriter.write(frame)

        count_fps += 1

        

    videoWriter.release()
    cap.release()

def get_bbox_list(count_fps,bbox_dict_list):

    re_list = []

    for i,bbox_dict in enumerate(bbox_dict_list):

        startfps = int(bbox_dict['startfps'])
        stopfps = int(bbox_dict['stopfps'])

        if count_fps in range(startfps,stopfps):

            labels = bbox_dict['labels']
            bbox_list = bbox_dict['bbox']

            re_list.append([labels,bbox_list])

    return re_list

if __name__ == '__main__': 

    draw_video(video_path='E:/Bank_files/Bank_02/dataset/video_kf/02.mp4',
               labels_file='E:/Bank_files/Bank_02/process_file/test_video/02.json',
               video_save='E:/Bank_files/Bank_02/process_file/test_video')