import os
import time

import requests
import streamlit as st

API_URL = "https://api-inference.huggingface.co/models/NbAiLab/norwegian-gpt2-vgd"


def translate(text, wait=True):
    headers = {"Authorization": f"Bearer {os.environ['BEARER']}"}
    payload  = {
        "inputs": text,
        "options": {
            "wait_for_model": not wait
        }
    }
    response = requests.post(API_URL, headers=headers, json=payload)
    json_response = response.json()
    if (isinstance(json_response, dict)
        and "error" in json_response
        and "estimated_time" in json_response):
        st.write(json_response)
        if wait:
            with st.spinner(json_response["error"]):
                bar = st.progress(0)
                time_to_load = int(json_response["estimated_time"]) + 1
                for progress in range(time_to_load):
                    bar.progress(progress / time_to_load)
                    time.sleep(1)
                bar.empty()
            return translate(text, wait=False)
        else:
            return "We could not load the model"
    elif (isinstance(json_response, list)
        and "translation_text" in json_response[0]):
        return json_response[0]["translation_text"]
    else:
        return f"Oops, something went wrong: {str(json_response)}"


st.set_page_config(
    page_title='Spurningen',
    page_icon='translator-icon.png',
)
st.title("Bokmål ⇔ Nynorsk")
st.sidebar.title("Translation Demo")
st.sidebar.write("""
Spurningen  (P.Chr. Asbjørnsen og Jørgen Moe Norske folkeeventyr (1852). I senere utgaver er eventyret kalt "Prinsessen som ingen kunde målbinde". Spurning er et gammelt norrønt ord for gåte. En spurning er en utspørring. Denne GPT2-modellen er trent på aviser, wikipedia og diverse sosiale media. Bl.a. er den trent på spørsmål og svar fra VGDebatt. Den er trent for å ikke la seg målbinde!

Dette er en demo-versjon hovedsakelig ment for forskning på auto-regressive språkmodeller. Det er ingen filter-funksjon her. Modellen kan godt generere tekst som i en annen kontekst kan være ekstremt fornærmende. Husk at dette er en lek med sammensetning av ord. Ikke la deg fornærme.

""")


text = st.text_area(" ",
    "Nasjonalbibliotekets AI-lab har faktisk utviklet et språkmodell som ikke kan målbindes. Hvor tror dere dette ender?",
    height=None,
    max_chars=None,
    key=None,
    help="Enter your text here",
)

if st.button('Translate'):
    if str(text).strip() == "":
        st.warning('Please **enter text** for translation')
    else:
        st.info(str(translate(text)))