import gradio as gr
from simplemma import simple_tokenizer
from difflib import Differ
from icecream import ic
from app.webui.patch import model_load,num_tokens_in_string,one_chunk_initial_translation, one_chunk_reflect_on_translation, one_chunk_improve_translation
from app.webui.patch import calculate_chunk_size, multichunk_initial_translation, multichunk_reflect_on_translation, multichunk_improve_translation

from llama_index.core.node_parser import SentenceSplitter

from translatepy.translators.google import GoogleTranslate
from translatepy.exceptions import UnknownLanguage
from translatepy import Language

gtranslator = GoogleTranslate()

progress=gr.Progress()

def tokenize(text):
    # Use nltk to tokenize the text
    words = simple_tokenizer(text)
    # Check if the text contains spaces
    if ' ' in text:
        # Create a list of words and spaces
        tokens = []
        for word in words:
            tokens.append(word)
            if not word.startswith("'") and not word.endswith("'"):  # Avoid adding space after punctuation
                tokens.append(' ')  # Add space after each word
        return tokens[:-1]  # Remove the last space
    else:
        return words

def diff_texts(text1, text2):
    tokens1 = tokenize(text1)
    tokens2 = tokenize(text2)

    d = Differ()
    diff_result = list(d.compare(tokens1, tokens2))

    highlighted_text = []
    for token in diff_result:
        word = token[2:]
        category = None
        if token[0] == '+':
            category = 'added'
        elif token[0] == '-':
            category = 'removed'
        elif token[0] == '?':
            continue  # Ignore the hints line

        highlighted_text.append((word, category))

    return highlighted_text

#modified from src.translaation-agent.utils.tranlsate
def translator(
        source_lang: str,
        target_lang: str,
        source_text: str,
        country: str,
        max_tokens:int = 1000,
):
    """Translate the source_text from source_lang to target_lang."""
    num_tokens_in_text = num_tokens_in_string(source_text)

    ic(num_tokens_in_text)

    if num_tokens_in_text < max_tokens:
        ic("Translating text as single chunk")

        progress((1,3), desc="First translation...")
        init_translation = one_chunk_initial_translation(
            source_lang, target_lang, source_text
        )

        progress((2,3), desc="Reflecton...")
        reflection = one_chunk_reflect_on_translation(
            source_lang, target_lang, source_text, init_translation, country
        )

        progress((3,3), desc="Second translation...")
        final_translation = one_chunk_improve_translation(
            source_lang, target_lang, source_text, init_translation, reflection
        )

        return init_translation, reflection, final_translation

    else:
        ic("Translating text as multiple chunks")

        progress((1,5), desc="Calculate chunk size...")
        token_size = calculate_chunk_size(
            token_count=num_tokens_in_text, token_limit=max_tokens
        )

        ic(token_size)

        #using sentence splitter
        text_parser = SentenceSplitter(
           chunk_size=token_size,
        )

        progress((2,5), desc="Spilt source text...")
        source_text_chunks = text_parser.split_text(source_text)

        progress((3,5), desc="First translation...")
        translation_1_chunks = multichunk_initial_translation(
            source_lang, target_lang, source_text_chunks
        )

        init_translation = "".join(translation_1_chunks)

        progress((4,5), desc="Reflection...")
        reflection_chunks = multichunk_reflect_on_translation(
            source_lang,
            target_lang,
            source_text_chunks,
            translation_1_chunks,
            country,
        )

        reflection = "".join(reflection_chunks)

        progress((5,5), desc="Second translation...")
        translation_2_chunks = multichunk_improve_translation(
            source_lang,
            target_lang,
            source_text_chunks,
            translation_1_chunks,
            reflection_chunks,
        )

        final_translation = "".join(translation_2_chunks)

        return init_translation, reflection, final_translation


def translator_sec(
        endpoint2: str,
        model2: str,
        api_key2: str,
        context_window: int,
        num_output: int,
        source_lang: str,
        target_lang: str,
        source_text: str,
        country: str,
        max_tokens: int = 1000,
        gt: bool = False,
):

    """Translate the source_text from source_lang to target_lang."""
    num_tokens_in_text = num_tokens_in_string(source_text)

    ic(num_tokens_in_text)

    if num_tokens_in_text < max_tokens:
        ic("Translating text as single chunk")

        progress((1,3), desc="First translation...")
        if gt:
            try:
                language = Language(target_lang)
            except Exception as e:
                raise gr.Error(f"An unexpected error occurred: {e}")
            init_translation = gtranslator.translate(source_text, language).result
        else:
            init_translation = one_chunk_initial_translation(
                source_lang, target_lang, source_text
            )
            try:
                model_load(endpoint2, model2, api_key2, context_window, num_output)
            except Exception as e:
                raise gr.Error(f"An unexpected error occurred: {e}")
        
        progress((2,3), desc="Reflecton...")
        reflection = one_chunk_reflect_on_translation(
            source_lang, target_lang, source_text, init_translation, country
        )

        progress((3,3), desc="Second translation...")
        final_translation = one_chunk_improve_translation(
            source_lang, target_lang, source_text, init_translation, reflection
        )

        return init_translation, reflection, final_translation

    else:
        ic("Translating text as multiple chunks")

        progress((1,5), desc="Calculate chunk size...")
        token_size = calculate_chunk_size(
            token_count=num_tokens_in_text, token_limit=max_tokens
        )

        ic(token_size)

        #using sentence splitter
        text_parser = SentenceSplitter(
           chunk_size=token_size,
        )

        progress((2,5), desc="Spilt source text...")
        source_text_chunks = text_parser.split_text(source_text)

        progress((3,5), desc="First translation...")
        if gt:
            try:
                language = Language(target_lang)
            except Exception as e:
                raise gr.Error(f"An unexpected error occurred: {e}")
            translation_1_chunks = gtranslator.translate(source_text_chunks, language).result
        else:
            translation_1_chunks = multichunk_initial_translation(
                source_lang, target_lang, source_text_chunks
            )
            try:
                model_load(endpoint2, model2, api_key2, context_window, num_output)
            except Exception as e:
                raise gr.Error(f"An unexpected error occurred: {e}")
                
        init_translation = "".join(translation_1_chunks)
        progress((4,5), desc="Reflection...")
        reflection_chunks = multichunk_reflect_on_translation(
            source_lang,
            target_lang,
            source_text_chunks,
            translation_1_chunks,
            country,
        )

        reflection = "".join(reflection_chunks)

        progress((5,5), desc="Second translation...")
        translation_2_chunks = multichunk_improve_translation(
            source_lang,
            target_lang,
            source_text_chunks,
            translation_1_chunks,
            reflection_chunks,
        )

        final_translation = "".join(translation_2_chunks)

        return init_translation, reflection, final_translation