import operator
import os
from typing import Annotated, Sequence, TypedDict

from dotenv import load_dotenv
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
from langgraph.prebuilt import ToolExecutor, ToolInvocation

# --- 1. 환경 설정 및 도구 정의 ---
load_dotenv()

# .env 파일에 TAVILY_API_KEY 설정 필요
tavily_tool = TavilySearchResults(max_results=3)
tools = [tavily_tool]
tool_executor = ToolExecutor(tools)


# --- 2. 에이전트 상태(State) 정의 ---
# 그래프의 모든 노드가 공유하는 상태 객체
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    # 각 노드에서 다음 노드로 전달될 메시지 목록
    # `operator.add`는 새 메시지가 기존 목록에 추가되도록 함


# --- 3. 에이전트 노드(Node) 생성 ---
# LLM 모델 정의
model = ChatOpenAI(temperature=0, streaming=True)


def create_agent_node(llm, system_message: str):
    """
    주어진 시스템 메시지를 따르는 에이전트 노드를 생성하는 팩토리 함수.
    이 노드는 LLM을 호출하여 응답(액션 또는 최종 답변)을 결정합니다.
    """
    prompt = [
        ("system", system_message),
        ("placeholder", "{messages}"),
    ]
    agent_runnable = ChatOpenAI(model="gpt-4-turbo", temperature=0).bind_tools(tools)

    def agent_node(state: AgentState):
        result = agent_runnable.invoke(state["messages"])
        return {"messages": [result]}

    return agent_node


def create_tool_node():
    """도구를 실행하는 노드를 생성합니다."""

    def tool_node(state: AgentState):
        tool_invocations = []
        last_message = state["messages"][-1]
        for tool_call in last_message.tool_calls:
            tool_invocations.append(
                ToolInvocation(tool=tool_call["name"], tool_input=tool_call["args"])
            )

        responses = tool_executor.batch(tool_invocations, return_exceptions=True)
        # 응답을 BaseMessage 형태로 변환
        from langchain_core.messages import ToolMessage

        tool_messages = [
            ToolMessage(content=str(response), tool_call_id=inv.tool_call_id)
            for inv, response in zip(last_message.tool_calls, responses)
        ]
        return {"messages": tool_messages}

    return tool_node


# --- 4. 조건부 엣지(Conditional Edge)를 위한 함수 ---
def should_continue(state: AgentState):
    """
    마지막 메시지에 tool_calls가 있으면 'continue' (도구 실행),
    없으면 'end' (종료)를 반환하여 워크플로우를 분기합니다.
    """
    if "tool_calls" in state["messages"][-1].additional_kwargs:
        return "continue"
    return "end"


# --- 5. 그래프 구축 ---
# 각 역할을 수행할 에이전트 정의
researcher_node = create_agent_node(
    model,
    "당신은 전문 리서처입니다. 주어진 주제에 대해 신뢰할 수 있는 출처를 바탕으로 상세하고 깊이 있는 정보를 찾아 정리해주세요.",
)
blogger_node = create_agent_node(
    model,
    "당신은 IT 전문 블로거입니다. 주어진 리서치 결과를 바탕으로 독자들이 이해하기 쉽고 흥미를 느낄만한 블로그 포스트 초안을 작성해주세요.",
)
critic_node = create_agent_node(
    model,
    "당신은 날카로운 비평가입니다. 작성된 블로그 초안을 검토하고, 내용의 정확성, 논리적 흐름, 표현의 명확성 등을 평가하여 구체적인 수정 제안을 해주세요.",
)
tool_node = create_tool_node()

# StateGraph를 사용하여 워크플로우 정의
workflow = StateGraph(AgentState)

workflow.add_node("Researcher", researcher_node)
workflow.add_node("Blogger", blogger_node)
workflow.add_node("Critic", critic_node)
workflow.add_node("call_tool", tool_node)

# 각 노드 간의 흐름(엣지) 정의
workflow.set_entry_point("Researcher")  # 시작점은 Researcher
workflow.add_edge("Researcher", "Blogger")
workflow.add_edge("Blogger", "Critic")


# 조건부 엣지: 비평 결과에 따라 흐름 분기
def after_critic(state: AgentState):
    last_message = state["messages"][-1].content
    if "수정할 필요 없이 훌륭합니다" in last_message:
        return "end"
    else:
        # 피드백을 다시 Researcher에게 전달 (실제 구현에서는 더 정교한 로직 필요)
        return "continue_research"


workflow.add_conditional_edges(
    "Critic",
    after_critic,
    {"end": END, "continue_research": "Researcher"},  # 피드백을 바탕으로 리서치 재시작
)
# 도구 사용 후에는 다시 원래 노드로 돌아가야 함 (이 예제에서는 Researcher만 도구 사용 가능)
workflow.add_conditional_edges(
    "Researcher",
    should_continue,
    {
        "continue": "call_tool",
        "end": "Blogger",  # Researcher의 최종 결과물을 Blogger에게 전달
    },
)
workflow.add_edge("call_tool", "Researcher")


# --- 6. 그래프 컴파일 및 실행 ---
# 그래프 컴파일
app = workflow.compile()

# 그래프 시각화 (pygraphviz 또는 pydot 필요)
try:
    image_data = app.get_graph().draw_png()
    with open("autonomous_agent_workflow.png", "wb") as f:
        f.write(image_data)
    print("워크플로우 그래프를 'autonomous_agent_workflow.png' 파일로 저장했습니다.")
except ImportError:
    print("Graphviz가 설치되지 않아 그래프를 시각화할 수 없습니다.")
    print(
        "시각화를 원하시면 'pip install pygraphviz' 또는 'pip install pydot'을 실행하세요."
    )


# 실행
topic = "최신 AI 에이전트 기술 동향과 LangGraph의 역할"
initial_messages = [
    HumanMessage(content=f"주제: '{topic}'에 대한 리서치를 시작해주세요.")
]

# 스트리밍 출력을 위해 invoke 대신 stream 사용
for s in app.stream({"messages": initial_messages}):
    print(s)
    print("---")
