import gradio as gr
from huggingface_hub import InferenceClient

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
englishToFinnishClient = InferenceClient("Helsinki-NLP/opus-mt-en-fi")
googleClient = InferenceClient("google/flan-t5-large")#"facebook/nllb-200-1.3B")#"facebook/nllb-200-1.3B")
finnishToEnglishClient = InferenceClient("Helsinki-NLP/opus-mt-fi-en")


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    sourceLanguage,
    targetLanguage,
    useGoogle
):
    toFinnish=targetLanguage=='Finnish' and sourceLanguage=="English"
    fromFinnish=targetLanguage=='English' and sourceLanguage=="Finnish"

    #adminClient=englishToFinnishClient if(toFinnish)else googleClient 

    #candidateClient=finnishToEnglishClient if(fromFinnish)else googleClient 

    print("the text",message, 'the source',sourceLanguage,"the target",targetLanguage)
  
    if(useGoogle=="true"):
        print("using google")
        text="Translate to "f"{targetLanguage}: "f"{message}"
        googleResponse=  googleClient.text_generation(
            text, max_new_tokens=240
        )

        return googleResponse
    
    if(toFinnish):
        print("to finnish")
        fiResponse=  englishToFinnishClient.translation(
            message,
            #max_new_tokens=100
            tgt_lang= targetLanguage,
            src_lang=sourceLanguage  
        )

        return fiResponse.translation_text
    
    if(fromFinnish):
        print("From Finnish called")
        fiResponse=finnishToEnglishClient.translation(
            message,
            tgt_lang= targetLanguage,
            src_lang=sourceLanguage  
        )

        return fiResponse.translation_text


    """ if(toFinnish):
        fiResponse=  englishToFinnishClient.translation(
        message,
        #max_new_tokens=100
        tgt_lang= targetLanguage,
        src_lang=sourceLanguage  
        )

        text="Translate to "f"{targetLanguage}: "f"{message}"
        print("the text",text)
        googResponse=  googleClient.text_generation(
        text
        )

        response=f"{googResponse}<b>Translation 2</b><br/> {fiResponse.translation_text}"
 
    else:    
        text="Translate to "f"{targetLanguage}: "f"{message}"
        print("the text",text)

        googleResponse=  googleClient.text_generation(
          text
        )

        response=f"{googleResponse}"

        if(fromFinnish):
            fiResponse=finnishToEnglishClient.translation(
                message,
                tgt_lang= targetLanguage,
                src_lang=sourceLanguage  
            )

            response+=f'<b>Translation 2</b><br/>{fiResponse.translation_text}'

    print("the response",response)

    response """
        

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a language translator", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
        gr.Textbox(value="Finnish", label="Source Language"),
        gr.Textbox(value="English", label="Target Language"),
        gr.Textbox(value="false", label="Use Google"),
    ],
)


if __name__ == "__main__":
    demo.launch()