"""
Chat demo for local LLMs using Streamlit.


Run with:
```
streamlit run chat.py --server.address 0.0.0.0
```
"""

import logging
import os

import openai
import regex
import streamlit as st

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def convert_latex_brackets_to_dollars(text):
    """Convert LaTeX bracket notation to dollar notation for Streamlit."""

    def replace_display_latex(match):
        return f"\n<bdi> $$ {match.group(1).strip()} $$ </bdi>\n"

    text = regex.sub(r"(?r)\\\[\s*([^\[\]]+?)\s*\\\]", replace_display_latex, text)

    def replace_paren_latex(match):
        return f" <bdi> $ {match.group(1).strip()} $ </bdi> "

    text = regex.sub(r"(?r)\\\(\s*(.+?)\s*\\\)", replace_paren_latex, text)

    return text


# (CSS injection moved below and applied conditionally based on st.session_state.lang)


@st.cache_resource
def openai_configured():
    return {
        "model": os.getenv("MY_MODEL", "Intel/hebrew-math-tutor-v1"),
        "api_base": os.getenv("AWS_URL", "http://localhost:8111/v1"),
        "api_key": os.getenv("MY_KEY"),
    }


config = openai_configured()


@st.cache_resource
def get_client():
    return openai.OpenAI(api_key=config["api_key"], base_url=config["api_base"])


client = get_client()

# Language toggle state: 'he' (Hebrew) or 'en' (English)
if "lang" not in st.session_state:
    st.session_state.lang = "he"

# Localized UI strings
labels = {
    "he": {
        "title": "מתמטיבוט 🧮",
        "intro": """
ברוכים הבאים לדמו! 💡 כאן תוכלו להתרשם **ממודל השפה החדש** שלנו; מודל בגודל 4 מיליארד פרמטרים שאומן לענות על שאלות מתמטיות בעברית, על המחשב שלכם, ללא חיבור לרשת.

קישור למודל, פרטים נוספים, יצירת קשר ותנאי שימוש:

https://huggingface.co/Intel/hebrew-math-tutor-v1

-----
""",
        "select_label": "בחרו שאלה מוכנה או צרו שאלה חדשה:",
        "new_question": "שאלה חדשה...",
        "text_label": "שאלה:",
        "placeholder": "הזינו את השאלה כאן...",
        "send": "שלח",
        "reset": "שיחה חדשה",
        "toggle_to": "English 🇬🇧",
        "predefined": [
            "שאלה חדשה...",
            " מהו סכום הסדרה הבאה:  1 + 1/2 + 1/4 + 1/8 + ...",
            "פתח את הביטוי: (a-b)^4",
            "פתרו את המשוואה הבאה: sin(2x) = 0.5",
        ],
    },
    "en": {
        "title": "MathBot 🧮",
        "intro": """
Welcome to the demo! 💡 Here you can try our **new language model** — a 4-billion-parameter model trained to answer math questions in Hebrew while maintaining its English capabilities. It runs locally on your machine without requiring an internet connection.

For the model page and more details see:

https://huggingface.co/Intel/hebrew-math-tutor-v1

-----
""",
        "select_label": "Choose a prepared question or create a new one:",
        "new_question": "New question...",
        "text_label": "Question:",
        "placeholder": "Type your question here...",
        "send": "Send",
        "reset": "New Conversation",
        "toggle_to": "עברית 🇮🇱",
        "predefined": [
            "New question...",
            "What is the sum of the series: 1 + 1/2 + 1/4 + 1/8 + ...",
            "Expand the expression: (a-b)^4",
            "Solve the equation: sin(2x) = 0.5",
        ],
    },
}

L = labels[st.session_state.lang]

