app.py

"""
소방 복무관리 RAG 챗봇
깔끔하고 심플한 Gradio UI + 관리자 탭
"""

import gradio as gr
from rag import chat
from pdf_processor import process_pdf
from database import get_all_documents
import uuid
import os

# 세션별 대화 기록 저장
conversation_history = {}

# 관리자 비밀번호 (환경변수로 설정 권장)
ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "admin1234")


def respond(message: str, history: list, session_id: str) -> tuple:
    """채팅 응답 처리"""
    if not message.strip():
        return "", history

    if session_id not in conversation_history:
        conversation_history[session_id] = []

    response = chat(message, conversation_history[session_id])

    conversation_history[session_id].append({
        "user": message,
        "assistant": response
    })

    history.append((message, response))
    return "", history


def clear_chat(session_id: str) -> tuple:
    """대화 초기화"""
    if session_id in conversation_history:
        conversation_history[session_id] = []
    return [], ""


def verify_admin(password: str) -> tuple:
    """관리자 인증"""
    if password == ADMIN_PASSWORD:
        return gr.update(visible=False), gr.update(visible=True), "✅ 인증 성공"
    return gr.update(visible=True), gr.update(visible=False), "❌ 비밀번호가 틀렸습니다"


def upload_pdf(file, title: str, category: str) -> str:
    """PDF 업로드 및 처리"""
    if file is None:
        return "❌ PDF 파일을 선택해주세요."

    if not title.strip():
        return "❌ 문서 제목을 입력해주세요."

    result = process_pdf(
        pdf_path=file.name,
        title=title.strip(),
        category=category
    )

    return result["message"]


def get_documents_list() -> list:
    """저장된 문서 목록 조회"""
    docs = get_all_documents()
    if not docs:
        return [["데이터 없음", "-", "-"]]

    return [[d.get("title", ""), d.get("category", ""), str(d.get("id", ""))] for d in docs]


# 커스텀 CSS
custom_css = """
.gradio-container {
    max-width: 900px !important;
    margin: auto !important;
    font-family: 'Pretendard', -apple-system, BlinkMacSystemFont, sans-serif !important;
}

.header {
    text-align: center;
    padding: 20px 0;
    border-bottom: 1px solid #e5e7eb;
    margin-bottom: 20px;
}

.header h1 {
    color: #dc2626;
    font-size: 1.8rem;
    font-weight: 700;
    margin: 0;
}

.header p {
    color: #6b7280;
    font-size: 0.9rem;
    margin-top: 8px;
}

.chatbot {
    border-radius: 12px !important;
    border: 1px solid #e5e7eb !important;
    box-shadow: 0 1px 3px rgba(0,0,0,0.1) !important;
}

.input-area {
    border-radius: 12px !important;
    border: 1px solid #e5e7eb !important;
}

.primary-btn {
    background-color: #dc2626 !important;
    border: none !important;
    border-radius: 8px !important;
    color: white !important;
    font-weight: 600 !important;
}

.primary-btn:hover {
    background-color: #b91c1c !important;
}

.secondary-btn {
    background-color: #f3f4f6 !important;
    border: 1px solid #e5e7eb !important;
    border-radius: 8px !important;
    color: #374151 !important;
}

.admin-section {
    background: #fef2f2;
    border: 1px solid #fecaca;
    border-radius: 12px;
    padding: 20px;
    margin: 10px 0;
}

.upload-section {
    background: #f9fafb;
    border: 1px solid #e5e7eb;
    border-radius: 12px;
    padding: 20px;
    margin: 10px 0;
}

.footer {
    text-align: center;
    padding: 16px 0;
    color: #9ca3af;
    font-size: 0.8rem;
}
"""

