from typing import TypedDict, Annotated, Sequence
import operator
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langgraph.graph import StateGraph, END
from api.agent.super_agent import get_llm
from api.agent.tools import AVAILABLE_TOOLS
from langgraph.prebuilt import ToolExecutor, ToolInvocation
import json

# Definisikan state untuk Graph
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    next_node: str

tool_executor = ToolExecutor(AVAILABLE_TOOLS)

def researcher_node(state):
    """Agen Peneliti: Bertugas mencari data dan memanggil tools (seperti web_search atau crossref)."""
    messages = state["messages"]
    llm = get_llm()
    # Beri tahu LLM bahwa dia adalah Peneliti
    prompt = "Anda adalah Peneliti Akademik. Cari informasi yang relevan menggunakan tools yang ada. Jawab langsung jika Anda tahu."
    response = llm.invoke([SystemMessage(content=prompt)] + messages)
    
    # Cek apakah dia mau memanggil tool
    if getattr(response, "tool_calls", None):
        return {"messages": [response], "next_node": "tools"}
    return {"messages": [response], "next_node": "reviewer"}

def tools_node(state):
    """Menjalankan tools (seperti pencarian internet atau konversi Word)."""
    messages = state["messages"]
    last_message = messages[-1]
    
    tool_invocations = []
    for tool_call in getattr(last_message, "tool_calls", []):
        action = ToolInvocation(
            tool=tool_call["name"],
            tool_input=tool_call["args"],
        )
        tool_invocations.append(action)
        
    responses = tool_executor.batch(tool_invocations, return_exceptions=True)
    
    tool_messages = []
    for tc, response in zip(getattr(last_message, "tool_calls", []), responses):
        tool_messages.append(AIMessage(content=str(response), name=tc["name"]))
        
    return {"messages": tool_messages, "next_node": "researcher"}

def reviewer_node(state):
    """Agen Kritikus: Memeriksa dan merevisi teks agar sesuai standar akademik tinggi."""
    messages = state["messages"]
    llm = get_llm()
    prompt = "Anda adalah Kritikus Jurnal Senior. Perbaiki draf sebelumnya menjadi bahasa akademik yang sangat profesional dan EYD yang benar."
    response = llm.invoke([SystemMessage(content=prompt)] + messages)
    return {"messages": [response], "next_node": "editor"}

def editor_node(state):
    """Agen Editor: Memberikan finalisasi dan menyajikannya ke user."""
    messages = state["messages"]
    llm = get_llm()
    prompt = "Anda adalah Editor Utama. Rangkum dan finalisasi hasil kerja Peneliti dan Kritikus agar siap disajikan ke pengguna dengan ramah."
    response = llm.invoke([SystemMessage(content=prompt)] + messages)
    return {"messages": [response], "next_node": "end"}

# Build the Graph
workflow = StateGraph(AgentState)

workflow.add_node("researcher", researcher_node)
workflow.add_node("tools", tools_node)
workflow.add_node("reviewer", reviewer_node)
workflow.add_node("editor", editor_node)

workflow.set_entry_point("researcher")

# Routing
workflow.add_conditional_edges(
    "researcher",
    lambda x: x["next_node"],
    {
        "tools": "tools",
        "reviewer": "reviewer"
    }
)
workflow.add_edge("tools", "researcher")
workflow.add_edge("reviewer", "editor")
workflow.add_edge("editor", END)

multi_agent_executor = workflow.compile()

def process_multi_agent_chat(message: str):
    """Fungsi utama untuk memanggil sistem Multi-Agent."""
    try:
        inputs = {"messages": [HumanMessage(content=message)]}
        result = multi_agent_executor.invoke(inputs)
        return result["messages"][-1].content
    except Exception as e:
        return f"Error di Multi-Agent System: {str(e)}"