import streamlit as st
import complete
import wisdom_extract
import garood

st.markdown(
"""
<style>
body, html {
    direction: RTL;
    # unicode-bidi: bidi-override;
    # text-align: right;
}
p, div, input, label, h1, h2, h3, h4, h5, h6 {
    direction: RTL;
    # unicode-bidi: bidi-override;
    # text-align: right;
}
code {
    font-family: 'Courier New', monospace;
    font-weight: bold;
}
</style>
""", unsafe_allow_html=True)


def tab_idx():
    if 'tab_idx' not in st.session_state:
        st.session_state.tab_idx = 1
    return st.session_state.tab_idx
def tab_idx_set(idx):
    if idx != tab_idx():
        st.session_state.tab_idx = idx

@st.cache_resource
def sim_model():
    return wisdom_extract.hf_model()

POETS = {
    "sample/1.wis.txt": "زهير بن أبي سلمى",
    "sample/2.wis.txt": "المتنبي",
}
@st.cache_data
def wisdom():
    return wisdom_extract.read_wisdom(
        sim_model(), POETS.keys())


cols = st.columns([1, 1, 1])
for i, txt in enumerate(["بحث", "وزن", "إتمام"]):
    with cols[i]:
        st.button(txt, use_container_width=True,
                on_click=tab_idx_set, args=(i,),
                disabled=i==tab_idx())

if tab_idx() == 0:
    # READ INPUT
    query = "ما يحلّ سلمياً خير مما يحل حربياً"
    query = st.text_input(
        label="إبحث في الشعر",
        value=query,
    )

    with st.spinner('جاري تحميل النموذج...'):
        sim_model()
    with st.spinner('جاري تجهيز الأشعار المحفوظة...'):
        _, verses, meaning_embed, reason, srcs = wisdom()
    with st.spinner('جاري تحليل عبارة البحث...'):
        queries = wisdom_extract.query_transform(query)
    with st.spinner('جاري البحث في الأشعار...'):
        queries_embed = sim_model().encode(queries)
        closest_poems = wisdom_extract.closest_embedding(
                sim_model(), queries_embed, verses, meaning_embed, reason, srcs)

    table = []
    for score, verse, reason, src in closest_poems:
        table.append([verse, f"{score * 100:.1f}%", reason, src])

    for idx, val in enumerate(table[:5]):
        verse, score, reason, src = val
        cols = st.columns([.5, 6.5, 2.5, 1, .5])
        with cols[0]:
            st.write(idx)
        with cols[1]:
            st.write(verse)
        with cols[2]:
            st.write(f"– {POETS[src]}")
        with cols[3]:
            st.write(score)
        with cols[4]:
            st.markdown("", help="\n\n".join(r[1] for i, r in enumerate(reason) if i < 3))


if tab_idx() == 1:
    BH = garood.config()['bahr']['kamil_1']
    query = "هل غادر الشعراء من متردم"
    cols = st.columns([6, 3, 3])
    with cols[0]:
        query = st.text_input(
            label="قم بوزن الشطر",
            value=query,
        )

    with cols[1]:
        st.selectbox("البحر", ["الكامل"], key="tab1")
    cols = st.columns([6, 3, 3])
    with cols[0]:
        with st.spinner('جاري إضافة التشكيل...'):
            ARs = complete.shakk([query] * 3)

        for AR in ARs:
            EN, HR = garood.en_gen_enhr(garood.ar_gen_en(AR))
            c, s = garood.bh_gen_best(BH, EN, HR, True)
            if c >= 0: break
        st.code(f"{query}\n{AR}\n", language=None)
        if c >= 0:
            g = garood.state_to_str(s, EN)
            st.code(f"{garood.en2ar(g)}\n{garood.g2v(g)}\n", language=None)
        else:
            st.error("الشطر لا يلتزم ببحر الكامل", icon="🚨")


if tab_idx() == 2:
    BH = garood.config()['bahr']['kamil_1']
    query = "كم حدث الأصحاب"
    cols = st.columns([6, 3, 3])
    with cols[0]:
        query = st.text_input(
            label="قم بإتمام الشطر",
            value=query,
        )
    with cols[1]:
        st.selectbox("البحر", ["الكامل"], key="tab2")
    with cols[0]:
        with st.spinner('جاري توليد الإقتراحات...'):
            poss = complete.gen(query)
        with st.spinner(f"جاري التصفية من {len(poss)} اقتراحا..."):
            pshak = complete.shakk(poss)
    for i, val in enumerate(complete.filter(pshak)):
        c, p, s, g = val
        cols = st.columns([1, 4, 7])
        with cols[0]:
            st.write(i)
        with cols[1]:
            st.write(p)
        with cols[2]:
            st.code(f"{garood.en2ar(g)}\n{garood.g2v(g)}\n", language=None)