import os
from dotenv import load_dotenv
import gradio as gr
from AinaTheme import theme
from api_endpoint import invoke_translate_endpoint
from translate import translate_text

load_dotenv()

MODELS_PATH = "./models"
HF_CACHE_DIR = "./hf_cache"
MAX_INPUT_CHARACTERS = int(os.environ.get("MAX_INPUT_CHARACTERS", default=1000))
API_ENDPOINT_ENABLED = os.environ.get("API_ENDPOINT_ENABLED", default=True) == "True"


directions = {
    "Catalan": {
        "target": {
            "Spanish": {"src": "ca", "tgt":"es"},
            "English": {"src": "ca", "tgt":"en"},
            "French": {"src": "ca", "tgt":"fr"},
            "German": {"src": "ca", "tgt":"de"},
            "Italian": {"src": "ca", "tgt":"it"},
            "Portuguese": {"src": "ca", "tgt":"pt"}

        }
    },
    "Spanish": {
        "target": {
            "Catalan": {"src": "es", "tgt":"ca"},
        }
    },
    "English": {
        "target": {
            "Catalan": {"src": "en", "tgt":"ca"},
        }
    },
    "French": {
        "target": {
            "Catalan": {"src": "fr", "tgt":"ca"},
        }
    },
    "Portuguese": {
        "target": {
            "Catalan": {"src": "pt", "tgt":"ca"},
        }
    }
}

DEFAULT_SOURCE_LANGUAGE = list(directions.keys())[0]

def get_target_languages(source_language):
    return list(directions.get(source_language, {}).get("target", {}).keys())


def get_target_languege_model(source_language, target_language):
    # return directions.get(source_language, {}).get("target", {}).get(target_language, {}).get("model")
    return directions.get(source_language, {}).get("target", {}).get(target_language, {})


def translate(text, source_language, lang_pair):
    """Use CTranslate model to translate a sentence

    Args:
        source (str): Source sentences to translate
        translator (object): Object of Translator, with the CTranslate2 model
        sp_model (object): Object of SentencePieceProcessor, with the SentencePiece source model
    Returns:
        Translation of the source text
    """
  
    translation = translate_text(text, source_language, lang_pair)

    return translation


def translate_input(input, source_language, target_language):

    if input.strip() == "":  
        gr.Warning('Not possible to translate an empty input.')
        return None
    
    target_language_model = get_target_languege_model(source_language, target_language)

    translation =  invoke_translate_endpoint(input, target_language_model)
    if translation is not None:
        return translation
    
    return None


def clear():
    return None, None

def change_interactive(text):
    if len(text.strip()) > MAX_INPUT_CHARACTERS:
        return gr.update(interactive = True),  gr.update(interactive = False)
    return gr.update(interactive = True),  gr.update(interactive = True)    

def update_target_languages_dropdown(source_language):

    output_languages = get_target_languages(source_language)
    return gr.update(choices=output_languages, value=output_languages[0], interactive=True)


with gr.Blocks(theme=theme) as app:
    with gr.Row(variant="panel"):    
        with gr.Column(scale=2):
            placeholder_max_token = gr.Textbox(
                    visible=False,
                    interactive=False,
                    value= MAX_INPUT_CHARACTERS
                )
            source_language = gr.Dropdown(label="Source Language", choices=list(directions.keys()), value=DEFAULT_SOURCE_LANGUAGE)
            input = gr.Textbox(placeholder="Enter a text here to translate.",  max_lines=100, lines=12, show_label=False, interactive=True)
            with gr.Row(variant="panel", equal_height=True):
                    gr.HTML("""<span id="countertext" style="display: flex; justify-content: start; color:#ef4444; font-weight: bold;"></span>""")
                    gr.HTML(f"""<span id="counter" style="display: flex; justify-content: end;"> <span id="inputlenght">0</span>&nbsp;/&nbsp;{MAX_INPUT_CHARACTERS}</span>""")

        with gr.Column(scale=2):

            target_outputs = get_target_languages(DEFAULT_SOURCE_LANGUAGE)
            #target_language = gr.Dropdown(choices=target_outputs, label="Target Language", value=target_outputs[0])

            target_language = gr.Radio(choices=target_outputs, label="Target Language", value=target_outputs[0])
            output = gr.Textbox(max_lines=100, lines=12, show_label=False, interactive=False, show_copy_button=True)

            with gr.Row(variant="panel"):
                    clear_btn = gr.Button(
                        "Clear", 
                    )
                    submit_btn = gr.Button(
                        "Submit", 
                        variant="primary",
                    )

        source_language.change(fn=update_target_languages_dropdown, inputs=[source_language], outputs=target_language)

        input.change(
            fn=change_interactive, 
            inputs=[input], 
            outputs=[clear_btn, submit_btn],
            api_name=False
        )
      
        input.change(
            fn=None, 
            inputs=[input], 
            js=f"""(i) => document.getElementById('countertext').textContent =  i.length > {MAX_INPUT_CHARACTERS} && 'Max length {MAX_INPUT_CHARACTERS} characters. ' || '' """,
            api_name=False
        )

        input.change(
            fn=None, 
            inputs=[input, placeholder_max_token], 
            js="""(i, m) => {
                document.getElementById('inputlenght').textContent = i.length + '  '
                document.getElementById('inputlenght').style.color =  (i.length > m) ? "#ef4444" : "";
            }""",
            api_name=False
            )
    
        clear_btn.click(
            fn=clear, 
            inputs=[], 
            outputs=[input, output], 
            queue=False,
            api_name=False
        )
        
        submit_btn.click(
            fn=translate_input, 
            inputs=[input, source_language, target_language], 
            outputs=[output],
            api_name="translate",
            concurrency_limit=1,
        )

app.launch(show_api=True)