import re
import regex
import inflect
import unicodedata
from lingua import Language, LanguageDetectorBuilder
from builtins import str as unicode

try:
    from tn.chinese.normalizer import Normalizer as ZhNormalizer
    from tn.english.normalizer import Normalizer as EnNormalizer
except:
    print("window")

from fireredtts.modules.text_normalizer.regex_common import *
from fireredtts.modules.text_normalizer.utils import *


def preprocess_text(sentence):
    # preprocessing
    sentence = bytes(sentence, "utf-8").decode("utf-8", "ignore")
    sentence = regex.sub("[\p{Cf}--[\u200d]]", "", sentence, flags=regex.V1)
    sentence = regex.sub("\p{Co}", "", sentence)
    sentence = sentence.replace("\u00a0", " ")
    sentence = sentence.replace("\ufffd", "")
    sentence = regex.sub("\p{Zl}", "\n", sentence)
    sentence = regex.sub("\p{Zp}", "\n", sentence)

    sentence = unicode(sentence)
    sentence = "".join(
        char
        for char in unicodedata.normalize("NFD", sentence)
        if unicodedata.category(char) != "Mn"
    )  # Strip accents

    sentence = strip_kaomoji(sentence)
    # full to half with exemption (to be converted after number TN): 。，：
    sentence = f2b(sentence, exemption="。，：")

    # clean spaces
    sentence = sentence.replace("\n", "，")
    sentence = sentence.replace("\t", "，")
    sentence = sentence.replace("\r", "，")
    sentence = re.sub(r"[。.]{3,}", "…", sentence)
    sentence = re.sub(r"[…⋯]{1,}", "…", sentence)
    sentence = re.sub(r"[ ]+", " ", sentence)
    sentence = sentence.strip()

    # punctuation reduction
    result = ""
    for idx, char in enumerate(sentence):
        if char in symbol_reduction:
            char = symbol_reduction[char]

        if char == " ":
            if idx == 0:
                continue
            if is_chinese(sentence[idx + 1]) and (
                is_chinese(sentence[idx - 1]) or sentence[idx - 1] in '") '
            ):
                result += "，"
            else:
                result += " "
            continue

        if is_valid_char(char):
            result += char
    result = re.sub(r"[ ]+", " ", result)
    return result


def rettt(sentence):
    # handle abbreviations for all languages
    sentence = sentence.replace("&nd", "and")
    sentence = sentence.replace("Jan.", "january")
    sentence = sentence.replace("Feb.", "febrary")
    sentence = sentence.replace("Mar.", "march")
    sentence = sentence.replace("Apr.", "april")
    sentence = sentence.replace("May.", "may")
    sentence = sentence.replace("Jun.", "june")
    sentence = sentence.replace("Jul.", "july")
    sentence = sentence.replace("Aug.", "august")
    sentence = sentence.replace("Sept.", "september")
    sentence = sentence.replace("Sep.", "september")
    sentence = sentence.replace("Oct.", "october")
    sentence = sentence.replace("Nov.", "november")
    sentence = sentence.replace("Dec.", "december")
    sentence = sentence.replace("Mon.", "monday")
    sentence = sentence.replace("Tues.", "tuesday")
    sentence = sentence.replace("Wed.", "wednesday")
    sentence = sentence.replace("Thur.", "thursday")
    sentence = sentence.replace("Fri.", "friday")
    sentence = sentence.replace("Sat.", "saturday")
    if sentence != "Sun.":
        sentence = sentence.replace("Sun.", "sunday")
    sentence = re.sub(r" St\. ([A-Z])", r" saint \1", sentence)
    sentence = re.sub(r" St\.", " street", sentence)
    sentence = re.sub(r" Rd\.", " road", sentence)
    sentence = re.sub(r"[Aa]\.[Mm]\.", "A_M", sentence)
    sentence = re.sub(r"[Pp]\.[Mm]\.", "P_M", sentence)
    sentence = re.sub(r"[Bb]\.[Cc]\.", "B_C", sentence)
    sentence = re.sub(r"[Ad]\.[Dd]\.", "A_D", sentence)
    sentence = sentence.replace("Mr.", "mister")
    sentence = sentence.replace("Ms.", "miss")
    sentence = sentence.replace("Mrs.", "misses")
    sentence = sentence.replace("Ph.D", "P_H_D")
    sentence = sentence.replace("i.e.", "that is")
    sentence = sentence.replace("e.g.", "for example")
    sentence = sentence.replace("btw.", "by the way")
    sentence = sentence.replace("btw", "by the way")
    sentence = sentence.replace("b.t.w.", "by the way")
    sentence = sentence.replace("@", " at ")
    return sentence


class TextNormalizer:
    def __init__(self):
        self.language_detector = LanguageDetectorBuilder.from_languages(
            Language.ENGLISH, Language.CHINESE
        ).build()
        try:
            self.zh_normalizer = ZhNormalizer()
            self.en_normalizer = EnNormalizer()
        except:
            pass
        self.inflect_parser = inflect.engine()
        self.lang2token = {Language.ENGLISH: "en", Language.CHINESE: "zh"}

    def tn(self, text):
        text = preprocess_text(text)
        text = rettt(text)  # regex replacements
        # for non chinese languages
        language = self.language_detector.detect_language_of(text)
        # enforce chinese if text contains any chinese character
        if contains_chinese(text):
            language = Language.CHINESE
        text_lang = self.lang2token.get(language, "zh")

        if is_upper_eng_and_digit(text):
            language = Language.CHINESE

        if language == Language.CHINESE:
            try:
                text = self.zh_normalizer.normalize(text)
            except:
                pass
            text = text.replace("\n", "")
            text = re.sub(r"[，,]+$", "。", text)
        else:
            text = re.sub(r"[^ 0-9A-Za-z\[\]'.,:?!_\-]", "", text)
            try:
                text = self.en_normalizer.normalize(text)
            except:
                pass
            # fallback number normalization
            pieces = re.split(r"(\d+)", text)
            text = "".join(
                [
                    self.inflect_parser.number_to_words(p) if p.isnumeric() else p
                    for p in pieces
                    if len(p) > 0
                ]
            )

        # cleanup
        text = text.replace("_", " ")
        text = re.sub(r"[ ]+", " ", text)

        # spell caplital words
        pieces = re.split(r"([A-Z]{2,4}|[ ])", text)
        for idx, p in enumerate(pieces):
            if re.match("[A-Z]{2,4}", p):
                pieces[idx] = " ".join(p)
        text = " ".join([p for p in pieces if p != " "])

        # post TN full to half
        text = text.replace("。", ".")
        text = text.replace("，", ",")
        text = text.replace("：", ":")

        # model limitations
        text = text.lower().strip()
        text = text.replace('"', "")
        text = text.replace("·", " ")
        text = re.sub("[…~、！，？：；!?:;]+", ",", text)
        text = re.sub("[,]+", ",", text)
        text = re.sub(r"[,. ]+$", ".", text)
        if len(text) > 0 and text[-1] != ".":
            text = text + "."

        return text, text_lang