"""StateGraph: llm -> route -> tools or end; tools -> llm."""

from langchain_core.messages import SystemMessage
from langchain_core.language_models.chat_models import BaseChatModel
from langgraph.graph import END, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition

from foremet.agent.state import AgentState


def _get_llm(
    provider: str,
    model: str,
    api_key: str | None = None,
    ollama_base_url: str | None = None,
) -> BaseChatModel:
    """Create chat model for the given provider. Uses env for keys if not passed."""
    import os
    if provider == "openai":
        from langchain_openai import ChatOpenAI
        key = api_key or os.environ.get("OPENAI_API_KEY")
        return ChatOpenAI(model=model, api_key=key)
    if provider == "anthropic":
        from langchain_anthropic import ChatAnthropic
        key = api_key or os.environ.get("ANTHROPIC_API_KEY")
        return ChatAnthropic(model=model, api_key=key)
    if provider == "gemini":
        from langchain_google_genai import ChatGoogleGenerativeAI
        key = api_key or os.environ.get("GOOGLE_API_KEY")
        return ChatGoogleGenerativeAI(model=model, google_api_key=key)
    if provider == "ollama":
        from langchain_ollama import ChatOllama
        base_url = ollama_base_url or os.environ.get("OLLAMA_BASE_URL")
        kwargs = {"model": model}
        if base_url:
            kwargs["base_url"] = base_url
        return ChatOllama(**kwargs)
    raise ValueError(f"Unknown LLM provider: {provider}")


def build_graph(
    system_prompt: str,
    tools: list,
    llm_provider: str = "openai",
    model: str = "gpt-4o-mini",
    api_key: str | None = None,
    ollama_base_url: str | None = None,
):
    """Build and compile the agent StateGraph. tools: list of LangChain tools."""
    llm = _get_llm(llm_provider, model, api_key=api_key, ollama_base_url=ollama_base_url).bind_tools(tools)

    def agent_node(state: AgentState) -> dict:
        messages = state["messages"]
        if not any(isinstance(m, SystemMessage) for m in messages):
            messages = [SystemMessage(content=system_prompt)] + list(messages)
        response = llm.invoke(messages)
        return {"messages": [response]}

    tool_node = ToolNode(tools)
    builder = StateGraph(AgentState)
    builder.add_node("agent", agent_node)
    builder.add_node("tools", tool_node)
    builder.add_edge("__start__", "agent")
    builder.add_conditional_edges("agent", tools_condition, {"tools": "tools", "__end__": END})
    builder.add_edge("tools", "agent")
    return builder.compile()
