superintelligence

defmodule Superintelligence.AIAgent.ToolRegistry do
  use GenServer
  require Logger

  @moduledoc """
  Registry for dynamically created tools that agents can use.
  Supports creating, updating, and executing tools at runtime.
  """

  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end

  @impl true
  def init(_opts) do
    {:ok, %{tools: %{}, functions: %{}}}
  end

  # Client API

  def register_tool(name, func, description, parameters) do
    GenServer.call(__MODULE__, {:register_tool, name, func, description, parameters})
  end

  def create_or_update_tool(name, code, description, parameters) do
    GenServer.call(__MODULE__, {:create_or_update_tool, name, code, description, parameters})
  end

  def get_tool(name) do
    GenServer.call(__MODULE__, {:get_tool, name})
  end

  def list_tools do
    GenServer.call(__MODULE__, :list_tools)
  end

  def execute_tool(name, args) do
    GenServer.call(__MODULE__, {:execute_tool, name, args}, 30_000)
  end

  # Server Callbacks

  @impl true
  def handle_call({:register_tool, name, func, description, parameters}, _from, state) do
    tool = %{
      type: "function",
      function: %{
        name: name,
        description: description,
        parameters: %{
          type: "object",
          properties: parameters,
          required: Map.keys(parameters)
        }
      }
    }
    
    new_state = state
      |> put_in([:tools, name], tool)
      |> put_in([:functions, name], func)
    
    Logger.info("Registered tool: #{name}")
    {:reply, {:ok, tool}, new_state}
  end

  @impl true
  def handle_call({:create_or_update_tool, name, code, description, parameters}, _from, state) do
    case compile_and_load_function(name, code) do
      {:ok, func} ->
        handle_call({:register_tool, name, func, description, parameters}, nil, state)
      
      {:error, reason} ->
        {:reply, {:error, "Error creating/updating tool '#{name}': #{reason}"}, state}
    end
  end

  @impl true
  def handle_call({:get_tool, name}, _from, state) do
    tool = get_in(state, [:tools, name])
    {:reply, tool, state}
  end

  @impl true
  def handle_call(:list_tools, _from, state) do
    tools = Map.values(state.tools)
    {:reply, tools, state}
  end

  @impl true
  def handle_call({:execute_tool, name, args}, _from, state) do
    case get_in(state, [:functions, name]) do
      nil ->
        {:reply, {:error, "Tool '#{name}' not found"}, state}
      
      func ->
        result = safe_execute(func, args)
        {:reply, result, state}
    end
  end

  # Private Functions

  defp compile_and_load_function(name, code) do
    try do
      # Create a module dynamically
      module_name = String.to_atom("Elixir.Superintelligence.AIAgent.DynamicTools.#{name}")
      
      # Wrap the code in a module definition
      module_code = """
      defmodule #{inspect(module_name)} do
        def execute(args) do
          #{code}
        end
      end
      """
      
      # Compile the module
      Code.eval_string(module_code)
      
      # Create a function that calls the module
      func = fn args -> apply(module_name, :execute, [args]) end
      
      {:ok, func}
    rescue
      e -> {:error, Exception.message(e)}
    end
  end

  defp safe_execute(func, args) do
    try do
      result = func.(args)
      {:ok, result}
    rescue
      e -> {:error, "Error executing tool: #{Exception.message(e)}"}
    end
  end

  # Built-in Tools

  def init_builtin_tools do
    # Install package tool
    register_tool(
      "install_package",
      fn %{"package_name" => package_name} ->
        case System.cmd("mix", ["deps.get", package_name]) do
          {output, 0} -> {:ok, "Package '#{package_name}' installed successfully. Output: #{output}"}
          {output, _} -> {:error, "Failed to install package '#{package_name}': #{output}"}
        end
      end,
      "Installs an Elixir package using mix.",
      %{
        "package_name" => %{
          "type" => "string",
          "description" => "The name of the package to install"
        }
      }
    )

    # Task completed tool
    register_tool(
      "task_completed",
      fn _args -> {:ok, "Task marked as completed."} end,
      "Marks the current task as completed.",
      %{}
    )

    # Execute bash command tool
    register_tool(
      "execute_bash",
      fn %{"command" => command} ->
        case System.cmd("bash", ["-c", command]) do
          {output, 0} -> {:ok, output}
          {output, code} -> {:error, "Command failed with code #{code}: #{output}"}
        end
      end,
      "Executes a bash command",
      %{
        "command" => %{
          "type" => "string",
          "description" => "The bash command to execute"
        }
      }
    )

    # Read file tool
    register_tool(
      "read_file",
      fn %{"path" => path} ->
        case File.read(path) do
          {:ok, content} -> {:ok, content}
          {:error, reason} -> {:error, "Failed to read file: #{reason}"}
        end
      end,
      "Reads a file from the filesystem",
      %{
        "path" => %{
          "type" => "string",
          "description" => "The path to the file to read"
        }
      }
    )

    # Write file tool
    register_tool(
      "write_file",
      fn %{"path" => path, "content" => content} ->
        case File.write(path, content) do
          :ok -> {:ok, "File written successfully"}
          {:error, reason} -> {:error, "Failed to write file: #{reason}"}
        end
      end,
      "Writes content to a file",
      %{
        "path" => %{
          "type" => "string",
          "description" => "The path to write to"
        },
        "content" => %{
          "type" => "string",
          "description" => "The content to write"
        }
      }
    )
  end
end