from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, BaseMessage
from prompts import main_prompt, research_agent_prompt   # تأكد إن analyzer_prompt متعرف
from tools import search
from state import search_keys, AgentState
import sys
sys.stdout.reconfigure(encoding="utf-8")


# -------- Models -------- #
boss_model = ChatOpenAI(
    model="meta-llama/llama-4-maverick:free",
    openai_api_key="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",
    base_url="https://openrouter.ai/api/v1",
    temperature=0.5,
    max_tokens=1024,
    top_p=0.5,
).with_structured_output(search_keys)

analyzer_model = ChatOpenAI(
    model="openrouter/sonoma-sky-alpha",
    openai_api_key="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",
    base_url="https://openrouter.ai/api/v1",
)


# -------- Nodes -------- #
def boss_node(state: AgentState) -> AgentState:
    if not state.get("messages"):
        raise ValueError("No messages found in state. Please provide at least one HumanMessage.")

    last_message: BaseMessage = state["messages"][-1]
    user_text = getattr(last_message, "content", str(last_message))

    query = boss_model.invoke(f"{main_prompt}\nUser: {user_text}")
    result = search(query.query, query.Topic)
    state["search_content"] = result
    return state


def analyzer_node(state: AgentState) -> AgentState:
    state["search_results"] = analyzer_model.invoke(
        f"{research_agent_prompt}\n{state['search_content']}"
    )
    return state


# -------- Graph -------- #
graph = StateGraph(AgentState)
graph.add_node("boss", boss_node)
graph.add_node("analyzer", analyzer_node)
graph.add_edge("boss", "analyzer")
graph.add_edge("analyzer", END)
graph.set_entry_point("boss")
app = graph.compile()


# -------- Run Tests -------- #
if __name__ == "__main__":
    # Test stream with Chinese input
    for event in app.stream({"messages": [HumanMessage(content="what is best player in football in all time ")]}):
        if "analyzer" in event:
            print(":: answer is -->")

    # Test invoke with football query
    result = app.invoke(
        {"messages": [HumanMessage(content="what is capital of egypt")]}
    )
    print(result["search_results"].content)