import anthropic
import streamlit as st
from streamlit.components.v1 import html
from streamlit_extras.stylable_container import stylable_container
import re
import urllib.parse
import traceback
import os

st.title("Claude Chat UI")

if "api_key" not in st.session_state and os.path.exists("api_key.dat"):
    with open("api_key.dat", "r", encoding="utf-8") as f:
        st.session_state.api_key = f.readline().strip()

if "api_key" not in st.session_state:
    api_key = st.text_input("Enter your API Key", type="password")
    
    if not api_key:
        st.warning("Please enter your API key to use the app.")
        st.stop()
    
    if not api_key.isascii():
        st.warning("Please enter your API key correctly.")
        st.stop()
    
    st.session_state.api_key = api_key
    client = anthropic.Anthropic(api_key=api_key)
    st.rerun()
else:
    client = anthropic.Anthropic(api_key=st.session_state.api_key)

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

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

if "use_continue" not in st.session_state:
    st.session_state.use_continue = False

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

if "exception" not in st.session_state:
    st.session_state.exception = None

if "use_thinking" not in st.session_state:
    st.session_state.use_thinking = False

if "thinking_budget_tokens" not in st.session_state:
    st.session_state.thinking_budget_tokens = 1024

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

def count_tokens(text):
    if len(text) > 0:
        response = client.messages.count_tokens(
            model=model,
            messages = [{
                "role": "user",
                "content": text
            }]
        )
        tokens = response.input_tokens
        return tokens
    else:
        return 0
    
def get_truncated_index(messages, limit_val, limit_unit, additional_tokens):
    """ Find index of first message within truncated context

    Args:
        messages (list): Input messages for anthropic API
        limit_val (int): Context length limitation
        limit_unit (str): "Turns" or "Tokens"
        additional_tokens (int): Token count of system prompt

    Returns:
        int: index of first message within truncated context
    """
    if limit_val == 0:
        return 0

    if limit_unit == "Turns":
        # Unit: Turns
        count_turn = 0
        for i in reversed(range(len(messages))):
            if messages[i]["role"] == "user":
                count_turn += 1
                if count_turn == limit_val:
                    return i
        return 0
    else:
        # Unit: Tokens
        last_user_index = len(messages)
        total_tokens = additional_tokens
        for i in reversed(range(len(messages))):
            total_tokens += count_tokens(messages[i]["content"])
            if total_tokens > limit_val:
                return last_user_index
            if messages[i]["role"] == "user":
                last_user_index = i
        return last_user_index

