import json
from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage
from .state import AgentState, tools_config

class ToolSelection(BaseModel):
    tool_name: str | None = Field(
        description="Name of the selected tool, or null if no tool matches."
    )

llm = ChatOpenAI(model="gpt-4o")

selector_llm = llm.with_structured_output(
    ToolSelection
)


def select_tool(state: AgentState):
    system_prompt = f"""
    You are a tool selection agent.
    Available tools:
    {json.dumps(tools_config)}

    Current extracted parameters:
    {json.dumps(state["global_params"])}

    Rules:
    - Select the tool that matches the user's intent.
    - Use conversation history for context.
    - Do not select a tool only because its parameters
      are available.
    - Do not invent tool names.
    - If no tool is appropriate, return null.
    - Do not execute any tool.
    """

    result = selector_llm.invoke([
        SystemMessage(content=system_prompt),
        *state["messages"]
    ])

    allowed_tools = {
        tool["name"] for tool in tools_config
    }

    selected = (
        result.tool_name
        if result.tool_name in allowed_tools
        else None
    )

    return {
        "selected_tool": selected
    }