import re
import urllib.parse

import json
import requests
import streamlit as st
from streamlit.components.v1 import html
from streamlit_extras.stylable_container import stylable_container

url = "https://app.premai.io/v1/chat/completions"

st.title("Prem Chat UI")

if "api_key" not in st.session_state or "project_id" not in st.session_state:
    api_key = st.text_input("Enter your API Key", type="password")
    project_id = st.text_input("Enter your Project ID")
    
    if not api_key or not project_id:
        st.warning("Please enter your API key and Project ID to use the app.")
        st.stop()

    if not project_id.isdecimal():
        st.warning("Please enter your Project ID correctly.")
        st.stop()
    
    if not api_key.isascii():
        st.warning("Please enter your API key correctly.")
        st.stop()
    
    st.session_state.api_key = api_key
    st.session_state.project_id = int(project_id)
    st.rerun()

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

if "prefill" not in st.session_state:
    st.session_state.prefill = ""

def get_ai_response(messages):
    st.session_state.is_streaming = True
    st.session_state.response = ""
    shown_message = ""

    st.session_state.prefill = st.session_state.prefill.strip()
    if st.session_state.prefill:
        messages.append({"role": "assistant", "content": st.session_state.prefill})
        st.session_state.response += st.session_state.prefill
        shown_message = st.session_state.prefill.replace("\n", "  \n")

    with st.chat_message("assistant", avatar=st.session_state.assistant_avatar):

        payload = {
            "project_id": st.session_state.project_id,
            "messages": messages,
            "model": model,
            "system_prompt": system_prompt,
            "max_tokens": 4000,
            "stream": True,
            "temperature": temperature,
        }
        headers = {
            "Authorization": st.session_state.api_key,
            "Content-Type": "application/json"
        }

        placeholder = st.empty()

        with stylable_container(
            key="stop_generating",
            css_styles="""
                button {
                    position: fixed;
                    bottom: 100px;
                    left: 50%;
                    transform: translateX(-50%);
                    z-index: 1;
                }
                """,
        ):
            st.button("Stop generating")

        response = requests.post(url, headers=headers, json=payload, stream=True)

        for chunk in response.iter_lines(decode_unicode=True):
            if chunk:
                if chunk.startswith("data: "):
                    data_str = chunk[len("data: "):].strip()
                    if data_str == "[DONE]":
                        break
                    try:
                        sse_data = json.loads(data_str)
                        if sse_data.get("choices") and sse_data["choices"][0].get("delta", {}).get("content"):
                            content = sse_data["choices"][0]["delta"]["content"]

                            st.session_state.response += content
                            shown_message += (
                                content
                                .replace("\n", "  \n")
                                .replace("<", "\\<")
                                .replace(">", "\\>")
                            )
                            placeholder.markdown(shown_message)

                    except json.JSONDecodeError:
                        pass


    if st.session_state.prefill == st.session_state.response:
        st.session_state.is_error = True
    else:
        st.session_state.is_error = False

    st.session_state.is_streaming = False
    return st.session_state.response

def normalize_code_block(match):
    return match.group(0).replace("  \n", "\n")\
                         .replace("\\<", "<")\
                         .replace("\\>", ">")

def normalize_inline(match):
    return match.group(0).replace("\\<", "<")\
                         .replace("\\>", ">")

code_block_pattern = r"(```.*?```)"
inline_pattern = r"`([^`\n]+?)`"

