import streamlit as st
import pandas as pd
import json
from io import StringIO
from datetime import datetime
import requests
import sseclient

st.set_page_config(page_title="Solar Prompt Comparator", page_icon="🌞")
st.title("🌞 System Prompt Comparison Chat")

# --- Input UI ---
api_key = st.text_input("🔑 Upstage API Key (starts with 'up_')", type="password")
user_input = st.text_input("💬 Your Message", "")

st.markdown("### 🧠 Custom System Prompt")
st.markdown("""
**What is a System Prompt?**  
A system prompt is a special instruction given to the AI before the conversation begins.  
It tells the AI how it should behave — like setting the tone, role, or personality.  
For example, you can ask it to be a polite assistant, a creative writer, or a strict grammar checker.  
Changing the system prompt can lead to very different responses, even if the user's question stays the same.
""")

custom_prompt = st.text_area("✏️ Write your own system prompt below:", "You are a helpful assistant.", height=100)

# --- Default prompt setting ---
default_prompt = "You are a helpful assistant."

# --- Session Reset ---
if "default_messages" not in st.session_state:
    st.session_state.default_messages = [{"role": "system", "content": default_prompt}]
if "custom_messages" not in st.session_state or st.session_state.custom_prompt != custom_prompt:
    st.session_state.custom_messages = [{"role": "system", "content": custom_prompt}]
    st.session_state.custom_prompt = custom_prompt

# --- Solar Pro API call ---
def solar_pro_chat(messages, api_key):
    url = "https://api.upstage.ai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Accept": "text/event-stream",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "solar-pro", 
        "messages": messages,
        "stream": True
    }

    try:
        response = requests.post(url, headers=headers, json=payload, stream=True)
        response.raise_for_status()
        client = sseclient.SSEClient(response)
        for event in client.events():
            if event.data == "[DONE]":
                break
            try:
                content = json.loads(event.data)["choices"][0]["delta"].get("content", "")
                yield content
            except Exception as e:
                st.error("⚠️ An error occurred while processing the response.")
                print(f"[SSE parsing error] {e}")
                continue
    except requests.exceptions.RequestException as e:
        st.error("❌ API call failed: Check your API key or network status")
        print(f"[API request error] {e}")
        return

if not api_key or not user_input:
    st.info("👆 Please enter your API Key and a message to start comparing prompts!")
    
# --- Response processing for comparison ---
if st.button("🚀 Compare Responses") and api_key and user_input:
    st.session_state.default_messages.append({"role": "user", "content": user_input})
    st.session_state.custom_messages.append({"role": "user", "content": user_input})

    col1, col2 = st.columns(2)

    with col1:
        st.subheader("🔹 Default Prompt")
        default_response = ""
        default_area = st.empty()
        with st.spinner("Generating default response..."):
            for chunk in solar_pro_chat(st.session_state.default_messages, api_key):
                default_response += chunk
                default_area.markdown(f"**🤖 Bot:** {default_response}")
        st.session_state.default_messages.append({"role": "assistant", "content": default_response})

    with col2:
        st.subheader("🔸 Custom Prompt")
        custom_response = ""
        custom_area = st.empty()
        with st.spinner("Generating custom response..."):
            for chunk in solar_pro_chat(st.session_state.custom_messages, api_key):
                custom_response += chunk
                custom_area.markdown(f"**🤖 Bot:** {custom_response}")
        st.session_state.custom_messages.append({"role": "assistant", "content": custom_response})

# --- Download chat history ---
def generate_csv(messages, prompt_label):
    rows = [{"role": "system", "content": prompt_label}]
    for msg in messages:
        if msg["role"] != "system":
            rows.append(msg)
    df = pd.DataFrame(rows)
    output = StringIO()
    df.to_csv(output, index=False)
    return output.getvalue()

now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
st.markdown("### ⬇️ Download Chat Histories")
col1, col2 = st.columns(2)
with col1:
    st.download_button(
        label="Download Default Chat",
        data=generate_csv(st.session_state.default_messages, default_prompt),
        file_name=f"default_chat_{now}.csv",
        mime="text/csv",
    )
with col2:
    st.download_button(
        label="Download Custom Chat",
        data=generate_csv(st.session_state.custom_messages, custom_prompt),
        file_name=f"custom_chat_{now}.csv",
        mime="text/csv",
    )