YiemAgent/src/toolRegistry.jl

module toolRegistry

export ToolStore, loadTools, registerTool, getTools, clearTools, listTool

using Dates
using JSON, DataStructures
using ..type

"""
Per-agent isolated tool storage.

Each agent gets its own `ToolStore` so tool registration is independent —
`registerTool(store, tool)` only affects that agent's tool set.

# Fields
- `tools::OrderedDict{String, agentTool}` — keyed by name for O(1) lookup + ordered iteration
- `name::String` — identifier for debugging/logs
"""
struct ToolStore
    tools::OrderedDict{String, agentTool}
    name::String
end

"""
Create a new isolated tool store.

# Keyword Arguments
- `name::String`: Identifier for this store (default: "default")

# Examples
```julia
store = ToolStore(name="agent1")
tools = loadTools(store, "src/tools")
registerTool(store, my_tool)
agent = yiemAgent(tools=getTools(store), llmCall=..., _tool_store=store)
```
"""
function ToolStore(; name::String="default")::ToolStore
    ToolStore(OrderedDict{String, agentTool}(), name)
end

"""
List tool definition — lets the agent query available tools for collision detection
when creating new tools via writeTool.

# Arguments
- `store::ToolStore`: The tool store to list from

Each `ToolStore` gets its own `listTool` instance bound to that store,
so each agent sees only its own tools.
"""
function listTool(store::ToolStore)::agentTool
    return agentTool(
        name = "listTools",
        label = "List Tools",
        description = "List all available tools with their names, labels, and descriptions. Use this before creating a new tool to check for name collisions.",
        inputSchema = Dict{String,Any}(
            "type" => "object",
            "properties" => Dict{String,Any}(),
            "required" => Any[]
        ),
        execute = (toolCallId, args, signal, onPartialResult) -> begin
            tools = getTools(store)
            if isempty(tools)
                result_text = "No tools registered."
            else
                lines = String["- $(t.name): $(t.label) — $(t.description)" for (k, t) in tools]
                result_text = "Available tools:\n" * join(lines, "\n")
            end
            return agentToolResult(
                [textContent(result_text)],
                Dict{Any,Any}("count" => length(tools)),
                nothing, false
            )
        end,
        prepareArguments = nothing,
        validateRequiredArgs = nothing,
        parallelToolExecute = false
    )
end

"""
Load all tool modules from a directory into a specific ToolStore.

Scans `dir` for `.jl` files. Each file must define a function named
`getTool()::agentTool`. Files are sorted alphabetically so tool
registration order is deterministic.

Each `.jl` file is loaded into its own **submodule** so that all functions
defined in the file (`validateRequiredArgs`, `prepareArguments`, `executeTool`,
and any helper functions) are namespaced and never collide with other tools.

# Arguments
- `store::ToolStore`: The tool store to register tools into
- `dir::String`: Directory path to scan for `.jl` tool files

# Returns
- `OrderedDict{String, agentTool}`: All loaded tools keyed by name

# Errors
- Throws `ArgumentError` if a tool file does not define a `getTool` function

# Examples
```julia
julia> store = ToolStore(name="agent1")
julia> tools = loadTools(store, "src/tools")
OrderedDict{String, agentTool} with 3 entries:
  "getWeather" => agentTool(...)
  "getTime" => agentTool(...)
  "listTools" => agentTool(...)
```
"""
function loadTools(store::ToolStore, dir::String)::OrderedDict{String, agentTool}
    if !isdir(dir)
        throw(ArgumentError("Tool directory does not exist: $dir"))
    end

    jl_files = filter(f -> endswith(f, ".jl") && !occursin(r"(?i)registry", f), readdir(dir))
    sort!(jl_files)

    for filename in jl_files
        filepath = joinpath(dir, filename)

        # Derive a unique module name from the filename only (not full path).
        # e.g. "getWeather.jl" -> "_tool_getWeather"
        mod_name = Symbol("_tool_", replace(rstrip(filename, '.'), ".jl" => ""))

        # Build the complete module as a string and eval the parsed code.
        # Julia does not allow `module ... end` inside eval(quote ...),
        # and constructing the module AST by hand is fragile.
        # Instead, we generate the full module source as a string,
        # parse it, and eval the resulting expression.
        # Each tool file declares its own dependencies via `using` statements
        # at the top of the file — the registry only injects `using ..type`
        # to make core types (agentTool, textContent, etc.) available.
        file_content = read(filepath, String)
        module_code = """
        module $(mod_name)
        using ..type
        $(file_content)
        end
        """
        mod = eval(Meta.parse(module_code))

        # Call getTool() via Core.eval in the submodule's scope.
        # This evaluates getTool() entirely within the new module's world,
        # completely avoiding world-age issues — no invokelatest needed.
        # Note: all uses of `tool` must be inside the `try` block because
        # Julia 1.12's SSA form doesn't track `tool` as definitely assigned
        # after a `try-catch` where it's only assigned inside `try`.
        try
            tool = Core.eval(mod, :(getTool()))
            if !(tool isa agentTool)
                throw(ArgumentError(
                    "getTool() in $(filepath) did not return an agentTool instance, got: $(typeof(tool))"
                ))
            end
            store.tools[tool.name] = tool
            println("[$(store.name)] Loaded tool: $(tool.name) ($(tool.label))")
        catch e
            if e isa UndefVarError || occursin("getTool", sprint(showerror, e))
                throw(ArgumentError(
                    "Tool file $(filepath) does not define a `getTool()` function in module $(mod_name). " *
                    "Each tool file must define: function getTool()::agentTool ... end"
                ))
            end
            rethrow(e)
        end
    end

    return store.tools
end

"""
Register a single agentTool into a specific ToolStore.

# Arguments
- `store::ToolStore`: The tool store to register into
- `tool::agentTool`: The tool to register

# Returns
- `OrderedDict{String, agentTool}`: Updated tool dict for this store

# Examples
```julia
julia> store = ToolStore(name="agent1")
julia> registerTool(store, my_tool)
[toolRegistry:agent1] Registered tool: my_tool
```
"""
function registerTool(store::ToolStore, tool::agentTool)::OrderedDict{String, agentTool}
    store.tools[tool.name] = tool
    println("[$(store.name)] Registered tool: $(tool.name)")
    return store.tools
end

"""
Get the registered tools from a specific ToolStore.

Returns the internal `OrderedDict` directly — O(1) lookup by name,
ordered iteration preserving registration order.

# Arguments
- `store::ToolStore`: The tool store to query

# Returns
- `OrderedDict{String, agentTool}`: Tools keyed by name, in registration order

# Examples
```julia
julia> getTools(store)
OrderedDict{String, agentTool} with 3 entries:
  "listTools" => agentTool(...)
  "getWeather" => agentTool(...)
  "getTime" => agentTool(...)
```
"""
function getTools(store::ToolStore)::OrderedDict{String, agentTool}
    return store.tools
end

"""
Clear all registered tools from a specific ToolStore.

# Arguments
- `store::ToolStore`: The tool store to clear

# Returns
- `nothing`

# Examples
```julia
julia> clearTools(store)
[toolRegistry:agent1] Registry cleared
```
"""
function clearTools(store::ToolStore)::Nothing
    empty!(store.tools)
    println("[$(store.name)] Registry cleared")
    return nothing
end

end # module