from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.tools import tool
from langchain_core.vectorstores import InMemoryVectorStore
from langchain.agents import create_agent

from agent_utils import run_agent

load_dotenv()

DOCS = [
    "회사 연차는 입사 1년 후 15일 부여되며, 이후 매년 1일씩 증가한다.",
    "재택근무는 주 3일까지 가능하고, 사전에 팀장 승인이 필요하다.",
    "경조사 휴가는 본인 결혼 5일, 직계가족 사망 5일이다.",
    "사내 카페는 평일 오전 8시부터 오후 6시까지 운영한다.",
]
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# (1) 정답
vectorstore = InMemoryVectorStore.from_texts(DOCS, embedding=embeddings)


# (2) 정답
@tool
def search_handbook(query: str) -> str:
    """사내 규정 핸드북에서 관련 내용을 검색한다 (연차/재택/휴가/카페 등)."""
    hits = vectorstore.similarity_search(query, k=2)
    return "\n".join(f"- {d.page_content}" for d in hits) or "관련 규정 없음"


llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# (3) 정답
agent = create_agent(
    llm,
    [search_handbook],
    system_prompt="사내 규정 질문은 search_handbook 로 근거를 찾아 답하라. "
                  "규정과 무관한 질문은 검색 없이 그냥 답하라.",
)

if __name__ == "__main__":
    # (4) 정답 — 규정 질문(검색 발생) vs 무관한 질문(검색 없음)
    run_agent(agent, "연차는 며칠 받을 수 있어?")
    print()
    run_agent(agent, "2 더하기 3은?")