def get_truncated_context(messages, limit_val, limit_unit, additional_tokens):
    """ Return truncated context

    Args:
        messages (list): Input messages for anthropic API
        limit_val (int): Context length limitation
        limit_unit (str): "Turns" or "Tokens"
        additional_tokens (int): Token count of system prompt

    Returns:
        list: Truncated input messages for anthropic API
    """
    first_message_index = get_truncated_index(messages, limit_val, limit_unit, additional_tokens)

    if first_message_index >= len(messages):
        return []
    
    return messages[first_message_index:]

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

    st.session_state.message_continue = st.session_state.message_continue.strip()
    st.session_state.prefill = st.session_state.prefill.strip()
    if st.session_state.message_continue != "":
        messages.append({"role": "assistant", "content": st.session_state.message_continue})
        st.session_state.response += st.session_state.message_continue
        shown_message = st.session_state.message_continue.replace("\n", "  \n")
    elif 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")

    api_messages = []
    for msg in messages:
        api_msg = {"role": msg["role"], "content": msg["content"]}
        api_messages.append(api_msg)

    api_messages = get_truncated_context(api_messages, limit_val, limit_unit, count_tokens(system_prompt))
    st.session_state.exception = None

    try:
        with st.chat_message("assistant", avatar=st.session_state.assistant_avatar): 
            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")

            # Configure thinking parameter
            thinking_param = None
            if st.session_state.use_thinking:
                thinking_param = {
                    "type": "enabled",
                    "budget_tokens": st.session_state.thinking_budget_tokens
                }
                
            # Prepare kwargs for the API call
            stream_kwargs = {
                "messages": api_messages,
                "model": model,
                "max_tokens": max_tokens,
                "system": system_prompt,
                "temperature": temperature,
            }
            
            # Add thinking parameter if enabled
            if thinking_param:
                stream_kwargs["thinking"] = thinking_param
            elif "-4-5-" not in model:
                # Only add top_p and top_k when thinking is disabled and not a 4-5 model
                # Claude 4-5 models don't support temperature + top_p/top_k simultaneously
                stream_kwargs["top_p"] = top_p
                stream_kwargs["top_k"] = top_k
            
            with client.messages.stream(**stream_kwargs) as stream:
                # Stream event handling
                current_block_type = None
                
                for event in stream:
                    if event.type == "content_block_start":
                        current_block_type = event.content_block.type
                    
                    elif event.type == "content_block_delta":
                        if event.delta.type == "thinking_delta" and current_block_type == "thinking":
                            content = str(event.delta.thinking) if event.delta.thinking is not None else ""
                            st.session_state.last_thinking += content
                            shown_message += content.replace("\n", "  \n")\
                                                  .replace("<", "\\<")\
                                                  .replace(">", "\\>")
                            placeholder.markdown(shown_message)
                        elif event.delta.type == "text_delta" and current_block_type == "text":
                            content = str(event.delta.text) if event.delta.text is not None else ""
                            st.session_state.response += content
                            shown_message += content.replace("\n", "  \n")\
                                                  .replace("<", "\\<")\
                                                  .replace(">", "\\>")
                            placeholder.markdown(shown_message)

    except Exception as e:
        st.session_state.exception = e

    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):
            # Add expander for thinking if it exists
            if message["role"] == "assistant" and "thinking" in message and message["thinking"]:
                with st.expander("Show Claude's thinking process"):
                    thinking_text = message["thinking"].replace("\n", "  \n")\
                                                     .replace("<", "\\<")\
                                                     .replace(">", "\\>")
                    if "```" in thinking_text:
                        thinking_text = re.sub(code_block_pattern, normalize_code_block, thinking_text, flags=re.DOTALL)
                    if "`" in thinking_text:
                        thinking_text = re.sub(inline_pattern, normalize_inline, thinking_text)
                    st.markdown(thinking_text)
            
            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)

            if st.session_state.get("show_message_tokens"):
                message_tokens = count_tokens(message["content"])
                st.info(f'Tokens: {message_tokens}')

            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:
                            if st.session_state.get("use_continue"):
                                st.session_state.message_continue = st.session_state.messages[-1]["content"]
                            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 st.session_state.exception:
        st.exception(st.session_state.exception)
    
