import math
import sys
from dataclasses import dataclass
from typing import List, Union

from langcodes import Language

from openlrc.chatbot import GPTBot
from openlrc.exceptions import SameLanguageException
from openlrc.logger import logger
from openlrc.prompter import BaseTranslatePrompter, format_texts
from openlrc.utils import extend_filename, json2dict


@dataclass
class LRCElement:
    """
    Save a LRC format element.
    """
    start: float
    end: Union[float, None]
    text: str

    @property
    def duration(self):
        if self.end:
            return self.end - self.start
        else:
            return sys.maxsize  # Fake int infinity


class LRC:
    """
    Save a LRC format data.
    """

    def __init__(self, lrc_name=None):
        self.lrc_name = lrc_name
        self.lang = self.get_lang(lrc_name)
        self.elements: List[LRCElement] = []

        if lrc_name:
            self.elements = self.read_lrc(file=lrc_name)

    def read_lrc(self, file):
        """
        Read lrc file and return a list of LRCElement.
        """
        with open(file, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        elements = []

        for line in lines:
            if line.startswith('['):
                start, text = line.split(']', 1)
                start = self.parse_timestamp(start[1:])
                text = text.strip()
                elements.append(LRCElement(start, None, text))

        for i, element in enumerate(elements[:-1]):
            element.end = elements[i + 1].start

        return elements

    def save_lrc(self, file):
        """
        Save lrc file.
        """
        with open(file, 'w', encoding='utf-8') as f:
            print(f'LRC generated by https://github.com/zh-plus/Open-Lyrics, lang={self.lang}', file=f, flush=True)
            for i, element in enumerate(self.elements):
                print(
                    f'[{self.format_timestamp(element.start)}] {element.text}',
                    file=f,
                    flush=True,
                )

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

        return file

    def get_system_prompt(self, src_lang, target_lang, prompter):
        system_prompt = str(prompter).format(
            src_lang=Language.get(self.lang if not src_lang else src_lang).display_name('en'),
            target_lang=Language.get(target_lang).display_name('en')
        )
        # Prevent translating text into Traditional Chinese
        if target_lang == 'zh-cn':
            system_prompt.replace(Language.get(target_lang).display_name('en'), 'Mandarin Chinese')

        return system_prompt

    def translate(self, prompter=BaseTranslatePrompter(), chunk_size=20, src_lang=None, target_lang='zh-cn',
                  intercept_line=None, force_translate=False):
        """
        Use GPT-3.5 to translate lyrics.
        TODO: dynamically adjust the chunk size.
        :param prompter: Translate prompter.
        :param chunk_size: Use smaller chunk size to avoid exceeding the token limit & output complete message.
        :param src_lang: Source language.
        :param target_lang: Target language.
        :param intercept_line: Intercepted lyrics line number.
        :param force_translate: Force translation even if the source language is the same as the target language.
        :return: The translated lrc file path.
        """
        if not force_translate and \
                Language.get(self.lang if not src_lang else src_lang).language_name() \
                == Language.get(target_lang).language_name():
            raise SameLanguageException()

        # Resulted json content
        json_content = {'total_number': 0, 'list': []}

        system_prompt = self.get_system_prompt(src_lang, target_lang, prompter)
        translate_bot = GPTBot(system_prompt=system_prompt)

        lyrics = [element.text for element in self.elements[:intercept_line]]
        # Split lyrics into different chunks
        chunks = [lyrics[i:i + chunk_size] for i in range(0, len(lyrics), chunk_size)]
        user_prompts = [format_texts(chunk) for chunk in chunks]  # Format the chunks into a single string

        logger.info(f'Translating {len(user_prompts)} user_prompts of lyrics with async call.')

        # Start translate
        responses = translate_bot.message(user_prompts)
        for i, response in enumerate(responses):
            content = response.choices[0].message.content
            logger.debug(f'Target content - chunk{i}: {content}')

            chunk_json_content = json2dict(content)
            logger.debug(f'Length of the translated chunk: {len(chunk_json_content["list"])}')

            chunk_size = len(chunks[i])
            # Helping OpenAI clean up their mess.
            if len(chunk_json_content['list']) < chunk_size:
                logger.warning('The number of translated sentences is less than that of the original list. '
                               'Add <MANUALLY-ADDED> label')
                chunk_json_content['list'] += ['<MANUALLY-ADDED>'] * (chunk_size - len(chunk_json_content['list']))
            elif len(chunk_json_content['list']) > chunk_size:
                logger.warning('The number of translated sentences is more than that of the original list. Truncated')
                chunk_json_content['list'] = chunk_json_content['list'][:chunk_size]

            json_content['total_number'] += chunk_json_content['total_number']
            json_content['list'] += chunk_json_content['list']

        # Remove the order number at the front of each sentence
        for i, text in enumerate(json_content['list']):
            json_content['list'][i] = text[text.find('-') + 1:]

        # Replace the original lyrics with the translated lyrics
        for i in range(len(self.elements)):
            self.elements[i].text = json_content['list'][i]
        self.lang = target_lang

        return self.save_lrc(extend_filename(self.lrc_name, '_translated'))

    @staticmethod
    def format_timestamp(seconds: float):
        assert seconds >= 0, "non-negative timestamp expected"
        milliseconds = round(seconds * 1000.0)

        minutes = milliseconds // 60_000
        milliseconds -= minutes * 60_000

        seconds = milliseconds // 1_000
        milliseconds -= seconds * 1_000

        return f"{minutes:02d}:{seconds:02d}.{milliseconds:02d}"

    @staticmethod
    def parse_timestamp(time_stamp):
        minutes, seconds = time_stamp.split(':')
        seconds, milliseconds = seconds.split('.')
        return int(minutes) * 60 + int(seconds) + int(milliseconds) / 1000.0

    @staticmethod
    def get_lang(lrc_name):
        with open(lrc_name, 'r', encoding='utf-8') as f:
            first_line = f.readline()

        if first_line.startswith('LRC generated by'):
            lang = first_line.split('lang=')[-1].strip()
            return lang


class LRCOptimizer:
    def __init__(self, lrc: Union[str, LRC]):
        if isinstance(lrc, str):
            lrc = LRC(lrc)

        self.lrc = lrc

    @property
    def lrc_name(self):
        return self.lrc.lrc_name

    def merge_same_lyrics(self):
        """
        Merge the same text.
        """
        new_elements = []

        for i, element in enumerate(self.lrc.elements):
            if i == 0 or element.text != new_elements[-1].text:
                new_elements.append(element)
            else:
                new_elements[-1].end = element.end

        logger.debug(f'Merge same text: {len(self.lrc.elements)} -> {len(new_elements)}')

        self.lrc.elements = new_elements

    def merge_short_lyrics(self, threshold=2):
        """
        Merge the short text.
        """
        new_elements = []

        for i, element, in enumerate(self.lrc.elements):
            if i == 0 or element.duration >= threshold:
                new_elements.append(element)
            else:
                new_elements[-1].text += ' ' + element.text
                new_elements[-1].end = element.end

        logger.debug(f'Merge short text: {len(self.lrc.elements)} -> {len(new_elements)}')

        self.lrc.elements = new_elements

    def merge_same_words(self):
        """
        Merge the same pattern in one lyric.
        :return:
        """
        elements = self.lrc.elements

        def get_repeat(text):
            """
            Check if the text is repeated for [1-4] words.
            """
            for i in range(1, 5):
                repeating_num = math.floor(len(text) / float(i))
                if text[:i] * repeating_num == text[:i * repeating_num]:
                    return text[:i]
            return None

        for i in range(len(elements)):
            repeat_text = get_repeat(elements[i].text)
            if repeat_text:
                elements[i].text = repeat_text + '...(Repeat)'
                logger.debug(f'Merge same words: {repeat_text}')

        logger.debug('Merge same words done.')

        self.lrc.elements = elements

    def cut_long_lyrics(self, threshold=125, keep=20):
        elements = self.lrc.elements

        for i, element in enumerate(elements):
            if len(element.text) > threshold:
                logger.warning(f'Cut long text: {element.text}\nInto: {element.text[:keep]}...')
                elements[i].text = element.text[:keep] + f'(Cut to {keep})'

        logger.debug('Cut long text done.')

        self.lrc.elements = elements

    def perform_all(self):
        for _ in range(2):
            self.merge_same_lyrics()
            self.merge_short_lyrics()
            self.merge_same_words()
            self.cut_long_lyrics()

    def save(self, output_lrc_name=None):
        optimized_name = extend_filename(self.lrc_name, '_optimized') if not output_lrc_name else output_lrc_name
        self.lrc.save_lrc(optimized_name)
        logger.info(f'Optimized LRC file saved to {optimized_name}')
        return optimized_name
