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 프롬프트 비교기", page_icon="🌞")
st.title("🌞 시스템 프롬프트 비교 챗봇")

# --- 입력 UI ---
api_key = st.text_input("🔑 Upstage API Key를 입력하세요 ('up_'로 시작)", type="password")
user_input = st.text_input("💬 사용자 메시지를 입력하세요", "")

# --- 시스템 프롬프트 안내 및 입력 ---
st.markdown("### 🧠 커스텀 시스템 프롬프트")
st.markdown("""
**시스템 프롬프트란 무엇인가요?**  
시스템 프롬프트는 AI에게 대화를 시작하기 전에 주어지는 특별한 지시문입니다.  
AI가 어떤 역할이나 말투, 성격으로 대화에 임해야 하는지를 알려주는 역할을 합니다.  
예를 들어, AI에게 예의 바른 비서처럼 행동하게 하거나, 창의적인 작가, 엄격한 문법 검사자처럼 행동하게 만들 수 있습니다.  
같은 질문이라도 시스템 프롬프트에 따라 완전히 다른 응답이 나올 수 있습니다.
""")

# 기본 시스템 프롬프트와 사용자 커스텀 프롬프트 입력
custom_prompt = st.text_area("✏️ 아래에 원하는 시스템 프롬프트를 입력하세요:", "You are a helpful assistant.", height=100)
default_prompt = "You are a helpful assistant."

# --- 세션 상태 초기화 ---
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 호출 함수 ---
def solar_pro_chat(messages, api_key):
    """
    Solar Pro API에 메시지를 보내고 스트리밍 형태로 응답을 받아오는 함수입니다.

    Parameters:
        messages (list): 시스템/사용자 메시지 목록
        api_key (str): Upstage API 키

    Returns:
        generator: AI의 응답을 스트리밍 형태로 한 줄씩 반환
    """
    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("⚠️ 응답 처리 중 오류가 발생했습니다.")
                print(f"[SSE 파싱 오류] {e}")
                continue
    except requests.exceptions.RequestException as e:
        st.error("❌ API 호출 실패: API 키나 네트워크 상태를 확인해주세요.")
        print(f"[API 요청 오류] {e}")
        return

# --- 비교 실행 ---
if st.button("🚀 응답 비교하기") 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_response = ""
        default_area = st.empty()
        with st.spinner("기본 응답 생성 중..."):
            for chunk in solar_pro_chat(st.session_state.default_messages, api_key):
                default_response += chunk
                default_area.markdown(f"**🤖 봇:** {default_response}")
        st.session_state.default_messages.append({"role": "assistant", "content": default_response})

    # 커스텀 프롬프트 응답
    with col2:
        st.subheader("🔸 커스텀 프롬프트")
        custom_response = ""
        custom_area = st.empty()
        with st.spinner("커스텀 응답 생성 중..."):
            for chunk in solar_pro_chat(st.session_state.custom_messages, api_key):
                custom_response += chunk
                custom_area.markdown(f"**🤖 봇:** {custom_response}")
        st.session_state.custom_messages.append({"role": "assistant", "content": custom_response})

# --- 채팅 내역 CSV로 저장 ---
def generate_csv(messages, prompt_label):
    """
    주어진 메시지를 CSV로 변환하는 함수입니다.

    Parameters:
        messages (list): 시스템/사용자/AI 메시지 리스트
        prompt_label (str): 시스템 프롬프트 텍스트

    Returns:
        str: CSV 형식의 문자열
    """
    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")

# --- 다운로드 UI ---
st.markdown("### ⬇️ 채팅 내역 다운로드")
col1, col2 = st.columns(2)
with col1:
    st.download_button(
        label="기본 응답 다운로드",
        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="커스텀 응답 다운로드",
        data=generate_csv(st.session_state.custom_messages, custom_prompt),
        file_name=f"custom_chat_{now}.csv",
        mime="text/csv",
    )