import streamlit as st
from translator import Translators, translate_text, all_langs

models = Translators.MODELS
# Language options and mappings
favourite_langs = {"German": "de", "Romanian": "ro", "English": "en", "-----": "-----"}
langs = list(favourite_langs.keys())
langs.extend(list(all_langs.keys())) # Language options as list, add favourite languages first

# App layout
st.header("Text Machine Translation", divider="gray", help="Text Machine Translation Streamlit App with Open Source Models")
input_text = st.text_area("Enter text to translate:", placeholder="Enter text to translate, maximum 512 characters!", max_chars=512)

# Initialize session state
if "sselected_language" not in st.session_state:
    st.session_state["sselected_language"] = langs[0]
if "tselected_language" not in st.session_state:
    st.session_state["tselected_language"] = langs[1]
if "model_name" not in st.session_state:
    st.session_state["model_name"] = models[1]

# Model selection FIRST
model_name = st.selectbox("Select a model:", models, 
                          index=models.index(st.session_state["model_name"]))

# Create columns for language selection
scol, swapcol, tcol = st.columns([3, 1, 3])

with scol:
    sselected_language = st.selectbox("Source language:", langs, 
                                      index=langs.index(st.session_state["sselected_language"]))
with swapcol:
    if st.button("🔄 Swap"):
        st.session_state["model_name"] = model_name  # Preserve model
        st.session_state["sselected_language"], st.session_state["tselected_language"] = \
            st.session_state["tselected_language"], st.session_state["sselected_language"]
        st.rerun()
with tcol:
    tselected_language = st.selectbox("Target language:", langs, 
                                      index=langs.index(st.session_state["tselected_language"]))

# Language codes
sl = all_langs[st.session_state["sselected_language"]][0]
tl = all_langs[st.session_state["tselected_language"]][0]

# Store selections
st.session_state["sselected_language"] = sselected_language
st.session_state["tselected_language"] = tselected_language
st.session_state["model_name"] = model_name

st.write(f'Selected language combination: {sselected_language} - {tselected_language}. Selected model: {model_name}')

with st.container(border=None, width="stretch", height="content", horizontal=False, horizontal_alignment="center", vertical_alignment="center", gap="small"):
    submit_button = st.button("Translate")
# Show text area with placeholder also before translating
# translated_textarea = st.empty()
# message_textarea = st.empty()
# translated_textarea.text_area(":green[Translation:]", placeholder="Translation area", value='')
# message_textarea.text_input(":blue[Messages:]", placeholder="Messages area", value='')

if submit_button: # Handle the submit button click
    with st.spinner("Translating...", show_time=True):
        translated_text, message = translate_text(model_name, sselected_language, tselected_language, input_text)       
    print(f"Translated from {sselected_language} to {tselected_language} using {model_name}.", input_text, translated_text)
    # Display the translated text
    # translated_textarea.text_area(":green[Translation:]", value=translated_text)
    # message_textarea.text_input(":blue[Message:]", value=message)
    st.text_area(":green[Translation:]", value=translated_text)
    # st.success(message, icon=":material/check:") st.info(message, icon="ℹ️"), st.warning(message, icon=":material/warning:"), error(message, icon=":material/error:"), st.exception
    st.info(message, icon=":material/info:")
    # st.text_input(":blue[Messages:]", value=message)
    # st.rerun()