import time

import streamlit as st
import subprocess
from highlighter import show_highlights
from utils import (
    download_quantized_model,
    download_spacy_model,
    errant_tokenize,
    get_annotator,
    load_model_and_tokenizer,
)

model_names = [
    "aseifert/t5-base-jfleg-wi",
    "aseifert/byt5-base-jfleg-wi",
    "team-writing-assistant/t5-base-c4jfleg",
    "Modfiededition/t5-base-fine-tuned-on-jfleg",
    "prithivida/grammar_error_correcter_v2",
]

st.write("Output of `pip freeze`:")
p = subprocess.Popen(["pip", "freeze"], stdout=subprocess.PIPE)
output = p.communicate()[0]
st.code(output.decode("utf-8"))


def predict(model, args, text: str):
    return model.generate_text(text, args=args).text


def main():
    st.title("🤗 Writing Assistant")
    st.markdown(
        """This writing assistant will proofread any text for you! See my [GitHub repo](https://github.com/aseifert/hf-writing-assistant) for implementation details."""
    )

    model_name = st.selectbox("Choose model", model_names)
    download_quantized_model(model_name)
    model, tokenizer = load_model_and_tokenizer(model_name)
    download_spacy_model()
    annotator = get_annotator("en")

    default_text = "A dog is bigger then mouse."
    default_text = "it gives him many apprtunites in the life, and i think that being knowledge person is a very wouderful thing to have so we can spend our lives in a successful way and full of happenis."
    input_text = st.text_area(
        label="Original text",
        value=default_text,
    )

    start = None
    if st.button("✍️ Check"):
        start = time.time()
        with st.spinner("Checking for errors 🔍"):
            output_text = predict(model, tokenizer, input_text)

            try:
                tokenized_input_text = errant_tokenize(input_text)
                tokenized_output_text = errant_tokenize(output_text)
                show_highlights(annotator, tokenized_input_text, tokenized_output_text)
                st.write("")
                st.success(output_text)
            except Exception as e:
                st.error("Some error occured!" + str(e))
                st.stop()

    st.write("---")
    st.markdown(
        "Built by [@therealaseifert](https://twitter.com/therealaseifert) during the HF community event – 👨\u200d💻 [GitHub repo](https://github.com/aseifert/hf-writing-assistant) – 🤗 Team Writing Assistant"
    )
    st.markdown(
        "_Highlighting code thanks to [Gramformer](https://github.com/PrithivirajDamodaran/Gramformer)_"
    )

    if start is not None:
        st.text(f"prediction took {time.time() - start:.2f}s")


if __name__ == "__main__":
    main()