#  Copyright (C) 2024. Hao Zheng
#  All rights reserved.

import re
from pathlib import Path
from typing import Union

import zhconv

from openlrc.logger import logger
from openlrc.subtitle import Subtitle
from openlrc.utils import extend_filename

# Different language may need different threshold
cut_long_threshold = {
    350: 'en',
    125: 'cn, ja'
}


class SubtitleOptimizer:
    """
    SubtitleOptimizer class is used to optimize subtitles by performing various operations. For example, merging same
    text, merging short duration subtitles, merging repeated patterns, cutting long texts, converting traditional
    Chinese to Mandarin, optimizing punctuation, removing '<unk>' tags, removing empty elements, and stripping
    whitespace.

    Args:
        subtitle (Union[Path, Subtitle]): The subtitle to be optimized.
    """

    def __init__(self, subtitle: Union[Path, Subtitle]):
        if isinstance(subtitle, Path):
            subtitle = Subtitle.from_json(subtitle)

        self.subtitle = subtitle
        self.lang = self.subtitle.lang

    @property
    def filename(self):
        return self.subtitle.filename

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

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

        self.subtitle.segments = new_elements

    def merge_short(self, threshold=1.2):
        """
        Merge the short duration subtitle.
        """
        new_elements = []

        merged_element = None
        for i, element, in enumerate(self.subtitle.segments):
            if i == 0 or element.duration >= threshold:
                if merged_element:
                    if merged_element.duration < 1.5:
                        # If the merged elements is still too small, find the closer element nearby and merge it
                        if element.start - merged_element.end < merged_element.start - new_elements[-1].end:
                            # merge to the next
                            element.text = merged_element.text + element.text
                            element.start = merged_element.start
                        else:
                            # merge to the last
                            new_elements[-1].text += merged_element.text
                            new_elements[-1].end = merged_element.end
                    else:
                        new_elements.append(merged_element)
                    new_elements.append(element)
                    merged_element = None
                else:
                    new_elements.append(element)
            else:
                if not merged_element:
                    merged_element = element
                else:
                    merged_element.text += element.text
                    merged_element.end = element.end

        self.subtitle.segments = new_elements

    def merge_repeat(self):
        """
        Merge the same pattern in one lyric.
        :return:
        """
        new_elements = self.subtitle.segments

        for i, element in enumerate(new_elements):
            text = element.text
            text = re.sub(r'(.)\1{4,}', r'\1\1...', text)
            text = re.sub(r'(.+)\1{4,}', r'\1\1...', text)
            new_elements[i].text = text

        self.subtitle.segments = new_elements

    def cut_long(self, keep=20):
        threshold = 150
        for threshold, lang in cut_long_threshold.items():
            if self.lang.lower() in lang:
                threshold = threshold
                break

        new_elements = self.subtitle.segments

        for i, element in enumerate(new_elements):
            if len(element.text) > threshold and len(element.text) / len(set(element.text)) > 3.0:
                # text length larger than threshold and text density is larger than 3.0
                logger.warning(f'Cut long text: {element.text}\nInto: {element.text[:keep]}...')
                new_elements[i].text = element.text[:keep]

        self.subtitle.segments = new_elements

    def traditional2mandarin(self):
        new_elements = self.subtitle.segments

        for i, element in enumerate(new_elements):
            new_elements[i].text = zhconv.convert(element.text, locale='zh-cn')

        self.subtitle.segments = new_elements

    def punctuation_optimization(self):
        def replace_punctuation_with_chinese(text):
            # Mapping of English punctuation to Chinese punctuation
            punctuation_mapping = {
                ',': '，',
                '.': '。',
                '?': '？',
                '!': '！',
                ':': '：',
                ';': '；',
                '"': '”',
                "'": '’',
                '(': '（',
                ')': '）',
                '[': '【',
                ']': '】',
                '{': '｛',
                '}': '｝'
            }

            # Compile a regular expression pattern for all English punctuation marks
            pattern = re.compile("|".join(map(re.escape, punctuation_mapping.keys())))

            # Replace the punctuation in the text
            result = pattern.sub(lambda match: punctuation_mapping[match.group(0)], text)

            # Replace consecutive (>3) 。 with ...
            result = re.sub(r'(。){3,}', '...', result)

            return result

        new_elements = self.subtitle.segments

        for i, element in enumerate(new_elements):
            new_elements[i].text = replace_punctuation_with_chinese(element.text)

        self.subtitle.segments = new_elements

    def remove_unk(self):
        new_elements = self.subtitle.segments

        for i, element in enumerate(new_elements):
            new_elements[i].text = element.text.replace('<unk>', '')

        self.subtitle.segments = new_elements

    def remove_empty(self):
        self.subtitle.segments = [element for element in self.subtitle.segments if element.text]

    def strip(self):
        new_elements = self.subtitle.segments

        for i, element in enumerate(new_elements):
            new_elements[i].text = element.text.strip()

        self.subtitle.segments = new_elements

    def perform_all(self):
        for _ in range(2):
            self.merge_same()
            self.merge_short()
            self.merge_repeat()
            self.cut_long()
            self.remove_unk()
            self.remove_empty()
            self.strip()

            if self.subtitle.lang.lower() in ['zh-cn', 'zh']:
                self.traditional2mandarin()
            if self.subtitle.lang.lower() in ['zh-cn', 'zh', 'zh-tw']:
                self.punctuation_optimization()

    def save(self, output_name=None, update_name=False):
        optimized_name = extend_filename(self.filename, '_optimized') if not output_name else output_name
        self.subtitle.save(optimized_name, update_name=update_name)
        logger.info(f'Optimized json file saved to {optimized_name}')
