import os
import subprocess
import csv
import shutil
import unicodedata
import re
import torch
import torchaudio
import librosa
import numpy as np
from webvtt import WebVTT

def normalize_czech(text):
    # Normalize Unicode characters
    text = unicodedata.normalize('NFKD', text)
    # Remove non-ASCII characters
    text = text.encode('ASCII', 'ignore').decode('ASCII')
    # Remove spaces
    text = text.replace(' ', '')
    return text

def convert_to_seconds(timestamp):
    # Convert a VTT timestamp (hh:mm:ss.ms) to total seconds
    hours, minutes, seconds = timestamp.split(':')
    seconds_parts = seconds.split('.')
    if len(seconds_parts) == 2:
        seconds, milliseconds = seconds_parts
        total_seconds = int(hours) * 3600 + int(minutes) * 60 + int(seconds) + float('0.' + milliseconds)
    else:
        seconds = seconds_parts[0]
        total_seconds = int(hours) * 3600 + int(minutes) * 60 + int(seconds)
    return total_seconds

def cut_audio(wav_path, vtt_path, output_dir, csv_writer, speaker_name, silence_method, silero_params):
    # Ensure the output directory exists
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)  # Delete the existing folder
    os.makedirs(output_dir)

    # Load the VTT file
    vtt = WebVTT().read(vtt_path)
    # Loop through each caption in the VTT file and extract timestamps
    for i, caption in enumerate(vtt):
        start = caption.start
        end = caption.end
        text = caption.text.strip()
        # Convert time format from VTT (hh:mm:ss.ms) to ffmpeg format (seconds)
        start_sec = convert_to_seconds(start)
        end_sec = convert_to_seconds(end)
        duration = end_sec - start_sec
        # Skip if the audio duration is not within the desired range (3.2 to 11.55 seconds)
        if duration < 3.2 or duration > 11.55:
            continue
        # Define the output path for the audio segment
        output_file = os.path.join(output_dir, f'{os.path.basename(output_dir)}{i+1}.wav')
        # Use ffmpeg to cut the audio to a temporary file
        temp_output_file = os.path.join(output_dir, f'temp{os.path.basename(output_dir)}_{i+1}.wav')
        cmd_cut = [
            'ffmpeg', '-i', wav_path,
            '-ss', str(start_sec), '-to', str(end_sec),
            '-ar', '22050',  # Resample to 22050 Hz
            '-ac', '1',  # Mono audio
            '-c', 'pcm_s16le',
            temp_output_file
        ]
        # Run the command to cut the initial segment
        subprocess.run(cmd_cut, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
        # Process the audio segment based on selected method
        success = process_audio_segment(temp_output_file, output_file, silence_method, silero_params)
        if success:
            # Write the entry to the CSV file
            csv_writer.writerow([output_file, text, speaker_name])
        # Remove the temporary file
        if os.path.exists(temp_output_file):
            os.remove(temp_output_file)

def process_audio_segment(input_file, output_file, silence_method, silero_params):
    if silence_method == 'silero':
        # Use Silero VAD
        return process_with_silero(input_file, output_file, silero_params)
    elif silence_method == 'energy':
        # Use energy-based method
        return process_with_energy(input_file, output_file)
    else:
        # No silence removal; apply loudness normalization only
        return process_without_silence_removal(input_file, output_file)

def process_with_silero(input_file, output_file, silero_params):
    # Load the audio segment
    wav, sample_rate = torchaudio.load(input_file)
    # Ensure the sample rate is 16000 Hz for Silero VAD
    if sample_rate != 16000:
        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
        wav = resampler(wav)
        sample_rate = 16000
    # Load the Silero VAD model
    model, utils = torch.hub.load(
        repo_or_dir='snakers4/silero-vad',
        model='silero_vad',
        force_reload=False,
        onnx=False
    )
    (get_speech_timestamps, _, _, _, _) = utils
    # Run VAD with custom parameters
    speech_timestamps = get_speech_timestamps(
        wav.squeeze(), 
        model, 
        sampling_rate=sample_rate,
        **silero_params
    )
    if speech_timestamps:
        # Get the start and end times of speech in seconds
        start_idx = speech_timestamps[0]['start']
        end_idx = speech_timestamps[-1]['end']
        adjusted_start_sec = start_idx / sample_rate
        adjusted_end_sec = end_idx / sample_rate
        # Use ffmpeg to extract the adjusted segment, apply loudness normalization
        cmd_process = [
            'ffmpeg', '-i', input_file,
            '-ss', str(adjusted_start_sec), '-to', str(adjusted_end_sec),
            '-af', 'loudnorm=I=-14:LRA=11:TP=-1.5',
            output_file
        ]
        subprocess.run(cmd_process, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
        return True
    else:
        # No speech detected
        return False

def process_with_energy(input_file, output_file):
    # Load the audio
    y, sr = librosa.load(input_file, sr=22050)
    # Compute short-time energy
    frame_length = 2048
    hop_length = 512
    energy = np.array([
        sum(abs(y[i:i+frame_length]**2))
        for i in range(0, len(y), hop_length)
    ])
    # Normalize energy
    energy = energy / np.max(energy)
    # Identify silent frames
    threshold = 0.01  # Energy threshold (adjustable)
    frames = np.where(energy > threshold)[0]
    if frames.size > 0:
        # Compute time positions
        start_frame = frames[0]
        end_frame = frames[-1]
        adjusted_start_sec = start_frame * hop_length / sr
        adjusted_end_sec = min(len(y), (end_frame * hop_length + frame_length)) / sr
        # Add padding to avoid cutting off word endings/beginnings
        padding = 0.2  # 200 ms padding (adjustable)
        adjusted_start_sec = max(0, adjusted_start_sec - padding)
        adjusted_end_sec = min(len(y) / sr, adjusted_end_sec + padding)

        # Use ffmpeg to extract the adjusted segment, apply loudness normalization
        cmd_process = [
            'ffmpeg', '-i', input_file,
            '-ss', str(adjusted_start_sec), '-to', str(adjusted_end_sec),
            '-af', 'loudnorm=I=-14:LRA=11:TP=-1.5',
            output_file
        ]
        subprocess.run(cmd_process, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
        return True
    else:
        # No speech detected
        return False

def process_without_silence_removal(input_file, output_file):
    # Apply loudness normalization only
    cmd_process = [
        'ffmpeg', '-i', input_file,
        '-af', 'loudnorm=I=-14:LRA=11:TP=-1.5',
        output_file
    ]
    subprocess.run(cmd_process, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
    return True

def process_directory(directory, csv_file_path, silence_method='silero', silero_params=None):
    # Set default Silero parameters if not provided
    if silero_params is None:
        silero_params = {
            'threshold': 0.5,
            'min_speech_duration_ms': 250,
            'min_silence_duration_ms': 100,
            'window_size_samples': 1536,
            'speech_pad_ms': 30
        }
    # Delete the old CSV file if it exists
    if os.path.exists(csv_file_path):
        os.remove(csv_file_path)
    # Open the CSV file for writing
    with open(csv_file_path, mode='w', newline='', encoding='utf-8') as csv_file:
        csv_writer = csv.writer(csv_file, delimiter='|')
        # Write the header row
        csv_writer.writerow(['audio_file', 'text', 'speaker_name'])
        # Loop through all .wav files in the directory
        for file in os.listdir(directory):
            if file.endswith('.wav'):
                wav_path = os.path.join(directory, file)
                vtt_path = os.path.join(directory, file.replace('.wav', '.vtt'))
                # Check if the corresponding VTT file exists
                if os.path.exists(vtt_path):
                    # Normalize the folder name
                    normalized_name = normalize_czech(file.replace('.wav', ''))
                    output_dir = os.path.join(directory, normalized_name)
                    speaker_name = 'Random'  # Use 'Random' as speaker_name
                    cut_audio(wav_path, vtt_path, output_dir, csv_writer, speaker_name, silence_method, silero_params)

if __name__ == "__main__":
    # Specify the directory containing the .wav and .vtt files
    directory = '.'
    # Specify the output CSV file path
    csv_file_path = os.path.join(directory, 'audio_dataset.csv')
    # Set the silence removal method ('silero', 'energy', 'none')
    silence_method = 'silero'  # Change this to 'energy' or 'silero' or 'none' as needed

    # Silero VAD parameters (Adjust as needed)
    silero_params = {
        'threshold': 0.35,
        'min_speech_duration_ms': 50,
        'min_silence_duration_ms': 200,
        'window_size_samples': 512,
        'speech_pad_ms': 30
    }

    # Process the directory
    process_directory(directory, csv_file_path, silence_method, silero_params) do the audio length check 3.2 and 11.55 seconds only after the silence cutting