# Generates translations
import asyncio
import ctranslate2
import time

from preprocessors.strip_split import splitAndClean
from preprocessors.tokenize import Tokenizer
from typing import List

MODEL_PATH = "/model"

translator = ctranslate2.Translator(
    MODEL_PATH, device="cpu",
    intra_threads=2
)
tokenizer = Tokenizer(MODEL_PATH)

async def translate_batch_async(src_lang: List[str], tgt_lang: List[str], batches_textArrArr: List[List[str]]) -> List[str]:

    # e.g. batches_textArrArr =: [["I hate eggs. I love fish", "I love dogs"], ["The cat is on the table"]]
    # Flatten the passed batches down to simple flat arrays
    flatBatches_textArr: List[str] = []
    # lang codes..
    flatBatches_src_lang: List[str] = []
    flatBatches_tgt_lang: List[str] = []

    for i, batch_textArr in enumerate(batches_textArrArr):

        for text in batch_textArr:
            flatBatches_textArr.append(text)
            flatBatches_src_lang.append(src_lang[i])
            flatBatches_tgt_lang.append(tgt_lang[i])

    # flatBatches_textArr = ["I hate eggs. I love fish", "I love dogs", "The cat is on the table"]
    # Further divide each string to sentences using a sentence splitter..

    sentences: List[str] = []
    # How many sentences were in each passed string..
    sentences_counts: List[int] = []
    sentences_src_lang: List[str] = []
    sentences_tgt_lang: List[str] = []

    for i, thisTranslateText in enumerate(flatBatches_textArr):

        thisSrcLang = flatBatches_src_lang[i]
        thisTgtLang = flatBatches_tgt_lang[i]
        theseSentences: List[str] = splitAndClean(
            thisTranslateText, thisSrcLang)

        sentences_counts.append(len(theseSentences))
        sentences.extend(theseSentences)
        sentences_src_lang.extend(
            [thisSrcLang]*len(theseSentences))
        sentences_tgt_lang.extend(
            [thisTgtLang]*len(theseSentences))

    # Tokenize the sentences

    sentences_tokensied: List[List[str]] = []

    for i, sentence in enumerate(sentences):

        thisSrcLang= sentences_src_lang[i]
        thisSentenceTokens = tokenizer.encode(sentence, thisSrcLang)
        sentences_tokensied.extend([thisSentenceTokens])

    # ok, let's translate already..

    print(f"Processing batch: {len(sentences_tokensied)}")

    def sync_func():
        return translator.translate_batch(
            sentences_tokensied,
            target_prefix=[[thisDestLang]
                           for thisDestLang in sentences_tgt_lang],
            max_batch_size=128
        )

    # Run sync_func asyncronously, so we don't block the event loop.
    # Allows other requests to be handled meanwhile.
    loop = asyncio.get_event_loop()
    results = await loop.run_in_executor(None, lambda: sync_func())

    targets = [result.hypotheses[0][1:] for result in results]

    sentences_translations = [tokenizer.decode(target) for target in targets]

    # Let's reconstruct back to where we split with the sentence splitter..

    flatBatches_translations: List[str] = []

    for count in sentences_counts:
        # Joining with a space, ideally this would be language specific..
        flatBatches_translations.append(
            ' '.join(sentences_translations[:count]))
        # Remove these items from the list
        sentences_translations[:count] = []

    # Let's assemble back into the passed batches..
    batches_translationsArrArr: List[List[str]] = []

    # Loop over the input batches
    for batch_textArr in batches_textArrArr:
        batch_translationArr: List[str] = []
        # loop over the strings passed in each batch
        for text in batch_textArr:
            batch_translationArr.append(flatBatches_translations.pop(0))
        batches_translationsArrArr.append(batch_translationArr)

    return batches_translationsArrArr