def display_messages():
    for i, message in enumerate(st.session_state.messages):
        if message["role"] == "user":
            avatar = st.session_state.user_avatar
        else:
            avatar = st.session_state.assistant_avatar
        
        with st.chat_message(message["role"], avatar=avatar):
            shown_message = message["content"].replace("\n", "  \n")\
                                              .replace("<", "\\<")\
                                              .replace(">", "\\>")
            if "```" in shown_message:
                # Replace "  \n" with "\n" within code blocks
                shown_message = re.sub(code_block_pattern, normalize_code_block, shown_message, flags=re.DOTALL)
            if "`" in shown_message:
                shown_message = re.sub(inline_pattern, normalize_inline, shown_message)
            st.markdown(shown_message)
            
            col1, col2, col3, col4 = st.columns([1, 1, 1, 1])
            with col1:
                if st.button("Edit", key=f"edit_{i}_{len(st.session_state.messages)}"):
                    st.session_state.edit_index = i
                    st.rerun()
            with col2:
                if st.session_state.is_delete_mode and st.button("Delete", key=f"delete_{i}_{len(st.session_state.messages)}"):
                    del st.session_state.messages[i]
                    st.rerun()
            with col3:
                text_to_copy = message["content"]
                # Encode the string to escape
                text_to_copy_escaped = urllib.parse.quote(text_to_copy)

                copy_button_html = f"""
                <button id="copy-msg-btn-{i}" style='font-size: 1em; padding: 0.5em;' onclick='copyMessage("{i}")'>Copy</button>
                
                <script>
                function copyMessage(index) {{
                    navigator.clipboard.writeText(decodeURIComponent("{text_to_copy_escaped}"));
                    let copyBtn = document.getElementById("copy-msg-btn-" + index);
                    copyBtn.innerHTML = "Copied!";
                    setTimeout(function(){{ copyBtn.innerHTML = "Copy"; }}, 2000);
                }}
                </script>
                """
                html(copy_button_html, height=50)

            if i == len(st.session_state.messages) - 1 and message["role"] == "assistant":
                with col4:
                    if st.button("Retry", key=f"retry_{i}_{len(st.session_state.messages)}"):
                        if len(st.session_state.messages) >= 2:
                            del st.session_state.messages[-1]
                            st.session_state.retry_flag = True
                            st.rerun()
        
        if "edit_index" in st.session_state and st.session_state.edit_index == i:
            with st.form(key=f"edit_form_{i}_{len(st.session_state.messages)}"):
                new_content = st.text_area("Edit message", height=200, value=st.session_state.messages[i]["content"])
                col1, col2 = st.columns([1, 1])
                with col1:
                    if st.form_submit_button("Save"):
                        st.session_state.messages[i]["content"] = new_content
                        del st.session_state.edit_index
                        st.rerun()
                with col2:
                    if st.form_submit_button("Cancel"):
                        del st.session_state.edit_index
                        st.rerun()
    
    if "is_error" in st.session_state and st.session_state.is_error:
        st.error("""
                   Something went wrong. To resolve this error:
                   1. Use the Retry button.
                   2. Update your API key or Project ID correctly.
                   3. Edit any messages that contain only whitespace characters (e.g. spaces, tabs, newlines).
                   4. Remove any consecutive messages with the same role (the same icon).
                   5. Check the Traces within the project you are using to see the details of the error.
                   """)
    
