from ffprobe import FFProbe
import ffmpeg
import shutil
import json
import math
import os
def execute_trim_video(js_path, start_frame, end_frame, empty_video = False):
    with open(js_path,'r') as jj:
        data = json.load(jj)

    movie_dir = os.path.dirname(js_path)
    archive_dir = os.path.join(movie_dir, 'original') + '/'
                                 
        
    if not os.path.exists(archive_dir):
        os.mkdir(archive_dir)

    rt_name = os.path.splitext(js_path)[0]
    movie_path = rt_name+'.mp4'
    movie_new_path = rt_name+'_trimmed.mp4'
    js_new_path = rt_name + '_trimmed.json'

    frame_pad = 5 #seconds
    vid_info = FFProbe(movie_path)
    vid_stream = vid_info.video[0]

    framerate = vid_stream.framerate
    duration = float(vid_stream.duration)
    max_frames = duration * framerate

    start_frame = max(0, start_frame - frame_pad * framerate)
    end_frame = min( max_frames, end_frame + frame_pad * framerate)

    start_time = start_frame / framerate
    end_time = end_frame / framerate



    to_keep = list()
    for x in data:
        json_frame_num = x['frame_number']
        if isinstance(json_frame_num, str) and json_frame_num == 'thumbnail':
            to_keep.append(x)
        elif json_frame_num >=start_frame and json_frame_num <= end_frame:
            x['frame_number_original'] = x['frame_number']
            x['frame_number'] -= start_frame
            
            to_keep.append(x)


    trim_duration = end_time - start_time

    if os.path.exists(movie_new_path):
        os.remove(movie_new_path)

    if not empty_video:
        cmd = f'ffmpeg -ss {start_time} -i {movie_path} -ss 0 -t {trim_duration} -c copy -map 0 {movie_new_path}'
        return_code = os.system(cmd)

        with open(js_new_path,'w') as jnp:
            json.dump(to_keep, jnp, indent=4)

    if empty_video or return_code == 0:

        shutil.move(movie_path, archive_dir)
        shutil.move(js_path ,archive_dir)

    


        



def trim_video(jspath):
    with open(jspath,'r') as jj:
        data = json.load(jj)

    start_frame = math.inf
    end_frame = -math.inf


    skip_trimming = False
    empty_video = True
    for x in data:
        if len(x['boxes']) > 0:
            json_frame_num = x['frame_number']

            if isinstance(json_frame_num, str) and json_frame_num == 'thumbnail':
                skip_trimming = True
            elif max(x['scores']) > 0.05:
                start_frame = min(start_frame, json_frame_num)
                end_frame = max(end_frame, json_frame_num)
                empty_video = False

        

    if not skip_trimming:
        execute_trim_video(jspath, start_frame, end_frame, empty_video = empty_video)