import os

from openlrc.logger import logger
from openlrc.lrc import LRC, LRCOptimizer
from openlrc.prompter import prompter_map
from openlrc.transcribe import Transcriber
from openlrc.utils import Timer, change_ext, extend_filename, get_audio_duration, format_timestamp


class LRCer:
    def __init__(self, model_name='large-v2'):
        self.transcriber = Transcriber(model_name=model_name)

    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('Transcription process'):
                logger.info(f'Audio length: {audio_path}: {get_audio_duration(audio_path)}')
                segments, info = self.transcriber.transcribe(audio_path, batch_size=2)
                logger.info(f'Detected language: {info.language}')

                # From generator to list with progress bar shown
                seg_list = segments['sentences']  # [{'text': ..., 'start_word': ..., 'end_word':...}, ...]
                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'), t2m=target_lang == 'zh-cn',
                          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'[{format_timestamp(segment["start_word"]["start"])}] {segment["text"]}',
                    file=f,
                    flush=True,
                )

                print(f'[{format_timestamp(segment["end_word"]["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, t2m=False):
        lrc_optimizer = LRCOptimizer(lrc_name)
        lrc_optimizer.perform_all(t2m=t2m)
        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
