import os
from faster_whisper import WhisperModel
from pydub import AudioSegment

class TranscriptionMaker():
    #書き起こしファイルを吐き出すディレクトリを指定
    def __init__(self,output_dir="/tmp/data/transcriptions"):
        self.model = WhisperModel("base", device="cpu")
        self.output_dir = output_dir
        try:
            if not os.path.exists(self.output_dir):
                os.makedirs(self.output_dir)
        except OSError as e:
            print(f"Error creating directory {self.output_dir}: {e}")
            raise

    #音声ファイルのディレクトリを受け取り、書き起こしファイルを作成する
    def create_transcription(self,segments_directory):
        results = []
        #細切れ音声をくっつける
        audio_directory = self.merge_segments(segments_directory)
        #ディレクトリ内のファイルを全て取得
        if not os.path.isdir(audio_directory):
            raise ValueError(f"The specified path is not a valid directory: {audio_directory}")
        audio_files = os.listdir(audio_directory)
        for audio_file in audio_files:
            if os.path.splitext(audio_file)[-1].lower() != '.wav':
                continue
            audio_path =  os.path.join(audio_directory, audio_file)
            try:
                segments,info = list(self.model.transcribe(audio_path))
            except Exception as e:
                print(f"Error transcripting file {audio_path}: {e}")
                raise
            for segment in segments:
                results.append({
                    "start": segment.start,
                    "end": segment.end,
                    "text": segment.text
                })
        #ファイルの書き込み。ファイル名は"読み込みディレクトリ名_transcription.txt"
        output_file=os.path.join(self.output_dir,os.path.basename(audio_directory)+"_transcription.txt")
        try:
            with open(output_file,"w",encoding="utf-8") as f:
                for result in results:
                    f.write(f"{result['text']}\n")
        except OSError as e:
            print(f"Error writing transcription file: {e}")
            raise
        return output_file
    
    #ファイル名が連続しているならくっつける
    def merge_segments(self,segments_dir,output_dir = "/tmp/data/merged_audio"):
        if not os.path.exists(output_dir):
            os.makedirs(output_dir, exist_ok=True)

        files = sorted([f for f in os.listdir(segments_dir) if f.endswith('.wav')])

        merged_files = []
        current_group = []
        previous_index = None

        for file in files:
            # ファイル名から番号を抽出（例: "0.wav" -> 0）
            file_index = int(file.split('.')[0])
        
            # 番号が連続していない場合、新しいグループを作成
            if previous_index is not None and file_index != previous_index + 1:
                # 現在のグループを結合して保存
                if current_group:
                    merged_files.append(current_group)
                current_group = []
        
            # 現在のファイルをグループに追加
            current_group.append(file)
            previous_index = file_index

        # 最後のグループを追加
        if current_group:
            merged_files.append(current_group)
    
        # グループごとに結合して保存
        for i, group in enumerate(merged_files):
            combined_audio = AudioSegment.empty()
            for file in group:
                file_path = os.path.join(segments_dir, file)
                segment = AudioSegment.from_file(file_path)
                combined_audio += segment
            # 出力ファイル名を設定して保存
            output_file = os.path.join(output_dir, f'merged_{i}.wav')
            combined_audio.export(output_file, format='wav')

        return output_dir