import os
import streamlit as st
import json
import anthropic

# API 설정
api_key = os.environ.get("API_KEY")
client = anthropic.Anthropic(api_key=api_key)

def get_system_prompt():
    return """당신은 친절하고 전문적인 AI 어시스턴트입니다. 
    다음 가이드라인을 따라주세요:
    - 정확하고 유용한 정보를 제공합니다
    - 불확실한 경우 그 사실을 인정합니다
    - 위험하거나 해로운 조언은 하지 않습니다
    - 개인정보는 보호합니다
    - 예의 바르고 공손하게 응답합니다
    - 한국어로 자연스럽게 대화합니다"""

def chatbot_interface():
    st.title("AI 어시스턴트와 대화하기")
    
    # 모델 고정 설정 (선택 박스 제거)
    if "ai_model" not in st.session_state:
        st.session_state["ai_model"] = "claude-3-5-sonnet-20241022"

    # 세션 상태 초기화
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # 대화 기록 불러오기
    st.sidebar.title("대화 기록")
    uploaded_file = st.sidebar.file_uploader("대화 기록 JSON 파일 불러오기")
    if uploaded_file is not None:
        st.session_state.messages = json.load(uploaded_file)

    # 메시지 표시
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # 사용자 입력
    if prompt := st.chat_input("무엇을 도와드릴까요?"):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        # AI 응답 생성
        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            full_response = ""
            
            # Claude API 호출
            with client.messages.stream(
                max_tokens=1024,
                system=get_system_prompt(),
                messages=[{"role": m["role"], "content": m["content"]} for m in st.session_state.messages],
                model=st.session_state["ai_model"]
            ) as stream:
                for text in stream.text_stream:
                    full_response += str(text) if text is not None else ""
                    message_placeholder.markdown(full_response + "▌")
                    
            message_placeholder.markdown(full_response)
            
        st.session_state.messages.append({"role": "assistant", "content": full_response})

    # 대화 기록 다운로드
    if st.button("대화 기록 다운로드"):
        json_history = json.dumps(st.session_state.messages, indent=4)
        st.download_button(
            label="대화 기록 저장하기",
            data=json_history,
            file_name="chat_history.json",
            mime="application/json"
        )

def main():
    chatbot_interface()

if __name__ == "__main__":
    main()