import os
from pathlib import Path
from tqdm import tqdm
from faster_whisper import WhisperModel

## %% Configuration du chemin et des répertoires
# Obtenir le répertoire du fichier courant
try:
    script_dir = Path(__file__).parent.parent
except NameError:
    script_dir = Path().absolute()

project_root = script_dir
podcast_dir = os.path.join(project_root, 'import_data', 'data', 'Podcast', 'audio')
output_dir = os.path.join(project_root, 'import_data', 'data', 'Podcast', 'transcripts')

# Créer le répertoire de sortie s'il n'existe pas
os.makedirs(output_dir, exist_ok=True)

## %% Chargement du modèle Faster-Whisper
# Charger le modèle Faster-Whisper
model = WhisperModel("small", device="cpu", compute_type="int8")

## %% Définition des fonctions de transcription et de création de fichiers SRT
def transcribe_audio(audio_path):
    l_segments, _ = model.transcribe(audio_path, language="fr", task="transcribe")
    return list(l_segments)  # Convertir le générateur en liste

def create_srt(l_segments, output_path):
    with open(output_path, 'w', encoding='utf-8') as f:
        for i, segment in enumerate(l_segments, 1):
            start_time = format_time(segment.start)
            end_time = format_time(segment.end)
            text = segment.text.strip()
            f.write(f"{i}\n{start_time} --> {end_time}\n{text}\n\n")

def format_time(seconds):
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    seconds = int(seconds % 60)
    milliseconds = int((seconds % 1) * 1000)
    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"

## %% Traitement de tous les fichiers MP3
# Traiter tous les fichiers MP3
for filename in tqdm(os.listdir(podcast_dir)):
    if filename.endswith(".mp3"):
        mp3_path = os.path.join(podcast_dir, filename)
        srt_path = os.path.join(output_dir, filename.replace(".mp3", ".srt"))

        print(f"Transcription de {filename}...")
        segments = transcribe_audio(mp3_path)
        create_srt(segments, srt_path)
        print(f"Transcription sauvegardée dans {srt_path}")

print("Tous les podcasts ont été transcrits.")