import gc
import json

import torch
import whisperx


class Transcriber:

    def __init__(self, language_code="en", device="cuda", compute_type="float16"):
        self.device = device if torch.cuda.is_available() else "cpu"

        # Initialize the transcription model
        self.transcription_model = whisperx.load_model(
            "large-v2", device=self.device, compute_type=compute_type)

        # Initialize the alignment model
        self.alignment_model, self.metadata = whisperx.load_align_model(
            language_code=language_code, device=self.device)
        self.transcript_output = None

    def transcribe(self, wav_filename, batch_size=16):
        # Load the audio file
        audio = whisperx.load_audio(wav_filename)

        # Transcribe the audio file
        result = self.transcription_model.transcribe(audio, batch_size=batch_size)

        # Use the alignment model
        result = whisperx.align(
            result["segments"], self.alignment_model, self.metadata, audio, self.device,
            return_char_alignments=False)

        # Clear GPU memory
        self._clear_memory()

        # Run Diarization
        diarize_model = whisperx.DiarizationPipeline(use_auth_token='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',
                                                     device=self.device)
        diarize_segments = diarize_model(audio)
        result = whisperx.assign_word_speakers(diarize_segments, result)

        # Clear GPU memory again
        self._clear_memory()
        self.transcript_output = result
        return result

    def save_transcript(self, output_json_file, output_txt_file):
        # Save full results to a json file
        with open(output_json_file, 'w') as f:
            json.dump(self.transcript_output, f)

        # Save just the transcripts to a .txt file
        transcripts = [seg['value'] for seg in self.transcript_output['segments']]
        with open(output_txt_file, 'w') as f:
            f.write('\n'.join(transcripts))

    def _clear_memory(self):
        del self.transcription_model
        del self.alignment_model
        del self.metadata
        gc.collect()
        torch.cuda.empty_cache()