#!/usr/bin/env python3

import argparse
import math
import os
from datetime import datetime, timezone
from pathlib import Path

try:
    import ffmpeg
except ImportError:
    print("Error: ffmpeg package not found. Please install it separately.")
    print("For installation instructions, visit: https://https://ffmpeg.org")
    ffmpeg = None

def seconds_to_timestamp(seconds):
    """Convert seconds to HH:MM:SS.milliseconds format."""
    return datetime.fromtimestamp(seconds, timezone.utc).strftime('%H:%M:%S.%f')[:-3]


def parse_ffmpeg_metrics_file(file_path):
    """
    Parse a text file generated by the ffmpeg ametadata filter.

    Samples are generated every 100ms. Each sample is represented by two lines:
        frame:0    pts:0       pts_time:0
        lavfi.r128.M=-120.691

    In this example:
        "pts_time" is the timestamp of the sample
        "lavfi.r128.M" is the name of the metric
        "-120.691" is the sample value
    """
    parsed_data = {}
    value4_sum = {}
    value4_count = {}

    with open(file_path, 'r') as file:
        lines = file.readlines()

    for i in range(0, len(lines), 2):
        line1 = lines[i].strip().split()
        line2 = lines[i + 1].strip().split('=')

        # Extract values
        value3 = int(float(line1[2].split(":")[-1]))  # Rounded down seconds
        value4 = float(line2[1])

        # Handle infinite and NaN values
        if not math.isfinite(value4):
            value4 = -120.0

        # Determine the 2-second interval start time
        interval_start = (value3 // 2) * 2

        # Initialize if not present
        if interval_start not in parsed_data:
            parsed_data[interval_start] = []
            value4_sum[interval_start] = 0
            value4_count[interval_start] = 0

        parsed_data[interval_start].append(value4)
        value4_sum[interval_start] += value4
        value4_count[interval_start] += 1

    # Calculate averages
    for key in parsed_data.keys():
        if value4_count[key] != 0:
            average_value4 = value4_sum[key] / value4_count[key]
            parsed_data[key] = round(average_value4, 3) if math.isfinite(average_value4) else -120.0
        else:
            parsed_data[key] = 0  # No data for this interval

    return parsed_data


def create_vtt_file(input_file_path, output_file_path, omp_measurement='', omp_comment=''):
    """Create an OMP 1.0 VTT file from parsed data in temp file."""
    parsed_data = parse_ffmpeg_metrics_file(input_file_path)

    with open(output_file_path, 'w') as output_file:
        # Write the initial lines
        output_file.write("WEBVTT\n\n")

        # If creating an OMP VTT 1.0, write required header
        output_file.write("NOTE\n")
        output_file.write("Omakase Player Web VTT\n")
        output_file.write("V1.0\n\n")

        # Write the parsed data
        for key, value in parsed_data.items():
            start_timestamp = seconds_to_timestamp(key)
            end_timestamp = seconds_to_timestamp(key + 1.999)
            omp_decoration = ""
            omp_decoration += f":MEASUREMENT={omp_measurement}" if omp_measurement else ""
            omp_decoration += f":COMMENT={omp_comment}" if omp_comment else ""

            output_file.write(f"{start_timestamp} --> {end_timestamp}\n")
            output_file.write(f"{value:.3f}{omp_decoration}\n\n")


def generate_audio_metrics(source_audio_file_path, vtt_directory='', samples_per_second=2):
    """
    For the source audio file, use ffmpeg to generate two OMP VTT 1.0 audio analysis VTT files
       EBU R128 metrics using the 'ebur128' audio metric filter
       RMS Level metrics using the 'astats' audio metric filter
    """
    if vtt_directory and not os.path.exists(vtt_directory):
        os.makedirs(vtt_directory)

    # Generate EBU R128 metrics from ametadata filter and convert ffmpeg output to an OMP VTT 1.0 file
    generate_ebur128_metrics(source_audio_file_path, vtt_directory)

    # Generate RMS level metrics with the astats filter and convert ffmpeg output to an OMP VTT 1.0 file
    generate_rms_level_metrics(source_audio_file_path, vtt_directory)


def generate_ebur128_metrics(source_audio_file_path, vtt_directory=''):
    """
    Generate EBU R128 metrics with ffmpeg:
        ffmpeg  -hide_banner \
            -i $file \
            -map 0:a:0 \
            -filter_complex ebur128=metadata=1,ametadata=direct=1:key=lavfi.r128.M:print:file=$metrics_name \
            -f null -
    """

    # If ffmpeg is none, import above failed skip generation
    if ffmpeg is None:
        print("Error: ffmpeg package not available. Skipping EBU R128 metrics generation.")
        return
    
    ffmpeg_metrics_file_path = Path(source_audio_file_path).stem + '.tmp'
    (
        ffmpeg
        .input(source_audio_file_path)
        .filter('ebur128', metadata=1)
        .filter('ametadata', direct=1, key='lavfi.r128.M', mode='print', file=ffmpeg_metrics_file_path)
        .output('-', format='null')
        .run()
    )

    # Convert the ffmpeg output to an OMP VTT 1.0 file
    create_vtt_file(
        ffmpeg_metrics_file_path,
        str(Path(vtt_directory) / f'{Path(source_audio_file_path).stem}_R128_2-SEC.vtt'),
        omp_measurement='lavfi.r128.M',
        omp_comment='2-sec avg'
    )

    # ffmpeg metrics file no longer needed
    os.remove(ffmpeg_metrics_file_path)


def generate_rms_level_metrics(source_audio_file_path, vtt_directory=''):
    """
    Generate RMS level metrics with ffmpeg
        ffmpeg  -hide_banner \
            -i $file \
            -map 0:a:0 \
            -filter_complex astats=metadata=1:reset=1:length=1,ametadata=direct=1:key=lavfi.astats.Overall.RMS_level:print:file=$metrics_name \
            -f null -
    """

    # If ffmpeg is none, import above failed skip generation
    if ffmpeg is None:
        print("Error: ffmpeg package not available. Skipping RMS level metrics generation.")
        return
    
    ffmpeg_metrics_file_path = Path(source_audio_file_path).stem + '.tmp'
    (
        ffmpeg
        .input(source_audio_file_path)
        .filter('astats', metadata=1, reset=1, length=1)
        .filter('ametadata', direct=1, key='lavfi.astats.Overall.RMS_level', mode='print',
                file=ffmpeg_metrics_file_path)
        .output('-', format='null')
        .run()
    )

    # Convert the ffmpeg output to an OMP VTT 1.0 file
    create_vtt_file(
        ffmpeg_metrics_file_path,
        str(Path(vtt_directory) / f'{Path(source_audio_file_path).stem}_RMS_2-SEC.vtt'),
        omp_measurement='lavfi.astats.Overall.RMS_level',
        omp_comment='2-sec avg')

    # ffmpeg metrics file no longer needed
    os.remove(ffmpeg_metrics_file_path)

