import os

from faster_whisper import WhisperModel
from tqdm import tqdm

from openlrc.logger import logger
from openlrc.lrc import LRC, LRCOptimizer
from openlrc.prompter import prompter_map
from openlrc.utils import Timer, change_ext, extend_filename


class LRCer:
    def __init__(self, model_name='large-v2'):
        self.model = WhisperModel(model_name, compute_type="float16")

    def __call__(self, audio_path, target_lang='zh-cn', prompter='base_trans'):
        if prompter not in prompter_map:
            raise ValueError(f'Prompter {prompter} not found.')
        prompter = prompter_map[prompter]()  # Initialize appropriate prompter

        transcribed_lrc_path = change_ext(extend_filename(audio_path, '_transcribed'), 'lrc')
        if not os.path.exists(transcribed_lrc_path):
            logger.info(f'Not found transcribed lrc file: {transcribed_lrc_path}')
            with Timer('Transcribing...'):
                # Temporarily workaround for https://github.com/guillaumekln/faster-whisper/issues/71
                # TODO: Should remove `temperature=0` to enhance whisper performance after issue fixed
                segments, info = self.model.transcribe(audio_path, word_timestamps=False, temperature=0)

                logger.info(f'Length of audio {audio_path}: {LRC.format_timestamp(info.duration)}')
                logger.info(f'Detected language: {info.language}, with probability: {info.language_probability}')

                # From generator to list with progress bar shown
                seg_list = []
                for segment in tqdm(segments, 'Transcribing'):
                    seg_list.append(segment)
                logger.debug(f'Transcribed fast-whisper Segments: {seg_list}')

            # Save the transcribed lrc
            self.to_lrc(seg_list, name=transcribed_lrc_path, lang=info.language)  # xxx_transcribed.lrc
        else:
            logger.info(f'Found transcribed lrc file: {transcribed_lrc_path}')

        transcribed_optimized_path = self.post_process(transcribed_lrc_path)  # xxx_transcribed_optimized.lrc

        # Translate the transcribed lrc
        transcribed_optimized_lrc = LRC(transcribed_optimized_path)

        with Timer('Translating...'):
            lrc_name = transcribed_optimized_lrc.translate(
                prompter=prompter, target_lang=target_lang
            )  # xxx_transcribed_optimized_translated.lrc

        self.post_process(lrc_name, output_lrc_name=change_ext(audio_path, 'lrc'),
                          remove_files=[
                              transcribed_optimized_path,  # xxx_transcribed_optimized.lrc
                              lrc_name  # xxx_transcribed_optimized_translated.lrc
                          ])  # xxx.lrc

    @staticmethod
    def to_lrc(segments, name, lang):
        """
        Convert the segments into lrc format.
        """
        with open(name, 'w', encoding='utf-8') as f:
            print(f'LRC generated by https://github.com/zh-plus/Open-Lyrics, lang={lang}', file=f, flush=True)
            for i, segment in enumerate(segments):
                print(
                    f'[{LRC.format_timestamp(segment.start)}] {segment.text}',
                    file=f,
                    flush=True,
                )

                if i != len(segments) - 1 and segment.end != segments[i + 1].start:
                    print(f'[{LRC.format_timestamp(segment.end)}]', file=f, flush=True)

        logger.info(f'File saved to {name}')

    @staticmethod
    def post_process(lrc_name, output_lrc_name=None, remove_files=None):
        lrc_optimizer = LRCOptimizer(lrc_name)
        lrc_optimizer.perform_all()
        optimized_name = lrc_optimizer.save(output_lrc_name=output_lrc_name)

        # Remove intermediate files
        if remove_files:
            for file in remove_files:
                os.remove(file)

        return optimized_name