# Add sidebar for advanced settings
with st.sidebar:
    settings_tab, appearance_tab = st.tabs(["Settings", "Appearance"])

    with settings_tab:
        st.markdown("Help (Japanese): https://rentry.org/9hgneofz")
    
        # Copy Conversation History button
        log_text = ""
        for message in st.session_state.messages:
            if message["role"] == "user":
                log_text += "<USER>\n"
                log_text += message["content"] + "\n\n"
            else:
                log_text += "<ASSISTANT>\n"
                log_text += message["content"] + "\n\n"
        log_text = log_text.rstrip("\n")
            
        # Encode the string to escape
        log_text_escaped = urllib.parse.quote(log_text)
        
        copy_log_button_html = f"""
        <button id="copy-log-btn" style='font-size: 1em; padding: 0.5em;' onclick='copyLog()'>Copy Conversation History</button>
        
        <script>
        function copyLog() {{
            navigator.clipboard.writeText(decodeURIComponent("{log_text_escaped}"));
            let copyBtn = document.getElementById("copy-log-btn");
            copyBtn.innerHTML = "Copied!";
            setTimeout(function(){{ copyBtn.innerHTML = "Copy Conversation History"; }}, 2000);
        }}
        </script>
        """
        html(copy_log_button_html, height=50)

        if st.session_state.get("is_history_shown") != True:
            if st.button("Display History as Code Block"):
                st.session_state.is_history_shown = True
                st.rerun()
        else:
            if st.button("Hide History"):
                st.session_state.is_history_shown = False
                st.rerun()
            st.code(log_text)
            
        st.session_state.is_delete_mode = st.toggle("Enable Delete button")
    
        st.header("Advanced Settings")
        model_list = ["claude-3.5-haiku",
                      "claude-3.5-sonnet",
                      "claude-3.5-sonnet-v2",
                      "claude-3-haiku",
                      "claude-3-opus",
                      "claude-3-sonnet",
                      "deepseek-r1",
                      "deepseek-r1-distill-llama-70b",
                      "gpt-4-eu",
                      "gpt-4o",
                      "gpt-4o-eu",
                      "gpt-4o-mini",
                      "gpt-4o-mini-eu",
                      "gpt-4-turbo",
                      "llama-3.1-8b",
                      "llama-3.2-1b",
                      "llama-3.2-3b",
                      "llama-3.3-70b",
                      "llama-3-70b",
                      "llama-3-8b",
                      "llama-3-8b-guard",
                      "prem-llama-3.1-8b",
                      "prem-llama-3.2-1b",
                      "prem-llama-3.2-3b",
                      "mixtral-8x7b",
                      ]
        model = st.selectbox("Model", options=model_list, index=0)
        system_prompt = st.text_area("System prompt", height=200)
    
        st.session_state.prefill = st.text_area("Prefill", height=68, value=st.session_state.prefill, placeholder="It only works well with the Claude models.",
                                                help="You can prefill the assistant's responses. You can also directly type the @prefill command into the chat field (e.g., \"Write a novel. @prefill Sure! I'd be happy to write a novel for you.\")")
        save_prefill = st.toggle("Save the @prefill command input in the sidebar", value=True)
        
        temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=1.0, step=0.1)
        top_p = st.slider("Top-P", min_value=0.01, max_value=1.00, value=1.00, step=0.01)
        penalty_type = st.selectbox("Penalty Type", options=["Frequency Penalty", "Presence Penalty"])
        penalty_value = st.slider("Penalty Value", min_value=0.0, max_value=1.0, value=0.0, step=0.1)
        
        st.header("Restore History")
        history_input = st.text_area("Paste conversation history:", height=200)
        if st.button("Restore History"):
            st.session_state.messages = []
            messages = re.split(r"^(<USER>|<ASSISTANT>)\n", history_input, flags=re.MULTILINE)
            role = None
            text = ""
            for message in messages:
                if message.strip() in ["<USER>", "<ASSISTANT>"]:
                    if role and text:
                        st.session_state.messages.append({"role": role, "content": text.strip()})
                        text = ""
                    role = "user" if message.strip() == "<USER>" else "assistant"
                else:
                    text += message
            if role and text:
                st.session_state.messages.append({"role": role, "content": text.strip()})
            st.rerun()
    
        st.header("Clear History")
        if st.button("Clear Chat History"):
            st.session_state.messages = []
            st.rerun()
        
        st.header("Change API Key or Project ID")
        new_api_key = st.text_input("Enter new API Key", type="password")
        if st.button("Update API Key"):
            if new_api_key:
                st.session_state.api_key = new_api_key
                st.success("API Key updated successfully!")
            else:
                st.warning("Please enter a valid API Key.")
    
        new_project_id = st.text_input("Enter new Project ID")
        if st.button("Update Project ID"):
            if new_project_id and new_project_id.isdecimal():
                st.session_state.project_id = int(new_project_id)
                st.success("Project ID updated successfully!")
            else:
                st.warning("Please enter a valid Project ID.")
    
    with appearance_tab:
        st.header("Font Selection")
        font_options = {
            "Zen Maru Gothic": "Zen Maru Gothic",
            "Noto Sans JP": "Noto Sans JP",
            "Sawarabi Mincho": "Sawarabi Mincho"
        }
        selected_font = st.selectbox("Choose a font", ["Default"] + list(font_options.keys()))

        st.header("Change the font size")
        st.session_state.font_size  = st.slider("Font size", min_value=16.0, max_value=50.0, value=16.0, step=1.0)
        
        st.header("Change the user's icon")
        st.session_state.user_avatar = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg", "webp", "gif", "bmp", "svg",], key="user_avatar_uploader")
    
        st.header("Change the assistant's icon")
        st.session_state.assistant_avatar = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg", "webp", "gif", "bmp", "svg",], key="assistant_avatar_uploader")

        st.header("Change the icon size")
        st.session_state.avatar_size  = st.slider("Icon size", min_value=2.0, max_value=20.0, value=2.0, step=0.2)
        