# Gradio 앱 구성
with gr.Blocks(css=custom_css, title="소방 복무관리 챗봇", theme=gr.themes.Soft()) as app:
    session_id = gr.State(value=str(uuid.uuid4()))

    # 헤더
    gr.HTML("""
        <div class="header">
            <h1>🚒 소방 복무관리 챗봇</h1>
            <p>복무, 근무, 휴가 등 소방 업무에 관해 질문해주세요</p>
        </div>
    """)

    # 탭 구성
    with gr.Tabs():
        # ========== 채팅 탭 ==========
        with gr.TabItem("💬 채팅"):
            chatbot = gr.Chatbot(
                label="",
                height=450,
                show_label=False,
                container=True,
                elem_classes=["chatbot"]
            )

            with gr.Row():
                msg = gr.Textbox(
                    placeholder="메시지를 입력하세요...",
                    show_label=False,
                    container=False,
                    scale=9,
                    elem_classes=["input-area"]
                )
                submit_btn = gr.Button("전송", scale=1, elem_classes=["primary-btn"])

            with gr.Row():
                clear_btn = gr.Button("대화 초기화", elem_classes=["secondary-btn"])

            gr.Examples(
                examples=[
                    "연차 휴가 일수는 어떻게 계산하나요?",
                    "당직 근무 시간은 어떻게 되나요?",
                    "병가 신청 절차가 어떻게 되나요?",
                    "초과 근무 수당 기준이 어떻게 되나요?"
                ],
                inputs=msg,
                label="예시 질문"
            )

        # ========== 관리자 탭 ==========
        with gr.TabItem("⚙️ 관리자"):
            # 로그인 섹션
            with gr.Column(visible=True) as login_section:
                gr.Markdown("### 🔐 관리자 인증")
                gr.Markdown("문서를 관리하려면 관리자 비밀번호를 입력하세요.")

                with gr.Row():
                    password_input = gr.Textbox(
                        label="비밀번호",
                        type="password",
                        placeholder="비밀번호 입력...",
                        scale=3
                    )
                    login_btn = gr.Button("로그인", elem_classes=["primary-btn"], scale=1)

                login_status = gr.Markdown("")

            # 관리자 패널 (인증 후 표시)
            with gr.Column(visible=False) as admin_panel:
                gr.Markdown("### 📄 PDF 문서 업로드")
                gr.Markdown("새로운 복무 관련 문서를 업로드하여 챗봇 지식을 확장하세요.")

                with gr.Group(elem_classes=["upload-section"]):
                    pdf_file = gr.File(
                        label="PDF 파일 선택",
                        file_types=[".pdf"],
                        file_count="single"
                    )

                    with gr.Row():
                        doc_title = gr.Textbox(
                            label="문서 제목",
                            placeholder="예: 소방공무원 복무규정 2024",
                            scale=2
                        )
                        doc_category = gr.Dropdown(
                            label="카테고리",
                            choices=["일반", "휴가", "근무", "수당", "규정", "교육"],
                            value="일반",
                            scale=1
                        )

                    upload_btn = gr.Button("📤 업로드 및 처리", elem_classes=["primary-btn"])
                    upload_result = gr.Markdown("")

                gr.Markdown("---")
                gr.Markdown("### 📚 저장된 문서 목록")

                refresh_btn = gr.Button("🔄 새로고침", elem_classes=["secondary-btn"])
                documents_table = gr.Dataframe(
                    headers=["제목", "카테고리", "ID"],
                    label="",
                    interactive=False
                )

    # 푸터
    gr.HTML("""
        <div class="footer">
            소방 복무관리 AI 어시스턴트 | Powered by HuggingFace & Supabase
        </div>
    """)

    # ========== 이벤트 핸들러 ==========
    # 채팅
    msg.submit(respond, [msg, chatbot, session_id], [msg, chatbot])
    submit_btn.click(respond, [msg, chatbot, session_id], [msg, chatbot])
    clear_btn.click(clear_chat, [session_id], [chatbot, msg])

    # 관리자
    login_btn.click(
        verify_admin,
        [password_input],
        [login_section, admin_panel, login_status]
    )
    upload_btn.click(
        upload_pdf,
        [pdf_file, doc_title, doc_category],
        [upload_result]
    )
    refresh_btn.click(
        get_documents_list,
        [],
        [documents_table]
    )


if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)