module toolRegistry

export registerTool, registerAllTools, clearTools, listTool

using Dates
using JSON, DataStructures
using ..type



# ── MCP helper functions ────────────────────────────────────────────

"""
Extract text from MCP tool result content array.

Handles JSON-RPC 2.0 result content format:
    {"content": [{"type": "text", "text": "..."}], "isError": false}
"""
function _extract_text_content(result::AbstractDict)::String
    content = get(result, "content", Any[])
    if content isa Vector && !isempty(content)
        lines = String[]
        for block in content
            if block isa Dict && get(block, "type", "") == "text"
                push!(lines, string(get(block, "text", "")))
            end
        end
        if !isempty(lines)
            return join(lines, "\n")
        end
    end
    return JSON.json(result)
end

"""
Wrap an MCP tool definition as an `agentTool`.

The returned tool's `execute` function calls the MCP server's "tools/call"
method with the validated arguments.
"""
function _wrap_mcp_tool(mcpserver, tool_def::AbstractDict{String, Any}; eventSink=nothing, sessionId=nothing)::agentTool
    name = tool_def["name"]
    title = get(tool_def, "title", get(tool_def, "label", name))
    desc = get(tool_def, "description", "")
    input_schema = get(tool_def, "inputSchema", Dict{String,Any}())

    # Normalize inputSchema to OpenAI function format
    if haskey(input_schema, "properties") && input_schema["type"] == "object"
        params = Dict(
            "type" => "object",
            "properties" => input_schema["properties"],
            "required" => get(input_schema, "required", Any[]),
        )
    else
        params = Dict(
            "type" => "object",
            "properties" => Dict{String,Any}(),
            "required" => Any[],
        )
    end

    return agentTool(
      name=name,
      label=title,
      description=desc,
      inputSchema=params,
      execute=(toolCallId::String, args::AbstractDict{String, Any},
                signal::Union{Nothing,abortSignal},
                eventSink) -> begin
        try
          response = mcpserver("tools/call", name, args, sessionId; eventSink=eventSink)

          # Parse JSON-RPC 2.0 response envelope
          if haskey(response, "error")
            rpc_error = response["error"]
            err_msg = get(rpc_error, "message", "Unknown MCP error")
            return agentToolResult(
                [textContent("MCP error: $err_msg")],
                Dict{Any,Any}("error" => err_msg),
                nothing, false
            )
          end

          result_data = get(response, "result", response)
          content_text = _extract_text_content(result_data)
          is_error = get(result_data, "isError", false)

          return agentToolResult(
              [textContent(content_text)],
              Dict{Any,Any}("isError" => is_error),
              nothing, false
          )
        catch e
          errMsg = sprint(showerror, e)
          return agentToolResult(
              [textContent("MCP call error: $errMsg")],
              Dict{Any,Any}("error" => errMsg),
              nothing, false
          )
        end
      end,
      prepareArguments=nothing,
      validateRequiredArgs=nothing,
      parallelToolExecute=false,
    )
end

"""
Discover and register MCP tools into `store.tools`.

Queries the MCP server via `mcpserver("tools/list")` (JSON-RPC 2.0 format),
parses the response envelope, and registers each discovered tool.
Handles pagination via `nextCursor`. Skips tools already registered.

# Returns
- `Int`: number of new tools registered
"""
function register_mcp_tools(tools::OrderedDict{String, agentTool}, mcpserver; eventSink=nothing, sessionId=nothing)::Int
  if mcpserver === nothing
      return 0
  end
  new_count = 0
  try
    eventSink("register_mcp_tools 1")
    response = mcpserver("tools/list", sessionId; eventSink=eventSink)
    eventSink("register_mcp_tools 2")
    # Parse JSON-RPC 2.0 response envelope
    if haskey(response, "error")
        rpc_error = response["error"]
        err_msg = get(rpc_error, "message", "Unknown MCP error")
        println("[toolRegistry] MCP tools/list failed: $err_msg")
        return 0
    end
    if haskey(response, "result")
        response = response["result"]
    end

    tools_array = response["tools"]
    cursor = get(response, "nextCursor", nothing)

    for tool_def in tools_array
      name = tool_def["name"]
      if haskey(tools, name)
        continue
      end

      wrapped = _wrap_mcp_tool(mcpserver, tool_def, eventSink=eventSink, sessionId=sessionId)
      tools[name] = wrapped
      new_count += 1
    end

    # Paginate: fetch remaining tools if nextCursor is present
    while cursor !== nothing && cursor !== ""
      response = sessionId !== nothing ? mcpserver("tools/list", sessionId; eventSink=eventSink) : mcpserver("tools/list"; eventSink=eventSink)
      if haskey(response, "result")
          response = response["result"]
      end
      tools_array = get(response, "tools", Any[])
      cursor = get(response, "nextCursor", nothing)
      for tool_def in tools_array
        name = tool_def["name"]
        if haskey(tools, name)
            continue
        end
        wrapped = _wrap_mcp_tool(mcpserver, tool_def, eventSink=eventSink, sessionId=sessionId)
        tools[name] = wrapped
        new_count += 1
      end
    end

    println("[toolRegistry] Discovered $new_count MCP tools. Total: $(length(tools))")
  catch e
    bt = catch_backtrace()
    err_msg = sprint() do io
      showerror(io, e, bt)
      println(io)
    end

    eventSink(err_msg)

    errMsg = sprint(showerror, e)
    println("[toolRegistry] MCP tools/list failed: $errMsg")
  end

  return new_count
