import streamlit as st
import models

def main():
    st.set_page_config(layout="wide", page_title="Review Toxicity Checker")
    st.title('Review Toxicity Checker')
    left_col, right_col = st.columns(2)

    content = left_col.empty()

    with st.sidebar:
        model_id = st.radio(
            "Select a model option:",
            ["stabilityai/stablelm-zephyr-3b", "mistralai/Mistral-7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.2", "microsoft/Phi-3-mini-4k-instruct", "harir/stablelm-zephyr-3b-review-toxicity", "harir/phi-3-mini-review-toxicity"]
        )
        #hf_api_key = st.text_input('HF API Key\nhttps://huggingface.co/settings/tokens')
        color = st.color_picker('Highlight Color', '#F44336')

    with right_col:
        st.markdown("<p style='text-align: left; color: white; margin-bottom: 5px; font-size: 14px;'>Revised review:</p>", unsafe_allow_html=True)

    with left_col:
        input_text = st.text_area('Enter your review here:', height=500)
        left_col2, right_col2 = st.columns([1,3])
        with left_col2:
            check = st.button('Check Review')
        with right_col2:
            clear = st.button('Clear Text')
        if check:
            hf_api_key = st.secrets["default_hf_api"]
            with st.spinner('Processing review...'):
                revision = models.revise_review(input_text, hf_api_key, model_id, color)
            try: 
                with right_col:
                    st.markdown(f""" <div style="border:1px solid white; padding:10px; height:500px; overflow:auto; border-radius:7px;"> {revision['data']['revision']} </div> """, unsafe_allow_html=True)
                    num_revised = int(revision['data']['revised_sentences'])/int(revision['data']['sentence_count'])
                    left_col2, right_col2 = st.columns([1,2])
                    st.write(f"Sentences Revised: {revision['data']['revised_sentences']}/{revision['data']['sentence_count']}")
                    st.progress(num_revised)
                    score = revision['data']['score']
                    if score == 1:
                        st.write(f"Score: 1 (toxic)")
                    if score == 0:
                        st.write(f"Score: 0 (non-toxic)")
                    st.progress(int(revision['data']['score'])/1)
            except Exception as e:
                 with right_col:
                    st.error("An error occured.")
                    st.error(e)
        elif clear:
            with right_col:
                st.markdown(f""" <div style="border:1px solid white; padding:10px; height:500px; overflow:auto; border-radius:7px;"> </div> """, unsafe_allow_html=True)
        else:
            with right_col:
                st.markdown(f""" <div style="border:1px solid white; padding:10px; height:500px; overflow:auto; border-radius:7px;"> </div> """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()