import cv2
import os
import yaml
import logging
#from yolov7_auto_annotation.yolov7processor import YoloV7Processor

# Set up logging
logging.basicConfig(filename='app.log', filemode='w', format='%(asctime)s - %(levelname)s - %(message)s')

def convert_video_to_frames(video_path, frame_skip, output_path):
    try:
        # Create a directory to store the frames
        directory = video_path.split('.')[0]
        if output_path:
            directory = os.path.join(output_path, os.path.basename(directory))
        if not os.path.exists(directory):
            os.makedirs(directory)
        logging.info(f"Created directory {directory}")
        # Open the video file
        cap = cv2.VideoCapture(video_path)
        # Get the total number of frames in the video
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        logging.info(f"Opened video file {video_path} with {total_frames} frames")

        # Iterate over the frames
        for i in range(total_frames):
            # Only process every "frame_skip" frame
            if i % frame_skip != 0:
                continue

            # Read the frame
            ret, frame = cap.read()
            if not ret:
                logging.error("Error reading video frame") 
                raise Exception("Error reading video frame")

            # Resize the frame to 416x416
            frame = cv2.resize(frame, (416, 416))

            # Write the frame to a file
            frame_path = os.path.join(directory, f"{i:06d}.jpg")
            cv2.imwrite(frame_path, frame)

        # Release the video file
        cap.release()
        logging.info("Released video file")

        # Rename the directory to include the date and time
        timestamp = os.path.basename(directory)
        new_directory = f"{timestamp}_{frame_skip}"
        if output_path:
            new_directory = os.path.join(output_path, new_directory)
        os.rename(directory, new_directory)
        logging.info(f"Renamed directory {directory} to {new_directory}")
        return new_directory

    except Exception as e:
        logging.error(f"Error: {e}")
        print(f"Error: {e}")
        return None


def rename_files_in_folder(folder_path, string_to_add):
    try:
        for filename in os.listdir(folder_path):
            # check if the file is a file and not a directory
            if os.path.isfile(os.path.join(folder_path, filename)):
                # check if the filename contains the word "classes"
                if "classes" in filename:
                    continue  # skip the file if it contains "classes"
                # get the file extension
                file_ext = os.path.splitext(filename)[1]
                # construct the new filename by adding the string and file extension
                new_filename = filename.replace(file_ext, string_to_add + file_ext)
                # rename the file
                os.rename(os.path.join(folder_path, filename), os.path.join(folder_path, new_filename))
                logging.info(f"Renamed file {filename} to {new_filename}")
        logging.info("Files renamed successfully!")
    except Exception as e:
        logging.error("An error occurred while renaming files:", e)
        print("An error occurred while renaming files:", e)

if __name__ == '__main__':
    # Load configuration from YAML file
    with open('config.yaml', 'r') as f:
        config = yaml.safe_load(f)

    # Extract parameters from configuration   
    video_path = config['video_path']
    frame_skip = config['frame_skip']
    output_path = config.get('output_path')
    string_to_add = config['string_to_add']

    if output_path and not os.path.exists(output_path):
        os.makedirs(output_path)          

    # Convert video to frames
    output_directory = convert_video_to_frames(video_path, frame_skip, output_path)
    logging.info("Videos to frames conversion complete!")


    if output_directory:
        # Rename files in output directory        
        rename_files_in_folder(output_directory, string_to_add)
        logging.info("Files renaming complete!")
