defmodule MyApp.AI.Actions.GetUserCount do
  @moduledoc "Tool: returns how many users signed up in a given period."

  use Jido.Action,
    name: "get_user_count",
    description: "Get the count of users who signed up in a period.",
    schema: [
      period: [
        type: {:in, ["today", "week", "month", "all"]},
        default: "week",
        doc: "Time window: today, week, month, or all."
      ]
    ]

  import Ecto.Query
  alias MyApp.{Accounts.User, Repo}

  @impl Jido.Action
  def run(%{period: period}, _context) do
    count =
      User
      |> filter_by_period(period)
      |> Repo.aggregate(:count, :id)

    {:ok, %{period: period, count: count}}
  end

  defp filter_by_period(query, "all"), do: query

  defp filter_by_period(query, period) do
    cutoff =
      case period do
        "today" -> DateTime.add(DateTime.utc_now(), -1, :day)
        "week" -> DateTime.add(DateTime.utc_now(), -7, :day)
        "month" -> DateTime.add(DateTime.utc_now(), -30, :day)
      end

    from(u in query, where: u.inserted_at >= ^cutoff)
  end
end

defmodule MyApp.AI.ToolRegistry do
  @moduledoc """
  Single source of truth for AI tools. Add an action here and it's
  instantly available to both ReqLLM tool-calling and MCP clients.
  """

  @actions [
    MyApp.AI.Actions.GetUserCount,
    MyApp.AI.Actions.SearchUsers
    # add more actions here
  ]

  def actions, do: @actions

  @doc "Build ReqLLM tool list from registered Jido actions."
  def req_llm_tools do
    Enum.map(@actions, fn action ->
      meta = action.__action_metadata__()

      ReqLLM.Tool.new!(
        name: meta.name,
        description: meta.description,
        parameter_schema: meta.schema,
        callback: fn args -> execute(meta.name, args) end
      )
    end)
  end

  @doc "Execute a tool by name. Atomizes string-keyed JSON args."
  def execute(tool_name, args) when is_map(args) do
    action = Enum.find(@actions, &(&1.__action_metadata__().name == tool_name))

    if action do
      atomized = atomize_keys(args, action.__action_metadata__().schema)
      action.run(atomized, %{})
    else
      {:error, "Unknown tool: #{tool_name}"}
    end
  end

  defp atomize_keys(args, schema) do
    valid_keys = Keyword.keys(schema)

    Map.new(args, fn {k, v} ->
      key = if is_binary(k), do: String.to_existing_atom(k), else: k
      if key in valid_keys, do: {key, v}, else: {key, v}
    end)
  end
end

defmodule MyApp.AI.Chat do
  @moduledoc "Multi-turn chat with tool calling."

  @model "anthropic:claude-haiku-4-5-20251001"

  def respond(user_message, history \\ []) do
    messages =
      history ++
        [
          %{role: "system", content: system_prompt()},
          %{role: "user", content: user_message}
        ]

    ReqLLM.put_key(:anthropic_api_key, System.get_env("ANTHROPIC_API_KEY"))

    ReqLLM.generate_text(@model, messages,
      tools: MyApp.AI.ToolRegistry.req_llm_tools(),
      max_steps: 5
    )
  end

  defp system_prompt do
    """
    You are an admin assistant. Use the available tools to answer questions
    about user counts and activity. If a tool fails, explain why in plain English.
    """
  end
end

# Drop into IEx and try it
iex> MyApp.AI.Chat.respond("How many users signed up this week?")
# → LLM calls GetUserCount.run(%{period: "week"}, _) under the hood
# → returns something like "42 users signed up in the last week."

iex> MyApp.AI.Chat.respond("Compare this week vs last month.")
# → LLM makes TWO tool calls (week, then month) and synthesizes the answer

def deps do
  [
    {:jido, "~> 2.0"},
    {:req, "~> 0.5"},
    {:req_llm, "~> 1.0"}
  ]
end