from dotenv import load_dotenv
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

from agent_utils import run_agent

load_dotenv()

# --- 작은 RAG (검색 대상 지식베이스) ---
DOCS = [
    "NVMe는 PCIe 버스를 사용하는 SSD 인터페이스로, SATA보다 훨씬 빠르다.",
    "SATA SSD의 최대 대역폭은 약 550MB/s로, NVMe(수 GB/s)보다 낮다.",
    "RAID 0은 속도를 위한 스트라이핑, RAID 1은 안정성을 위한 미러링 구성이다.",
]
splits = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50).split_documents(
    [Document(page_content=d) for d in DOCS])
vectorstore = Chroma.from_documents(
    splits, OpenAIEmbeddings(model="text-embedding-3-small"), collection_name="agent_finale")


# (1) 정답 — 검색 도구
@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 "관련 문서를 찾을 수 없습니다."


# (2) 정답 — 계산 도구
@tool
def calculator(expression: str) -> str:
    """수학 식을 계산한다. 예: '550 * 6'."""
    try:
        return str(eval(expression, {"__builtins__": {}}, {}))
    except Exception as e:
        return f"계산 오류: {e}"


TOOLS = [search_docs, calculator]
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# (3) 정답
agent = create_agent(
    llm,
    TOOLS,
    system_prompt=(
        "너는 문서 기반 도우미다. 저장장치 관련 질문은 search_docs 로 근거를 찾아 답하고, "
        "계산이 필요하면 calculator 를 써라. 문서에 없는 내용은 추측하지 말고 모른다고 답하라."
    ),
    checkpointer=MemorySaver(),
)


# (4) 정답 — 콘솔 멀티턴 루프 (각 턴마다 thinking 을 단계별로 출력)
def main():
    config = {"configurable": {"thread_id": "console"}}
    print("질문을 입력하세요 (종료: quit)")
    while True:
        q = input("\nQ> ").strip()
        if q in {"quit", "exit"}:
            break
        run_agent(agent, q, config=config)


if __name__ == "__main__":
    main()
