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):
    translate_text_prompt = None
    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 translate_text_prompt is 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:
        raise e

    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"
    DESCRIPTION = "This is a node that translates the prompt into another language using Argos Translate."
    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  ===== ###