# After Stop generating
if st.session_state.get("is_streaming"):
    st.session_state.messages.append({"role": "assistant", "content": st.session_state.response})
    st.session_state.is_error = False
    st.session_state.is_streaming = False
    if "retry_flag" in st.session_state and st.session_state.retry_flag:
        st.session_state.retry_flag = False
    st.rerun()

# Change the font
if selected_font != "Default":
    with open("style.css") as css:
        st.markdown(f'<style>{css.read()}</style>', unsafe_allow_html=True)
    st.markdown(f'<style>body * {{ font-family: "{font_options[selected_font]}", serif !important; }}</style>', unsafe_allow_html=True)

# Change font size
st.markdown(f'<style>[data-testid="stChatMessageContent"] .st-emotion-cache-kj6hex p{{font-size: {st.session_state.font_size}px;}}</style>', unsafe_allow_html=True)

# Change icon size
# (CSS element names may be subject to change.)
# (Contributor: ★31 >>538)
AVATAR_SIZE_STYLE = f"""
                <style>
                    [data-testid="stChatMessageAvatarUser"] {{
                        width: {st.session_state.avatar_size}rem;
                        height: {st.session_state.avatar_size}rem;
                    }}
                    [data-testid="stChatMessageAvatarAssistant"] {{
                        width: {st.session_state.avatar_size}rem;
                        height: {st.session_state.avatar_size}rem;
                    }}
                    [data-testid="stChatMessage"] .st-emotion-cache-1pbsqtx {{
                        width: {st.session_state.avatar_size / 1.6}rem;
                        height: {st.session_state.avatar_size / 1.6}rem;
                    }}
                    [data-testid="stChatMessage"] .st-emotion-cache-p4micv {{
                        width: {st.session_state.avatar_size}rem;
                        height: {st.session_state.avatar_size}rem;
                    }}
                </style>
"""
st.markdown(AVATAR_SIZE_STYLE, unsafe_allow_html=True)

display_messages()

# After Retry
if st.session_state.get("retry_flag"):
    if len(st.session_state.messages) > 0:
        messages = st.session_state.messages.copy()
        response = get_ai_response(messages)
        st.session_state.messages.append({"role": "assistant", "content": response})

        st.session_state.retry_flag = False
        st.rerun()
            
    else:
        st.session_state.retry_flag = False

if prompt := st.chat_input("Enter your message here..."):
    used_prefill = False
    prefill_pattern = r"([@＠](prefill|ぷれふぃる|プレフィル)\s?(.*))"
    prefill_match = re.search(prefill_pattern, prompt)

    if prefill_match:
        used_prefill = True
        if not save_prefill:
            original_prefill = st.session_state.prefill
        st.session_state.prefill = prefill_match.group(3)
        prompt = prompt.replace(prefill_match.group(1), '')

    st.session_state.messages.append({"role": "user", "content": prompt})
    messages = st.session_state.messages.copy()

    shown_message = prompt.replace("\n", "  \n")\
                          .replace("<", "\\<")\
                          .replace(">", "\\>")

    with st.chat_message("user", avatar=st.session_state.user_avatar):
        st.write(shown_message)

    response = get_ai_response(messages)

    st.session_state.messages.append({"role": "assistant", "content": response})

    if used_prefill and not save_prefill:
        st.session_state.prefill = original_prefill
    
    st.rerun()