import streamlit as st
from transformers import pipeline
from streamlit_language_texts import *
import torch
import os

st.set_page_config(layout="wide", page_title=app_title["eng"], page_icon="🇲🇦")

token = os.getenv("HF_TOKEN")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
languages = {"🇬🇧 English": "eng", "🇲🇦 Darija": "ary", "🇸🇦 Arabic": "ara", "🇫🇷 French": "fra"}

if "lang" not in st.query_params:
    st.query_params["lang"] = "eng"
    st.rerun()

def set_language():
    if "selected_language" in st.session_state:
        st.query_params["lang"] = languages.get(st.session_state["selected_language"])

with st.sidebar:
    sel_lang = st.selectbox(
        "Select a primary language for this app",
        options=languages,
        on_change=set_language,
        key="selected_language",
    )

if 'eng2ary_pipe' not in st.session_state:
    st.session_state['eng2ary_pipe'] = pipeline("text2text-generation", model="ychafiqui/english-to-darija-2", token=token, device=device)
if 'ary2eng_pipe' not in st.session_state:
    st.session_state['ary2eng_pipe'] = pipeline("text2text-generation", model="ychafiqui/darija-to-english-2", token=token, device=device)

def eng_to_ary(text):
    if text:
        return st.session_state.eng2ary_pipe(text)[0]['generated_text']
    else:
        return ""

def ary_to_eng(text):
    if text:
        return st.session_state.ary2eng_pipe(text)[0]['generated_text']
    else:
        return ""

st.title(app_title[languages[sel_lang]])
st.write(app_description[languages[sel_lang]])

direction = st.radio(".", ["English to Darija", "Darija to English"], key="lang", horizontal=True, label_visibility="collapsed")

if direction == "English to Darija":
    input_col, output_col = st.columns(2)
    input_col_container = input_col.container()

    eng_text = input_col_container.text_area("🇬🇧 Write English text here", key="eng_text")

    output_col_container = output_col.container()
    output_col_container.text_area("🇲🇦 Darija translation will appear here", value=eng_to_ary(eng_text))

elif direction == "Darija to English":
    input_col, output_col = st.columns(2)
    input_col_container = input_col.container()

    ary_text = input_col_container.text_area("🇲🇦 Write Darija text here", key="ary_text")

    output_col_container = output_col.container()
    output_col_container.text_area("🇬🇧 English translation will appear here", value=ary_to_eng(ary_text))

st.markdown(app_footer[languages[sel_lang]])