# 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"
                if "thinking" in message and message["thinking"]:
                    log_text += "<THINKING>\n"
                    log_text += message["thinking"] + "\n</THINKING>\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>
        const log_text_escaped = "{log_text_escaped}";
        function copyLog() {{
            navigator.clipboard.writeText(decodeURIComponent(log_text_escaped));
            const copyBtn = document.getElementById("copy-log-btn");
            copyBtn.innerHTML = "Copied!";
            setTimeout(function(){{ copyBtn.innerHTML = "Copy Conversation History"; }}, 2000);
        }}
        window.parent.document.addEventListener('keydown', (e) => {{
            if ( e.code == "Pause" ){{
                window.parent.navigator.clipboard.writeText(decodeURIComponent(log_text_escaped));
                const copyBtn = document.getElementById("copy-log-btn");
                copyBtn.innerHTML = "Copied!";
                setTimeout(function(){{ copyBtn.innerHTML = "Copy Conversation History"; }}, 2000);
            }}
        }} , false);
        </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.session_state.use_continue = st.toggle("Use Continue instead of Retry", value=st.session_state.use_continue)
        
        st.header("Advanced Settings")
        model_list = ["claude-sonnet-4-5-20250929",
                      "claude-opus-4-5-20251101",
                      "claude-haiku-4-5-20251001",
                      "claude-opus-4-1-20250805",
                      "claude-opus-4-20250514",
                      "claude-sonnet-4-20250514",
                      "claude-3-7-sonnet-20250219",
                      "claude-3-5-haiku-20241022",
                      "claude-3-haiku-20240307",
                      "claude-3-opus-20240229",
                      ]
        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,
                                                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_k = st.slider("Top-K", min_value=0, max_value=500, value=0, step=1)
        top_p = st.slider("Top-P", min_value=0.01, max_value=1.00, value=1.00, step=0.01)
        max_tokens = st.slider("Max Output Tokens", min_value=1, max_value=4096, value=4096, step=1)

        st.header("Extended Thinking")
        st.session_state.use_thinking = st.toggle(
            "Enable extended thinking", 
            value=st.session_state.use_thinking,
            help="Enable Claude's enhanced reasoning capabilities"
        )
        if st.session_state.use_thinking:
            st.session_state.thinking_budget_tokens = st.slider(
                "Thinking budget tokens", 
                min_value=1024, 
                max_value=4000, 
                value=st.session_state.thinking_budget_tokens, 
                step=100,
                help="Maximum tokens Claude can use for internal reasoning"
            )
            
            # Check relationship between max_tokens and budget_tokens
            if st.session_state.thinking_budget_tokens >= max_tokens:
                st.warning("Thinking budget tokens must be less than Max Output Tokens")
        
        st.header("Context limitation")
        col_limit_val, col_limit_unit = st.columns([1, 1])
        with col_limit_val:
            limit_val = st.number_input("Limit", min_value=0, max_value=200000, value="min", help="0 means no limit")
        with col_limit_unit:
            limit_unit = st.selectbox("Unit", options=["Turns","Tokens"], index=0, help="Tokens Unit is a rough estimate")

        st.header("Tokens")
        st.session_state.show_message_tokens = st.toggle("Show message tokens")
        system_prompt_tokens = count_tokens(system_prompt)
        prefill_tokens = count_tokens(st.session_state.prefill)
        messages_tokens = 0
        for message in st.session_state.messages:
            messages_tokens += count_tokens(message["content"])
        total_tokens = system_prompt_tokens + prefill_tokens + messages_tokens

        tokens_table_md = f"| Part | Tokens |\n"
        tokens_table_md += f"| ---- | ---- |\n"
        tokens_table_md += f"| System | {system_prompt_tokens} |\n"
        tokens_table_md += f"| Prefill | {prefill_tokens} |\n"
        tokens_table_md += f"| Message | {messages_tokens}({len(st.session_state.messages)}) |\n"
        tokens_table_md += f"| Total | {total_tokens} |\n\n"
        tokens_table_md += f"These token counts are very rough estimates.\n"
        st.markdown(tokens_table_md)

        st.header("Restore History")
        history_input = st.text_area("Paste conversation history:", height=200)
        if st.button("Restore History"):
            st.session_state.messages = []
            st.session_state.exception = None
            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:
                        if role == "assistant":
                            thinking_match = re.search(r"<THINKING>\n(.*?)</THINKING>\n", text, re.DOTALL)
                            if thinking_match:
                                thinking_content = thinking_match.group(1).strip()
                                content = re.sub(r"<THINKING>\n.*?</THINKING>\n", "", text, flags=re.DOTALL).strip()
                                st.session_state.messages.append({
                                    "role": role, 
                                    "content": content,
                                    "thinking": thinking_content
                                })
                            else:
                                st.session_state.messages.append({"role": role, "content": text.strip()})
                        else:
                            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:
                if role == "assistant":
                    thinking_match = re.search(r"<THINKING>\n(.*?)</THINKING>\n", text, re.DOTALL)
                    if thinking_match:
                        thinking_content = thinking_match.group(1).strip()
                        content = re.sub(r"<THINKING>\n.*?</THINKING>\n", "", text, flags=re.DOTALL).strip()
                        st.session_state.messages.append({
                            "role": role, 
                            "content": content,
                            "thinking": thinking_content
                        })
                    else:
                        st.session_state.messages.append({"role": role, "content": text.strip()})
                else:
                    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.session_state.exception = None
            st.rerun()
        
        st.header("Change API Key")
        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
                client = anthropic.Anthropic(api_key=new_api_key)
                st.success("API Key updated successfully!")
            else:
                st.warning("Please enter a valid API Key.")
    
    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"):
    message_content = {"role": "assistant", "content": st.session_state.response}
    if st.session_state.last_thinking:
        message_content["thinking"] = st.session_state.last_thinking
    
    st.session_state.messages.append(message_content)
    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.session_state.message_continue = ""
    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)
        
        message_content = {"role": "assistant", "content": response}
        if st.session_state.last_thinking:
            message_content["thinking"] = st.session_state.last_thinking
            
        st.session_state.messages.append(message_content)

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

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)

    message_content = {"role": "assistant", "content": response}
    if st.session_state.last_thinking:
        message_content["thinking"] = st.session_state.last_thinking
        
    st.session_state.messages.append(message_content)

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