end

"""
    registerTool(tools::OrderedDict{String, agentTool}, tool::agentTool) -> OrderedDict{String, agentTool}

Add `tool` to `tools`, overwriting any existing tool with the same name.

# Arguments
- `tools`: Tool dict to modify
- `tool`: The `agentTool` to register

# Returns
- The same `tools` dict (modified in place)
"""
function registerTool(tools::OrderedDict{String, agentTool}, tool::agentTool; eventSink=nothing
  )::OrderedDict{String, agentTool}
    tools[tool.name] = tool
    println("[toolRegistry] Registered tool: $(tool.name)")
    return tools
end

"""
Remove all tools from the dict.

# Arguments
- `tools`: Tool dict to clear

# Returns
- `nothing`
"""
function clearTools(tools::OrderedDict{String, agentTool})::Nothing
    empty!(tools)
    println("[toolRegistry] Registry cleared")
    return nothing
end

# ── listTools tool (auto-discover new MCP tools at runtime) ─────────

"""
    listTool(tools::OrderedDict{String, agentTool}, mcpserver) -> agentTool

MCP-aware listTools tool (JSON-RPC 2.0 protocol).

First call: queries the MCP server via `mcpserver("tools/list")`, registers
all discovered tools into `tools` (in-place mutation, with
pagination via nextCursor), then returns the full tool list.

Subsequent calls: returns the current list (tools remain registered).
"""
function listTool(tools::OrderedDict{String, agentTool}, mcpserver; eventSink=nothing, sessionId=nothing)::agentTool
    return agentTool(
        name="listTools",
        label="List Tools",
        description="List all available tools. First call discovers and registers all tools from the MCP server. After discovery, new tools become immediately available for use.",
        inputSchema=Dict{String,Any}(
            "type" => "object",
            "properties" => Dict{String,Any}(),
            "required" => Any[]
        ),
        execute=(toolCallId::String, args::AbstractDict{String, Any},
                 signal::Union{Nothing,abortSignal},
                 eventSink) -> begin
            # Discover and register MCP tools (idempotent — skips already registered)
            new_count = register_mcp_tools(tools, mcpserver; eventSink=eventSink, sessionId=sessionId)
            eventSink("tools dump: " * sprint(show, tools))

            # Always include listTools itself in the count
            total = length(tools)

            lines = String[
                "- $(t.name): $(t.label) — $(t.description)"
                for (k, t) in tools
            ]
            result_text = "Available tools ($total):\n" * join(lines, "\n")

            return agentToolResult(
                [textContent(result_text)],
                Dict{Any,Any}("count" => total),
                nothing, false
            )
        end,
        prepareArguments=nothing,
        validateRequiredArgs=nothing,
        parallelToolExecute=false,
    )
end

# ── High-level API ──────────────────────────────────────────────────

"""
    registerAllTools(tools::OrderedDict{String, agentTool}, mcpserver)

Register all tools for an agent:
  1. Auto-discover existing tools from the MCP server
  2. Register the `listTools` tool so the agent can discover new tools at runtime

# Arguments
- `tools`: The tool dict to populate
- `mcpserver`: A callable struct that communicates with the MCP server
"""
function registerAllTools(tools::OrderedDict{String, agentTool}, mcpserver=nothing; eventSink=nothing, sessionId=nothing)
    list_t = listTool(tools, mcpserver; eventSink=eventSink, sessionId=sessionId)
    registerTool(tools, list_t; eventSink=eventSink)
end

end # module