"""Czech punctuation restoration."""
import os
import sys
import time

import streamlit as st

st.set_page_config(page_title="Czech Punctuation", page_icon="😻")

REPO = "AILabTUL/electra-punct-cs-adaptive"
TOKEN = os.environ.get("HF_TOKEN_AILAB") or os.environ.get("HF_TOKEN")


def sentence_case(text):
    """Capitalise the first word and every word that follows '.' or '?'."""
    words, upper = [], True
    for word in text.split():
        words.append(word[0].upper() + word[1:] if upper and word else word)
        upper = word.endswith(".") or word.endswith("?")
    return " ".join(words)


@st.cache_resource(show_spinner="Loading the model...")
def load_model():
    from huggingface_hub import snapshot_download
    # The model code lives in the private repo, not here.
    root = snapshot_download(REPO, token=TOKEN)
    if root not in sys.path:
        sys.path.insert(0, root)
    from modeling import PunctuationRestorer
    return PunctuationRestorer.from_pretrained(root, token=TOKEN)


st.title("Czech Punctuation Restoration")
st.caption("Restores '.', ',' and '?' in Czech text and capitalises sentence "
           "starts. Write the input the way an ASR emits it - with diacritics, "
           "without punctuation.")

default = ("ministerstvo dopravy dnes představilo nový plán oprav dálnic "
           "který má podle úřadu zkrátit délku kolon práce začnou na jaře "
           "příštího roku a potrvají zhruba osmnáct měsíců řidiči se ale "
           "obávají že se situace nejprve zhorší kolik bude celá akce stát "
           "zatím není jasné mluvčí uvedl že přesné náklady se dozvíme až "
           "po vyhodnocení nabídek opozice plán kritizuje a tvrdí že peníze "
           "měly jít jinam")
text = st.text_area("Text", default, height=160)

if st.button("Restore punctuation", type="primary"):
    if not text.strip():
        st.warning("Enter some text first.")
    else:
        try:
            model = load_model()
            start = time.perf_counter()
            out, stats = model(text, return_stats=True)
            elapsed = time.perf_counter() - start
            st.success(sentence_case(out))
            st.metric("Tokens per second", "%.0f" % (stats["tokens"] / elapsed))
        except Exception as error:                                  # noqa: BLE001
            st.error("Could not run the model: %s" % error)