import streamlit as st
import requests
import json
import os

# --- 설정 및 보안 ---
# 환경변수 미설정 시 비밀번호 0000
st.set_page_config(page_title="한양대 학칙 챗봇", page_icon="🦁")
ACCESS_PASSWORD = os.getenv("APP_PASSWORD", "0000")

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

def check_password():
    if st.session_state.password_input == ACCESS_PASSWORD:
        st.session_state.authenticated = True
        del st.session_state.password_input
    else:
        st.error("비밀번호가 틀렸습니다.")

if not st.session_state.authenticated:
    st.title("🔒 접근 제한")
    st.text_input("비밀번호를 입력하세요", type="password", key="password_input", on_change=check_password)
    st.stop()

# --- UI 구현 ---
st.title("한양대학교 학칙 챗봇 (RAG)")
st.caption("GPT-5-mini & Real-time Hybrid Search & Reranking")

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

# 히스토리 표시
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        if message.get("docs"):
            with st.expander("참고 문서 (Evidence)"):
                for idx, doc in enumerate(message["docs"]):
                    st.text(f"[{idx+1}] {doc['source']}: {doc['content']}")

# 입력 처리
if prompt := st.chat_input("질문을 입력하세요..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        status_container = st.status("🚀 처리 중...", expanded=True)
        answer_placeholder = st.empty()
        full_answer = ""
        received_docs = []
        
        try:
            # stream=True로 요청하여 한 줄씩 받음
            response = requests.post(
                "http://localhost:8000/chat",
                json={"query": prompt, "history": []},
                stream=True
            )
            
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line.decode('utf-8'))
                        msg_type = data["type"]
                        content = data["content"]
                        
                        if msg_type == "log":
                            status_container.write(content)
                            status_container.update(label=content)
                            
                        elif msg_type == "answer":
                            # 답변은 실시간 타이핑 효과
                            full_answer += content
                            answer_placeholder.markdown(full_answer + "▌")
                            
                        elif msg_type == "docs":
                            received_docs = content
                            
                    except json.JSONDecodeError:
                        continue
            
            # 완료 후 처리
            answer_placeholder.markdown(full_answer)
            status_container.update(label="✅ 처리 완료", state="complete", expanded=False)
            
            if received_docs:
                with st.expander("참고 문서"):
                    for idx, doc in enumerate(received_docs):
                        st.info(f"**{doc['source']}**: {doc['content']}")

            st.session_state.messages.append({
                "role": "assistant", 
                "content": full_answer,
                "docs": received_docs
            })
            
        except Exception as e:
            status_container.update(label="❌ 오류 발생", state="error")
            st.error(f"Connection Error: {e}")