# Inject language-specific CSS so alignment follows the current UI language
if st.session_state.lang == "he":
    st.markdown(
        """
    <style>
    /* RTL: apply to Streamlit content */
    .stText, .stTextArea textarea, .stTextArea label, .stSelectbox select, .stSelectbox label, .stSelectbox div,
    select, option, [data-testid="stSelectbox"] select, [data-testid="stSelectbox"] option, .stSelectbox > div, .stSelectbox > div > div,
    .stSelectbox [role="listbox"], .stSelectbox [role="option"], div[role="listbox"], div[role="option"] {
        direction: rtl !important;
        text-align: right !important;
    }
    .stChatMessage { direction: rtl !important; text-align: right !important; }
    h1, .stTitle, [data-testid="stHeader"] h1 { direction: rtl !important; text-align: right !important; }
    .stMarkdown p:not(:has(.MathJax)):not(:has(mjx-container)):not(:has(.katex)) { direction: rtl !important; text-align: right !important; }
    .stMarkdown code, .stMarkdown pre { direction: ltr !important; text-align: left !important; }
    .stButton button { direction: rtl !important; }
    </style>
    """,
        unsafe_allow_html=True,
    )
else:
    # Ensure default LTR for English mode (override any residual RTL rules)
    st.markdown(
        """
    <style>
    .stText, .stTextArea textarea, .stTextArea label, .stSelectbox select, .stSelectbox label, .stSelectbox div,
    select, option, [data-testid="stSelectbox"] select, [data-testid="stSelectbox"] option, .stSelectbox > div, .stSelectbox > div > div,
    .stSelectbox [role="listbox"], .stSelectbox [role="option"], div[role="listbox"], div[role="option"] {
        direction: ltr !important;
        text-align: left !important;
    }
    .stChatMessage { direction: ltr !important; text-align: left !important; }
    h1, .stTitle, [data-testid="stHeader"] h1 { direction: ltr !important; text-align: left !important; }
    .stMarkdown code, .stMarkdown pre { direction: ltr !important; text-align: left !important; }
    .stButton button { direction: ltr !important; }
    </style>
    """,
        unsafe_allow_html=True,
    )

# Localized strings/templates for thinking/details and final answer
if st.session_state.lang == "he":
    _dir = "rtl"
    _align = "right"
    _summary_text = "לחץ כדי לראות את תהליך החשיבה"
    _thinking_prefix = "🤔 חושב"
    _thinking_done = "🤔 *תהליך החשיבה הושלם, מכין תשובה...*"
    _final_label = "📝 תשובה סופית:"
else:
    _dir = "ltr"
    _align = "left"
    _summary_text = "Click to view the thinking process"
    _thinking_prefix = "🤔 Thinking"
    _thinking_done = "🤔 *Thinking complete, preparing answer...*"
    _final_label = "📝 Final answer:"

# Helper HTML template for the collapsible thinking/details block
_details_template = (
    '<details dir="{dir}" style="text-align: {align};">'
    "<summary>🤔 <em>{summary}</em></summary>"
    '<div style="white-space: pre-wrap; margin: 10px 0; direction: {dir}; text-align: {align};">{content}</div>'
    "</details>"
)

st.title(L["title"])

st.markdown(L["intro"])

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Predefined options
predefined_options = L["predefined"]

# Dropdown for predefined options
selected_option = st.selectbox(L["select_label"], predefined_options)

# Text area for input
if selected_option == L["new_question"]:
    user_input = st.text_area(
        L["text_label"], height=100, key="user_input", placeholder=L["placeholder"]
    )
else:
    user_input = st.text_area(L["text_label"], height=100, key="user_input", value=selected_option)

# Buttons layout: Reset | Language Toggle | Send
col_left, col_mid, col_right = st.columns([4, 2, 4])

with col_left:
    if st.button(L["reset"], type="secondary", use_container_width=True):
        st.session_state.chat_history = []
        st.rerun()

with col_mid:
    # Button shows the language to switch TO (e.g. 'English' when current is Hebrew)
    if st.button(L["toggle_to"], use_container_width=True):
        st.session_state.lang = "en" if st.session_state.lang == "he" else "he"
        st.rerun()

with col_right:
    # Guard against None from text_area and ensure non-empty trimmed input
    send_clicked = st.button(L["send"], type="primary", use_container_width=True) and (
        user_input and user_input.strip()
    )

