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()

# --- 작은 RAG: 문서마다 '숫자'를 담아 검색+계산 조합 질문이 가능하게 한다 ---
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__)          # static_folder='static' (기본) → /static/* 자동 서빙
app.secret_key = "web_demo_secret"


def build_trace(messages):
    """에이전트 실행 메시지를 화면용 단계 리스트로 변환한다."""
    # (1) 정답
    trace = []
    for m in messages:
        t = getattr(m, "type", "")
        if t == "human":
            trace.append({"step": "질문", "text": m.content})
        elif t == "ai":
            for c in (getattr(m, "tool_calls", []) or []):
                trace.append({"step": "판단·도구 호출", "tool": c["name"], "args": c["args"]})
            if m.content:
                trace.append({"step": "최종 답변", "text": m.content})
        elif t == "tool":
            trace.append({"step": "도구 결과", "tool": getattr(m, "name", ""), "text": m.content})
    return trace


@app.route("/ask", methods=["POST"])
def ask():
    # (2) 정답
    question = (request.json or {}).get("question", "").strip()
    if not question:
        return jsonify({"error": "질문을 입력하세요"}), 400
    if "tid" not in session:
        session["tid"] = str(uuid.uuid4())          # 브라우저 세션별 대화 thread
    result = agent.invoke(
        {"messages": [("user", question)]},
        config={"configurable": {"thread_id": session["tid"]}},
    )
    return jsonify({
        "answer": result["messages"][-1].content,
        "trace": build_trace(result["messages"]),
    })


@app.route("/")
def index():
    # 템플릿 엔진 없이 static/ 의 정적 HTML 을 그대로 서빙한다.
    return send_from_directory(app.static_folder, "index.html")


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