from server import PromptServer
from aiohttp import web

import argostranslate.package
import argostranslate.translate

# Find packages https://www.argosopentech.com/argospm/index/

### =====  Argos Translate Node  ===== ###
ALL_CODES = {
    "english": {
        "code": "en",
        "targets": [
            "albanian",
            "arabic",
            "azerbaijani",
            "bengali",
            "bulgarian",
            "catalan",
            "chinese",
            "chinese (traditional)",
            "czech",
            "danish",
            "dutch",
            "esperanto",
            "estonian",
            "finnish",
            "french",
            "german",
            "greek",
            "hebrew",
            "hindi",
            "hungarian",
            "indonesian",
            "irish",
            "italian",
            "japanese",
            "korean",
            "latvian",
            "lithuanian",
            "malay",
            "norwegian",
            "persian",
            "polish",
            "portuguese",
            "romanian",
            "russian",
            "slovak",
            "slovenian",
            "spanish",
            "swedish",
            "tagalog",
            "thai",
            "turkish",
            "ukrainian",
            "urdu",
        ],
    },
    "albanian": {"code": "sq", "targets": ["english"]},
    "arabic": {"code": "ar", "targets": ["english"]},
    "azerbaijani": {"code": "az", "targets": ["english"]},
    "bengali": {"code": "bn", "targets": ["english"]},
    "bulgarian": {"code": "bg", "targets": ["english"]},
    "catalan": {"code": "ca", "targets": ["english"]},
    "chinese": {"code": "zh", "targets": ["english"]},
    "chinese (traditional)": {"code": "zt", "targets": ["english"]},
    "czech": {"code": "cs", "targets": ["english"]},
    "danish": {"code": "da", "targets": ["english"]},
    "dutch": {"code": "nl", "targets": ["english"]},
    "esperanto": {"code": "eo", "targets": ["english"]},
    "estonian": {"code": "et", "targets": ["english"]},
    "finnish": {"code": "fi", "targets": ["english"]},
    "french": {"code": "fr", "targets": ["english"]},
    "german": {"code": "de", "targets": ["english"]},
    "greek": {"code": "el", "targets": ["english"]},
    "hebrew": {"code": "he", "targets": ["english"]},
    "hindi": {"code": "hi", "targets": ["english"]},
    "hungarian": {"code": "hu", "targets": ["english"]},
    "indonesian": {"code": "id", "targets": ["english"]},
    "irish": {"code": "ga", "targets": ["english"]},
    "italian": {"code": "it", "targets": ["english"]},
    "japanese": {"code": "ja", "targets": ["english"]},
    "korean": {"code": "ko", "targets": ["english"]},
    "latvian": {"code": "lv", "targets": ["english"]},
    "lithuanian": {"code": "lt", "targets": ["english"]},
    "malay": {"code": "ms", "targets": ["english"]},
    "norwegian": {"code": "nb", "targets": ["english"]},
    "persian": {"code": "fa", "targets": ["english"]},
    "polish": {"code": "pl", "targets": ["english"]},
    "portuguese": {"code": "pt", "targets": ["english", "spanish"]},
    "romanian": {"code": "ro", "targets": ["english"]},
    "russian": {"code": "ru", "targets": ["english"]},
    "slovak": {"code": "sk", "targets": ["english"]},
    "slovenian": {"code": "sl", "targets": ["english"]},
    "spanish": {"code": "es", "targets": ["english", "portuguese"]},
    "swedish": {"code": "sv", "targets": ["english"]},
    "tagalog": {"code": "tl", "targets": ["english"]},
    "thai": {"code": "th", "targets": ["english"]},
    "turkish": {"code": "tr", "targets": ["english"]},
    "ukrainian": {"code": "uk", "targets": ["english"]},
    "urdu": {"code": "ur", "targets": ["english"]},
}


@PromptServer.instance.routes.get("/alekpet/argo_langs_support/{lang}")
async def argo_langs_support(request):
    lang = request.match_info["lang"]

    if lang:
        langs_support = ALL_CODES[lang]
        return web.json_response(
            {
                "langs_support": langs_support["targets"],
                "lang_code": langs_support["code"],
            }
        )

    return web.json_response({"langs_support": [], "lang_code": "en"})


def installPackages(srcTrans, toTrans="en"):
    argostranslate.package.update_package_index()
    available_packages = argostranslate.package.get_available_packages()
    package_to_install = next(
        filter(
            lambda x: x.from_code == srcTrans and x.to_code == toTrans,
            available_packages,
        )
    )
    argostranslate.package.install_from_path(package_to_install.download())


def preTranslate(prompt, srcTrans, toTrans):
    if prompt and prompt.strip() != "":
        installed_languages = argostranslate.translate.get_installed_languages()

        from_lang = list(filter(lambda x: x.code == srcTrans, installed_languages))[0]
        to_lang = list(filter(lambda x: x.code == toTrans, installed_languages))[0]

        translation = from_lang.get_translation(to_lang)
        translate_text_prompt = translation.translate(prompt)

    return translate_text_prompt if translate_text_prompt and not None else ""


def translate(prompt, srcTrans=None, toTrans="english"):
    translate_text_prompt = ""
    try:
        srcTransCode = ALL_CODES[srcTrans]["code"] if srcTrans is not None else None
        toTransCode = ALL_CODES[toTrans]["code"]
        installPackages(srcTransCode, toTransCode)
        translate_text_prompt = preTranslate(prompt, srcTransCode, toTransCode)

    except Exception as e:
        print(e)
        return "[Error] No translate text!"

    return translate_text_prompt


class ArgosTranslateCLIPTextEncodeNode:
    @classmethod
    def INPUT_TYPES(self):
        self.langs_support = ALL_CODES["russian"]["targets"]
        return {
            "required": {
                "from_translate": (list(ALL_CODES.keys()), {"default": "russian"}),
                "to_translate": (self.langs_support, {"default": "english"}),
                "text": ("STRING", {"multiline": True, "placeholder": "Input text"}),
                "clip": ("CLIP",),
            }
        }

    RETURN_TYPES = (
        "CONDITIONING",
        "STRING",
    )
    FUNCTION = "argos_translate_text"
    CATEGORY = "AlekPet Nodes/conditioning"

    def argos_translate_text(self, from_translate, to_translate, text, clip):
        self.langs_support = ALL_CODES[from_translate]["targets"]
        text = translate(text, from_translate, to_translate)
        tokens = clip.tokenize(text)
        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
        return ([[cond, {"pooled_output": pooled}]], text)

    @classmethod
    def VALIDATE_INPUTS(cls, from_translate, to_translate, text, clip):
        return True


class ArgosTranslateTextNode(ArgosTranslateCLIPTextEncodeNode):
    @classmethod
    def INPUT_TYPES(self):
        return_types = super().INPUT_TYPES()
        del return_types["required"]["clip"]
        return return_types

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "argos_translate_text"

    CATEGORY = "AlekPet Nodes/text"

    def argos_translate_text(self, from_translate, to_translate, text):
        self.langs_support = ALL_CODES[from_translate]["targets"]
        text_tranlsated = translate(text, from_translate, to_translate)
        return (text_tranlsated,)

    @classmethod
    def VALIDATE_INPUTS(cls, from_translate, to_translate, text):
        return True


### =====  Argos Translate Node  ===== ###