# ============================================================
# [실습] 통합 웹 데모로 에이전트 들여다보기 (Flask)
#
# 자기완결 RAG 에이전트를 Flask 웹으로 감싸, 질문에 대해 에이전트가 어떤 도구를 어떤
# 판단으로 호출했는지(trace)를 오른쪽 패널로 보여 준다. (화면은 static/ 으로 분리, 제공됨)
#
# 구현 항목:
#  1) build_trace(messages): result["messages"] 를 [질문 / 판단·도구 호출 / 도구 결과 / 최종 답변]
#     단계 리스트로 변환한다(우측 트레이스 패널용).
#  2) POST /ask: 질문을 받아 thread_id(session)로 agent.invoke 하고 {answer, trace} 를 JSON 으로 반환.
# ============================================================

import os
import uuid
from dotenv import load_dotenv
from flask import Flask, request, jsonify, session, send_from_directory
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.tools import tool
from langchain.agents import create_agent
from langgraph.checkpoint.memory import MemorySaver

load_dotenv()

DOCS = [
    "NVMe SSD의 순차 읽기 속도는 약 3500MB/s이다.",
    "SATA SSD의 순차 읽기 속도는 약 550MB/s이다.",
    "이 서버에는 NVMe SSD가 4개 장착되어 있다.",
    "RAID 0은 속도를 위한 스트라이핑, RAID 1은 안정성을 위한 미러링 구성이다.",
]
splits = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=30).split_documents(
    [Document(page_content=d) for d in DOCS])
vectorstore = Chroma.from_documents(
    splits, OpenAIEmbeddings(model="text-embedding-3-small"), collection_name="web_demo")


@tool
def search_docs(query: str) -> str:
    """기술 문서에서 질문과 관련된 내용을 검색한다 (SSD 속도, 장착 개수, RAID 등)."""
    hits = vectorstore.similarity_search(query, k=3)
    return "\n".join(f"- {d.page_content}" for d in hits) or "관련 문서를 찾을 수 없습니다."


@tool
def calculator(expression: str) -> str:
    """수학 식을 계산한다. 예: '3500 + 550', '3500 * 4'."""
    try:
        return str(eval(expression, {"__builtins__": {}}, {}))
    except Exception as e:
        return f"계산 오류: {e}"


agent = create_agent(
    ChatOpenAI(model="gpt-4o-mini", temperature=0),
    [search_docs, calculator],
    system_prompt=(
        "문서 속 수치 질문은 먼저 search_docs 로 값을 찾고, 그 값으로 계산이 필요하면 "
        "calculator 를 써라. 모르면 모른다고 답하라."
    ),
    checkpointer=MemorySaver(),
)

app = Flask(__name__)
app.secret_key = "web_demo_secret"


def build_trace(messages):
    """에이전트 실행 메시지를 화면용 단계 리스트로 변환한다."""
    # 1) 여기에 코드를 작성하세요
    #    (human→질문 / ai.tool_calls→판단·도구 호출 / tool→도구 결과 / ai.content→최종 답변)
    pass


@app.route("/ask", methods=["POST"])
def ask():
    # 2) 여기에 코드를 작성하세요
    #    (question 받기 → session["tid"] 로 thread 구분 → agent.invoke → {answer, trace} 반환)
    pass


@app.route("/")
def index():
    return send_from_directory(app.static_folder, "index.html")


if __name__ == "__main__":
    app.run(debug=True, port=5001)
