"""
Netra Translate — Streamlit demo for the netra-nmt English ↔ Khmer model.

A two-pane translation UI (mirroring the bundled FastAPI web app) intended for
deployment on Hugging Face Spaces. All inference is delegated to the
``NetraTranslator`` class from the installed ``netra-nmt`` package.
"""

from __future__ import annotations

from pathlib import Path

import streamlit as st

# ---------------------------------------------------------------------------
# Page config + branding
# ---------------------------------------------------------------------------

_ASSETS = Path(__file__).resolve().parent / "assets"
_WORDMARK = _ASSETS / "netra-wordmark-orange.png"

st.set_page_config(
    page_title="Netra Translate",
    page_icon="🪔",
    layout="wide",
)

LABELS = {"en": "English", "km": "Khmer"}

# Netra palette + Khmer webfont, adapted from netra_nmt/static/index.html.
st.markdown(
    """
    <style>
      @import url('https://fonts.googleapis.com/css2?family=Noto+Sans+Khmer:wght@400;500;600;700&display=swap');

      :root {
        --primary: #e75f1b;
        --primary-dark: #c84e13;
        --gold: #f6c416;
        --muted: #8a7872;
        --border: #ece4e0;
      }

      html, body, [class*="css"] { font-family: "Noto Sans Khmer", system-ui, -apple-system, sans-serif; }

      /* Primary buttons in Netra orange. */
      .stButton > button[kind="primary"] {
        background: var(--primary);
        border-color: var(--primary);
        font-weight: 600;
      }
      .stButton > button[kind="primary"]:hover {
        background: var(--primary-dark);
        border-color: var(--primary-dark);
      }

      /* The translation output card. */
      .netra-output {
        border: 1px solid var(--border);
        border-radius: 12px;
        padding: 18px 20px;
        min-height: 180px;
        font-size: 19px;
        line-height: 1.7;
        white-space: pre-wrap;
        word-break: break-word;
      }
      .netra-output.placeholder { color: var(--muted); }

      .netra-langlabel { font-weight: 600; color: var(--muted); font-size: 15px; }
    </style>
    """,
    unsafe_allow_html=True,
)


# ---------------------------------------------------------------------------
# Model (loaded once per Space process)
# ---------------------------------------------------------------------------

@st.cache_resource(show_spinner="Loading Netra model… (first run downloads weights)")
def get_translator():
    from netra_nmt import NetraTranslator

    # Spaces free tier is CPU-only; this also works locally on a GPU box.
    return NetraTranslator(device="cpu")


# ---------------------------------------------------------------------------
# Session state
# ---------------------------------------------------------------------------

if "direction" not in st.session_state:
    st.session_state.direction = "en2km"
if "source" not in st.session_state:
    st.session_state.source = ""
if "output" not in st.session_state:
    st.session_state.output = ""


def _src_tgt() -> tuple[str, str]:
    return ("en", "km") if st.session_state.direction == "en2km" else ("km", "en")


def _swap() -> None:
    """Flip the direction and carry the previous output into the input box."""
    st.session_state.direction = (
        "km2en" if st.session_state.direction == "en2km" else "en2km"
    )
    prev = st.session_state.output
    st.session_state.source = prev
    st.session_state.output = ""


def _clear() -> None:
    st.session_state.source = ""
    st.session_state.output = ""


# ---------------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------------

head_left, head_right = st.columns([3, 2], vertical_alignment="center")
with head_left:
    if _WORDMARK.exists():
        st.image(str(_WORDMARK), width=170)
    else:  # pragma: no cover - asset is shipped alongside the app
        st.markdown("### netra")
    st.caption("English ↔ Khmer neural machine translation")

st.divider()

# ---------------------------------------------------------------------------
# Language bar + swap
# ---------------------------------------------------------------------------

src, tgt = _src_tgt()
lang_l, lang_mid, lang_r = st.columns([5, 1, 5], vertical_alignment="center")
with lang_l:
    st.markdown(f"<span class='netra-langlabel'>{LABELS[src]}</span>", unsafe_allow_html=True)
with lang_mid:
    st.button("⇄", on_click=_swap, help="Swap languages", use_container_width=True)
with lang_r:
    st.markdown(
        f"<div style='text-align:right'><span class='netra-langlabel'>{LABELS[tgt]}</span></div>",
        unsafe_allow_html=True,
    )

# ---------------------------------------------------------------------------
# Two panes: source (left) / output (right)
# ---------------------------------------------------------------------------

pane_l, pane_r = st.columns(2)
with pane_l:
    source = st.text_area(
        "Source",
        key="source",
        height=200,
        placeholder="Enter text",
        label_visibility="collapsed",
    )
    st.caption(f"{len(source)} characters")

with pane_r:
    out = st.session_state.output
    if out:
        st.markdown(f"<div class='netra-output'>{out}</div>", unsafe_allow_html=True)
    else:
        st.markdown(
            "<div class='netra-output placeholder'>Translation</div>",
            unsafe_allow_html=True,
        )
    if out:
        st.code(out, language=None)  # gives a one-click copy affordance

# ---------------------------------------------------------------------------
# Controls
# ---------------------------------------------------------------------------

ctrl_l, ctrl_mid, ctrl_r = st.columns([3, 2, 2], vertical_alignment="bottom")
with ctrl_l:
    mode_label = st.selectbox(
        "Decoding",
        ["Greedy (fast)", "Beam search (quality)"],
    )
    mode = "greedy" if mode_label.startswith("Greedy") else "beam"
with ctrl_mid:
    st.button("Clear", on_click=_clear, use_container_width=True)
with ctrl_r:
    translate_clicked = st.button(
        "Translate", type="primary", use_container_width=True
    )

# ---------------------------------------------------------------------------
# Translate
# ---------------------------------------------------------------------------

if translate_clicked:
    text = st.session_state.source.strip()
    if not text:
        st.warning("Enter some text to translate.")
    else:
        try:
            translator = get_translator()
            with st.spinner("Translating…"):
                st.session_state.output = translator.translate(
                    text,
                    direction=st.session_state.direction,
                    mode=mode,
                    max_new_tokens=128,
                )
            st.rerun()
        except Exception as exc:  # noqa: BLE001 - surface any failure to the user
            st.error(f"Translation failed: {exc}")