import streamlit as st
from backend import load_model, translate

# -----------------------------
# PAGE CONFIG
# -----------------------------
st.set_page_config(page_title="Kurdish Translator", layout="wide")

# Initialize session state
if "theme" not in st.session_state:
    st.session_state.theme = "dark"
if "src_lang" not in st.session_state:
    st.session_state.src_lang = "English"
if "tgt_lang" not in st.session_state:
    st.session_state.tgt_lang = "Kurdish"
if "output_text" not in st.session_state:
    st.session_state.output_text = ""

# Theme colors
THEME = st.session_state.theme
bg = "#111111" if THEME == "dark" else "#ffffff"
fg = "#ffffff" if THEME == "dark" else "#000000"
card_bg = "rgba(255,255,255,0.07)" if THEME == "dark" else "rgba(0,0,0,0.05)"
input_bg = "#1e1e1e" if THEME == "dark" else "#f5f5f5"
button_bg = "#2d2d2d" if THEME == "dark" else "#e0e0e0"

st.markdown(
    f"""<style>
[data-testid="stAppViewContainer"] {{
    background-color: {bg};
    color: {fg};
}}
[data-testid="stHeader"] {{
    background-color: {bg};
}}
[data-testid="stToolbar"] {{
    background-color: {bg};
}}
.stTextArea textarea {{
    background-color: {input_bg} !important;
    color: {fg} !important;
}}
.stSelectbox {{
    color: {fg};
}}
h1, h2, h3, p, label {{
    color: {fg} !important;
}}
.dots-loader {{
    display: flex;
    justify-content: center;
    margin: 20px 0;
}}
.dots-loader div {{
    width: 12px;
    height: 12px;
    margin: 4px;
    background-color: #4A90E2;
    border-radius: 50%;
    animation: bounce 0.6s infinite alternate;
}}
.dots-loader div:nth-child(2) {{ animation-delay: 0.2s; }}
.dots-loader div:nth-child(3) {{ animation-delay: 0.4s; }}
@keyframes bounce {{
    from {{ transform: translateY(0); }}
    to {{ transform: translateY(-12px); }}
}}
.model-loader {{
    display: flex;
    flex-direction: column;
    justify-content: center;
    align-items: center;
    height: 60vh;
}}
.spinner {{
    width: 60px;
    height: 60px;
    border: 4px solid {card_bg};
    border-top: 4px solid #4A90E2;
    border-radius: 50%;
    animation: spin 1s linear infinite;
}}
@keyframes spin {{
    0% {{ transform: rotate(0deg); }}
    100% {{ transform: rotate(360deg); }}
}}
.loading-text {{
    margin-top: 20px;
    font-size: 18px;
    color: #4A90E2;
    animation: pulse 1.5s ease-in-out infinite;
}}
@keyframes pulse {{
    0%, 100% {{ opacity: 1; }}
    50% {{ opacity: 0.5; }}
}}
.output-box {{
    padding: 20px;
    border-radius: 12px;
    background: {card_bg};
    font-size: 18px;
    min-height: 120px;
    word-wrap: break-word;
    color: {fg};
}}
</style>""",
    unsafe_allow_html=True
)

# -----------------------------
# LOAD MODEL
# -----------------------------
@st.cache_resource
def get_model():
    return load_model()

# Show custom loader while model is loading
if "model_loaded" not in st.session_state:
    loader_container = st.empty()
    loader_container.markdown(
        """
        <div class='model-loader'>
            <div class='spinner'></div>
            <div class='loading-text'>Loading translation model...</div>
        </div>
        """,
        unsafe_allow_html=True
    )
    model = get_model()
    st.session_state.model_loaded = True
    loader_container.empty()
    st.rerun()
else:
    model = get_model()

# -----------------------------
# CALLBACK FUNCTIONS
# -----------------------------
def swap_languages():
    st.session_state.src_lang, st.session_state.tgt_lang = (
        st.session_state.tgt_lang,
        st.session_state.src_lang
    )

def toggle_theme():
    st.session_state.theme = "dark" if st.session_state.theme == "light" else "light"

# -----------------------------
# UI
# -----------------------------
st.title("Kurdish ↔ English Translator (NLLB + LoRA)")

col1, colSwap, colTheme, col2 = st.columns([1, 0.6, 0.8, 1])

with col1:
    src_lang = st.selectbox(
        "From:",
        ["English", "Kurdish"],
        key="src_lang"
    )

with colSwap:
    st.button("⇆ Swap", use_container_width=True, on_click=swap_languages)

with colTheme:
    theme_icon = "🌙" if THEME == "light" else "☀️"
    st.button(theme_icon, use_container_width=True, on_click=toggle_theme)

with col2:
    tgt_lang = st.selectbox(
        "To:",
        ["Kurdish", "English"],
        key="tgt_lang"
    )

lang_codes = {
    "English": "eng_Latn",
    "Kurdish": "ckb_Arab"
}

src_code = lang_codes[st.session_state.src_lang]
tgt_code = lang_codes[st.session_state.tgt_lang]

text = st.text_area("Enter text:", height=180)

# -----------------------------
# TRANSLATE BUTTON
# -----------------------------
if st.button("Translate", type="primary", use_container_width=True):
    if not text.strip():
        st.warning("Please enter some text.")
    else:
        # Show loader
        loader_placeholder = st.empty()
        loader_placeholder.markdown(
            "<div class='dots-loader'><div></div><div></div><div></div></div>",
            unsafe_allow_html=True
        )
        
        try:
            output = translate(src_code, tgt_code, model, text)
            st.session_state.output_text = output
        except Exception as e:
            st.error(f"Translation error: {str(e)}")
            st.session_state.output_text = ""
        finally:
            # Remove loader
            loader_placeholder.empty()

# Display output if available
if st.session_state.output_text:
    st.subheader("Output")
    st.markdown(
        f"<div class='output-box'>{st.session_state.output_text}</div>",
        unsafe_allow_html=True
    )
    
    # Copy section
    st.code(st.session_state.output_text, language=None)
    st.caption("👆 Select and copy the text above (Ctrl+C / Cmd+C)")