import streamlit as st
import requests
import json

st.set_page_config(page_title="ASI1-mini Chatbot", layout="centered")
st.title("ASI1-mini LLM Chat")
st.caption("Powered by [asi1.ai](https://asi1.ai/)")

# Sidebar API Key Input
st.sidebar.subheader("🔑 API Key")
asi_api_key = st.sidebar.text_input("ASI1 API Key", type="password")
st.sidebar.caption(
    "Get your API key from [asi1.ai](https://asi1.ai/dashboard/api-keys)"
)

# API Docs Button
st.sidebar.markdown(
    """
    <a href="https://docs.asi1.ai/docs/" target="_blank">
        <button style="width: 100%; padding: 0.5em; font-size: 16px; background-color: #4CAF50; color: white; border: none; border-radius: 5px;">
            🔍 Looking for API docs?
        </button>
    </a>
    """,
    unsafe_allow_html=True,
)

# Chat history initialization
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

history_to_render = st.session_state.chat_history.copy()
if (
    history_to_render
    and history_to_render[-1]["role"] == "assistant"
    and st.session_state.get("awaiting_response", False)
):
    history_to_render = history_to_render[:-1]

# Display past chat messages
for msg in history_to_render:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# Chat input box
user_input = st.chat_input("Ask me anything...")


def stream_asi1_api(api_key, messages):
    url = "https://api.asi1.ai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": "asi1-mini",
        "messages": messages,
        "temperature": 0.7,
        "stream": True,
        "max_tokens": 1024,
    }

    with requests.post(url, headers=headers, json=payload, stream=True) as response:
        response.raise_for_status()
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode("utf-8").strip()
                if decoded_line.startswith("data: "):
                    data = decoded_line[len("data: ") :]
                    if data == "[DONE]":
                        break
                    try:
                        parsed = json.loads(data)
                        content_piece = (
                            parsed.get("choices", [{}])[0]
                            .get("delta", {})
                            .get("content", "")
                        )
                        if content_piece:
                            yield content_piece
                    except json.JSONDecodeError:
                        continue


# When user submits a message
if user_input:
    if not asi_api_key:
        st.error("Please enter your API key in the sidebar.")
    else:
        # Mark that we're about to stream a new response
        st.session_state["awaiting_response"] = True

        # Add user message to chat history
        st.session_state.chat_history.append({"role": "user", "content": user_input})
        with st.chat_message("user"):
            st.markdown(user_input)

        with st.chat_message("assistant"):
            response_stream = stream_asi1_api(
                asi_api_key, st.session_state.chat_history
            )
            placeholder = st.empty()
            partial_text = ""

            for chunk in response_stream:
                partial_text += chunk
                placeholder.markdown(partial_text + "▌")

            placeholder.markdown(partial_text)

        # Save final assistant message to history
        st.session_state.chat_history.append(
            {"role": "assistant", "content": partial_text}
        )
        st.session_state["awaiting_response"] = False