import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain.agents import create_agent
from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper

from agent_utils import run_agent

load_dotenv()


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


# --- 도구 2: 위키백과 (백과사전 지식, 키 불필요) ---
_wiki = WikipediaQueryRun(
    api_wrapper=WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=400, lang="ko")
)


@tool
def search_wikipedia(query: str) -> str:
    """위키백과에서 사전적/역사적 사실을 검색한다 (인물·장소·개념 등)."""
    try:
        return _wiki.run(query)
    except Exception as e:
        return f"위키백과 검색 오류: {e}"


# --- 도구 3: 웹 검색 (Tavily, 최신 정보) ---
@tool
def web_search(query: str) -> str:
    """최신/실시간 정보를 웹에서 검색한다 (뉴스, 요즘 이슈 등)."""
    if not os.getenv("TAVILY_API_KEY"):
        return "웹 검색 불가: 루트 .env 에 TAVILY_API_KEY 를 추가하면 동작합니다."
    try:
        from langchain_tavily import TavilySearch
        return str(TavilySearch(max_results=2).invoke(query))
    except Exception as e:
        return f"웹 검색 오류: {e}"


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

agent = create_agent(
    llm,
    [calculator, search_wikipedia, web_search],
    system_prompt=(
        "너는 도구를 쓰는 도우미다. 사전·역사 사실은 search_wikipedia, "
        "최신·실시간 정보는 web_search, 계산은 calculator 를 골라 써라. "
        "도구가 필요 없으면 그냥 답하라."
    ),
)

if __name__ == "__main__":
    run_agent(agent, "에펠탑은 누가 언제 만들었어?")        # → 위키백과
    print()
    run_agent(agent, "123 * 456 은 얼마야?")                # → 계산기
    print()
    run_agent(agent, "요즘 인공지능 업계 최신 뉴스 하나만 알려줘")  # → 웹검색(Tavily)