if send_clicked:
    st.session_state.chat_history.append(("user", user_input))

    # Create a placeholder for streaming output
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""

        # System prompt - adapt to UI language; do not force Hebrew when UI is English
        if st.session_state.lang == "he":
            system_prompt = """\
You are a helpful AI assistant specialized in mathematics and problem-solving who can answer math questions with the correct answer.
Answer shortly, not more than 500 tokens, but outline the process step by step.
Answer ONLY in Hebrew!
"""
        else:
            system_prompt = """\
You are a helpful AI assistant specialized in mathematics and problem-solving who can answer math questions with the correct answer.
Answer shortly, not more than 500 tokens, but outline the process step by step.
"""

        # Create messages in proper chat format
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_input},
        ]

        # Build a single string prompt for OpenAI-compatible chat API
        # Keep the special thinking tokens (<think>...</think>) if the remote model supports them
        prompt_messages = messages

        # Stream from OpenAI-compatible API (vllm remote exposing openai-compatible endpoint)
        # Use the chat completions streaming interface
        in_thinking = True
        thinking_content = "<think>"
        final_answer = ""

        try:
            # openai.ChatCompletion.create with stream=True yields chunks with 'choices'
            stream = client.chat.completions.create(
                messages=prompt_messages,
                model=config["model"],
                temperature=0.6,
                max_tokens=2000,
                top_p=0.95,
                stream=True,
                extra_body={"top_k": 20},
            )

            for chunk in stream:
                # Each chunk is a dict; text delta at chunk['choices'][0]['delta'] for newer APIs
                delta = ""
                try:
                    # compatible with OpenAI response structure
                    delta = chunk.choices[0].delta.content
                except Exception:
                    # fallback for older/other shapes; use getattr to avoid dict-specific calls
                    delta = getattr(chunk, "text", None) or "HI "

                if not delta:
                    continue

                full_response += delta

                # Handle thinking markers
                if "<think>" in delta:
                    in_thinking = True

                if in_thinking:
                    thinking_content += delta
                    if "</think>" in delta:
                        in_thinking = False
                        thinking_text = (
                            thinking_content.replace("<think>", "").replace("</think>", "").strip()
                        )
                        display_content = _details_template.format(
                            dir=_dir, align=_align, summary=_summary_text, content=thinking_text
                        )
                        message_placeholder.markdown(display_content + "▌", unsafe_allow_html=True)
                    else:
                        dots = "." * ((len(thinking_content) // 10) % 6)
                        # thinking indicator
                        thinking_indicator = f"""
<div dir="{_dir}" style="padding: 10px; background-color: #f0f2f6; border-radius: 10px; border-right: 4px solid #1f77b4; text-align: {_align};">
    <p style="margin: 0; color: #1f77b4; font-style: italic;">
        {_thinking_prefix}{dots}
    </p>
</div>
"""
                        message_placeholder.markdown(thinking_indicator, unsafe_allow_html=True)
                else:
                    # Final answer streaming
                    final_answer += delta
                    converted_answer = convert_latex_brackets_to_dollars(final_answer)
                    message_placeholder.markdown(
                        f"{_thinking_done}\n\n**{_final_label}**\n\n" + converted_answer + "▌",
                        unsafe_allow_html=True,
                    )
        except Exception as e:
            # Show an error to the user
            message_placeholder.markdown(f"**Error contacting remote model:** {e}")

        # Final rendering: if there was thinking content include it
        if thinking_content and "</think>" in thinking_content:
            thinking_text = thinking_content.replace("<think>", "").replace("</think>", "").strip()
            message_placeholder.empty()
            with message_placeholder.container():
                thinking_html = _details_template.format(
                    dir=_dir, align=_align, summary=_summary_text, content=thinking_text
                )
                st.markdown(thinking_html, unsafe_allow_html=True)
                st.markdown(
                    f'<div dir="{_dir}" style="text-align: {_align}; margin: 10px 0;"><strong>{_final_label}</strong></div>',
                    unsafe_allow_html=True,
                )
                converted_answer = convert_latex_brackets_to_dollars(final_answer or full_response)
                st.markdown(converted_answer, unsafe_allow_html=True)
        else:
            converted_response = convert_latex_brackets_to_dollars(final_answer or full_response)
            message_placeholder.markdown(converted_response, unsafe_allow_html=True)

    st.session_state.chat_history.append(("assistant", final_answer or full_response))