from langchain.chat_models import init_chat_model
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_community.tools import TavilySearchResults
from langchain_core.tools import tool

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

embeddings = OpenAIEmbeddings()
vectorstore = Chroma(collection_name="documents", embedding_function=embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 4})

@tool
def vector_search(query: str) -> str:
    """벡터 데이터베이스에서 관련 문서를 검색합니다."""
    docs = retriever.invoke(query)
    return "\n\n".join(doc.page_content for doc in docs)

@tool
def web_search(query: str) -> str:
    """웹에서 최신 정보를 검색합니다."""
    search = TavilySearchResults(max_results=3)
    results = search.invoke(query)
    return "\n\n".join(r["content"] for r in results)

@tool
def calculator(expression: str) -> str:
    """수학 계산을 수행합니다. 예: '2 + 3 * 4'"""
    import ast
    import operator

    # eval() 대신 안전한 AST 기반 파서를 사용합니다.
    # eval()은 임의 코드 실행이 가능하므로 프로덕션에서 절대 사용하지 마세요.
    ops = {
        ast.Add: operator.add, ast.Sub: operator.sub,
        ast.Mult: operator.mul, ast.Div: operator.truediv,
        ast.Pow: operator.pow, ast.USub: operator.neg,
    }

    def _eval(node):
        if isinstance(node, ast.Expression):
            return _eval(node.body)
        elif isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
            return node.value
        elif isinstance(node, ast.BinOp) and type(node.op) in ops:
            return ops[type(node.op)](_eval(node.left), _eval(node.right))
        elif isinstance(node, ast.UnaryOp) and type(node.op) in ops:
            return ops[type(node.op)](_eval(node.operand))
        raise ValueError(f"지원하지 않는 연산: {ast.dump(node)}")

    try:
        tree = ast.parse(expression, mode="eval")
        result = _eval(tree)
        return str(result)
    except Exception as e:
        return f"계산 오류: {e}"

tools = [vector_search, web_search, calculator]

from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]

# 도구가 바인딩된 LLM
llm_with_tools = llm.bind_tools(tools)

SYSTEM_PROMPT = (
    "당신은 질문에 답변하는 AI 어시스턴트입니다.\n"
    "주어진 도구를 활용하여 정확한 정보를 수집한 뒤 답변하세요.\n"
    "- vector_search: 내부 문서에서 정보를 찾을 때 사용\n"
    "- web_search: 최신 정보나 내부 문서에 없는 정보가 필요할 때 사용\n"
    "- calculator: 수학 계산이 필요할 때 사용\n"
    "충분한 정보를 수집한 후에 답변하세요."
)

def agent(state: AgentState) -> AgentState:
    """에이전트가 도구 사용 여부를 판단하고 응답합니다."""
    messages = [SystemMessage(content=SYSTEM_PROMPT)] + state["messages"]
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}

# 그래프 구성
workflow = StateGraph(AgentState)

workflow.add_node("agent", agent)
workflow.add_node("tools", ToolNode(tools))

workflow.add_edge(START, "agent")
workflow.add_conditional_edges("agent", tools_condition)
workflow.add_edge("tools", "agent")

app = workflow.compile()

# 벡터 검색 활용
result = app.invoke({
    "messages": [HumanMessage(content="Self-RAG의 Reflection Token에 대해 설명해줘")]
})

# 웹 검색 활용
result = app.invoke({
    "messages": [HumanMessage(content="2024년 가장 주목받은 RAG 논문은?")]
})

# 복합 도구 활용
result = app.invoke({
    "messages": [HumanMessage(content="RAG 논문 수가 2023년 대비 2024년에 몇 배 증가했나?")]
})

# 마지막 메시지 출력
print(result["messages"][-1].content)

# 최대 반복 횟수 제한
app = workflow.compile()
result = app.invoke(
    {"messages": [HumanMessage(content="질문")]},
    config={"recursion_limit": 10},
)