import os
import json
from datetime import datetime
import subprocess

def cut_video_discard_audio(input_file, output_file, start_frame, end_frame):
    command = [
        "/data/miniconda3/bin/ffmpeg", # change this to your ffmpeg path
        "-i", input_file,
        "-vf", f"select='between(n,{start_frame},{end_frame})',setpts=PTS-STARTPTS",
        "-an",  # Discard audio (for AnimeShooter, we don't need audio)
        "-c:v", "libx264",
        "-crf", "23",
        "-preset", "fast",
        output_file
    ]
    subprocess.run(command)

def timestamp_to_seconds(timestamp):
    if not timestamp or timestamp == "00:00":
        return 0
    
    time = datetime.strptime(timestamp, "%M:%S")
    seconds = time.minute * 60 + time.second
    return seconds

# set root. change it to your own path
video_id = ""
video_path = f"videos/{video_id}.mp4"
annotation_path = f"annotations/{video_id}.json"
output_dir = f"outputs/dataset_samples/{video_id}"
os.makedirs(output_dir, exist_ok=True)

# print all segments available in the video
video_annotation = json.load(open(annotation_path, "r"))
for index, segment in enumerate(video_annotation["segments"]):
    print(f"Segment {index}: {segment['start frame index']} - {segment['end frame index']}")

# each segment is an independent story, we can crop the segment and get the corresponding segment annotation
segment_index = 0 # we take the first segment as an example
story_script = video_annotation["segments"][segment_index]
start_frame = story_script["start frame index"]
end_frame = story_script["end frame index"]
cut_video_discard_audio(video_path, f"{output_dir}/segment_{start_frame}_{end_frame}.mp4", start_frame, end_frame)
json.dump(story_script, open(f"{output_dir}/segment_{start_frame}_{end_frame}.json", "w"), indent=4)

# alternatively, you can avoid cropping the video and calculate the start and end timestamp of the shot in the original video
shot_index = 0 # we take the first shot as an example
fps = video_annotation["fps"]

start_timestamp_in_segment = timestamp_to_seconds(story_script["story script"]["shots"][shot_index]["start time"])
end_timestamp_in_segment = timestamp_to_seconds(story_script["story script"]["shots"][shot_index]["end time"])

start_timestamp_in_video = start_timestamp_in_segment + start_frame / fps
end_timestamp_in_video = end_timestamp_in_segment + start_frame / fps

print(f"\nVideo ID: {video_id}, segment index: {segment_index}, shot index: {shot_index}")
print(f"Start timestamp in segment: {start_timestamp_in_segment} seconds, End timestamp in segment: {end_timestamp_in_segment} seconds")
print(f"Start timestamp in video: {start_timestamp_in_video} seconds, End timestamp in video: {end_timestamp_in_video} seconds")