sagents

defmodule Sagents.Middleware.TodoList do
  @moduledoc """
  Middleware that adds TODO list management capabilities to agents.

  Provides the `write_todos` tool for creating, updating, and managing TODO items
  during complex multi-step tasks.

  ## Usage

      {:ok, agent} = Agent.new(
        model: model,
        middleware: [TodoList]
      )

      # Agent can now use write_todos tool to manage tasks

  ## Configuration

  - `:display_text` - Label shown in UI when the `write_todos` tool runs.
    Defaults to `"Updating task list"`. Useful when the todo list is used
    internally by the agent (e.g. self-organization during an onboarding
    flow) and the default label would leak implementation detail to users.

  - `:inline` - Controls how todo updates surface to the UI. Defaults to
    `false`. See "Display modes" below.

  ## Display modes

  Every successful `write_todos` call produces a state change. Two delivery
  channels are available — the choice is purely about *where* in the UI the
  todo list appears:

  - **Default (`inline: false`)** — only emits the `{:todos_updated, todos}`
    event over the AgentServer's broadcast channel. Subscribers (typically a
    sidebar or status panel in the LiveView) re-render from that snapshot.
    The chat transcript itself shows nothing about the todo list; the
    transcript stays focused on the agent's textual reasoning and tool
    calls. Best when the UI has dedicated "task list" real estate, or when
    todos are an internal planning detail you don't want surfacing to the
    end user.

  - **Inline (`inline: true`)** — emits the same `{:todos_updated, todos}`
    event *and additionally* persists a synthetic display message
    (content_type `"todo_snapshot"`) into the conversation transcript. Each
    snapshot is a separate message, so the chat shows the plan evolving
    over time threaded with the agent's other output. Best when there is no
    sidebar, or when seeing the plan unfold inside the conversation is part
    of the desired UX.

  Inline mode requires the AgentServer to be configured with a
  `Sagents.DisplayMessagePersistence` module that implements
  `save_synthetic_message/3`. Without it the inline-mode call is a silent
  no-op (the existing `{:todos_updated, _}` broadcast still fires).

  ## Usage

      # default: only the {:todos_updated, todos} broadcast — sidebar UIs
      {:ok, agent} = Agent.new(middleware: [TodoList])

      # custom label for a user-facing flow where todos are internal
      {:ok, agent} = Agent.new(
        middleware: [{TodoList, display_text: "Updating my notes"}]
      )

      # inline mode: also emit a todo_snapshot display message per update
      # so the plan appears inline in the chat transcript
      {:ok, agent} = Agent.new(middleware: [{TodoList, inline: true}])

      # inline mode + custom label, both options together
      {:ok, agent} = Agent.new(
        middleware: [{TodoList, inline: true, display_text: "Replanning"}]
      )
  """

  @behaviour Sagents.Middleware

  alias Sagents.{AgentServer, State, Todo}
  alias LangChain.Function

  @default_display_text "Updating task list"

  @system_prompt """
  ## `write_todos`

  You have access to the `write_todos` tool to help you manage and plan complex objectives.
  Use this tool for complex objectives to ensure that you are tracking each necessary step and giving the user visibility into your progress.
  This tool is very helpful for planning complex objectives, and for breaking down these larger complex objectives into smaller steps.

  **CRITICAL COMPLETION REQUIREMENT**: Once you create todos, you MUST work through ALL items until they are marked as completed or cancelled.
  Do not stop execution while todos remain in pending or in_progress status. Continue calling tools and working until the entire list is completed.

  When you have incomplete todos:
  - Always check your current todo list state before considering your work done
  - Always work on the next pending/in_progress item
  - Never stop with unfinished work - this is considered an incomplete response
  - If you cannot complete a todo, mark it as cancelled and explain why in the todo content

  It is critical that you mark todos as completed as soon as you are done with a step. Do not batch up multiple steps before marking them as completed.
  For simple objectives that only require a few steps, it is better to just complete the objective directly and NOT use this tool.
  Writing todos takes time and tokens, use it when it is helpful for managing complex many-step problems! But not for simple few-step requests.

  ## Important To-Do List Usage Notes to Remember
  - The `write_todos` tool should never be called multiple times in parallel.
  - Don't be afraid to revise the To-Do list as you go. New information may reveal new tasks that need to be done, or old tasks that are irrelevant.
  - Before finishing your work, verify that all todos are either completed or cancelled. Stopping with pending/in_progress todos is an error.
  """

  @tool_description """
  Use this tool to create and manage a structured task list for your current work session. This helps you track progress, organize complex tasks, and demonstrate thoroughness to the user.

  Only use this tool if you think it will be helpful in staying organized. If the user's request is trivial and takes less than 3 steps, it is better to NOT use this tool and just do the task directly.

  ## When to Use This Tool
  Use this tool in these scenarios:

  1. Complex multi-step tasks - When a task requires 3 or more distinct steps or actions
  2. Non-trivial and complex tasks - Tasks that require careful planning or multiple operations
  3. User explicitly requests todo list - When the user directly asks you to use the todo list
  4. User provides multiple tasks - When users provide a list of things to be done (numbered or comma-separated)
  5. The plan may need future revisions or updates based on results from the first few steps

  ## How to Use This Tool
  1. When you start working on a task - Mark it as in_progress BEFORE beginning work.
  2. After completing a task - Mark it as completed and add any new follow-up tasks discovered during implementation.
  3. You can also update future tasks, such as deleting them if they are no longer necessary, or adding new tasks that are necessary. Don't change previously completed tasks.
  4. You can make several updates to the todo list at once. For example, when you complete a task, you can mark the next task you need to start as in_progress.

  ## When NOT to Use This Tool
  It is important to skip using this tool when:
  1. There is only a single, straightforward task
  2. The task is trivial and tracking it provides no benefit
  3. The task can be completed in less than 3 trivial steps
  4. The task is purely conversational or informational

  ## Task States and Management

  1. **Task States**: Use these states to track progress:
     - pending: Task not yet started
     - in_progress: Currently working on (you can have multiple tasks in_progress at a time if they are not related to each other and can be run in parallel)
     - completed: Task finished successfully

  2. **Task Management**:
     - Update task status in real-time as you work
     - Mark tasks complete IMMEDIATELY after finishing (don't batch completions)
     - Complete current tasks before starting new ones
     - Remove tasks that are no longer relevant from the list entirely
     - IMPORTANT: When you write this todo list, you should mark your first task (or tasks) as in_progress immediately!
     - IMPORTANT: Unless all tasks are completed, you should always have at least one task in_progress to show the user that you are working on something.

  3. **Task Completion Requirements**:
     - ONLY mark a task as completed when you have FULLY accomplished it
     - If you encounter errors, blockers, or cannot finish, keep the task as in_progress
     - When blocked, create a new task describing what needs to be resolved
     - Never mark a task as completed if:
       - There are unresolved issues or errors
       - Work is partial or incomplete
       - You encountered blockers that prevent completion
       - You couldn't find necessary resources or dependencies
       - Quality standards haven't been met

  4. **Task Breakdown**:
     - Create specific, actionable items
     - Break complex tasks into smaller, manageable steps
     - Use clear, descriptive task names

  Being proactive with task management demonstrates attentiveness and ensures you complete all requirements successfully
  Remember: If you only need to make a few tool calls to complete a task, and it is clear what you need to do, it is better to just do the task directly and NOT call this tool at all.
  """

  @impl true
  def init(opts) do
    config = %{
      display_text: Keyword.get(opts, :display_text, @default_display_text),
      inline: Keyword.get(opts, :inline, false)
    }

    {:ok, config}
  end

  @impl true
  def system_prompt(_config) do
    @system_prompt
  end

  @impl true
  def tools(config) do
    [build_write_todos_tool(config)]
  end

  @impl true
  def on_server_start(state, _config) do
    # Broadcast initial todos when AgentServer starts
    # This handles restored conversations showing their saved/restored todos
    broadcast_todos(state.agent_id, state.todos)
    {:ok, state}
  end

  # Build the tool function - called at runtime to avoid compile-time ordering issues
  defp build_write_todos_tool(config) do
    display_text =
      case config do
        %{display_text: text} when is_binary(text) -> text
        _other -> @default_display_text
      end

    inline =
      case config do
        %{inline: value} when is_boolean(value) -> value
        _other -> false
      end

    Function.new!(%{
      name: "write_todos",
      description: @tool_description,
      display_text: display_text,
      parameters_schema: %{
        type: "object",
        properties: %{
          merge: %{
            type: "boolean",
            description: "If true, merge with existing todos by ID. If false, replace all todos."
          },
          todos: %{
            type: "array",
            description: "Array of TODO items to set or update",
            items: %{
              type: "object",
              properties: %{
                id: %{
                  type: "integer",
                  description:
                    "Positive integer that is unique within this todo list. If omitted, IDs are assigned from list order (first item = 1)."
                },
                content: %{
                  type: "string",
                  description: "Description of the task"
                },
                status: %{
                  type: "string",
                  enum: ["pending", "in_progress", "completed", "cancelled"],
                  description: "Current status of the task"
                }
              },
              required: ["id", "content", "status"]
            },
            min_items: 1
          }
        },
        required: ["merge", "todos"]
      },
      function: fn args, context -> execute_write_todos(args, context, inline) end
    })
  end

  defp execute_write_todos(args, context, inline) do
    # Args come from LLM as string-keyed map
    merge = get_boolean_arg(args, "merge", false)
    todos_data = get_arg(args, "todos")

    # Parse TODO data from tool call
    case parse_todos(todos_data) do
      {:ok, parsed_todos} ->
        # Compute the post-update state, then apply the auto-clear that
        # collapses a fully-completed list. We hold both: the inline-mode
        # snapshot uses the pre-clear state (so the chat preserves the
        # final celebratory checked list), while the actual state delta
        # and sidebar broadcast use the post-clear state (so the agent
        # and sidebar see the list as resolved).
        state_after_update = update_todos(context.state, parsed_todos, merge)
        updated_state = clear_if_all_completed(state_after_update)

        # Broadcast todos immediately for real-time UI updates
        # This is the point where todos are actually committed to state
        broadcast_todos(context.state.agent_id, updated_state.todos)

        # Inline mode: persist a synthetic display message so the snapshot
        # appears inline in the chat transcript. Snapshot the pre-clear
        # state so the user sees the completed list, not an empty card.
        # The AgentServer guards on display_message_persistence +
        # conversation_id being configured; without them the call is a
        # silent no-op.
        if inline do
          save_todo_snapshot(context.state.agent_id, state_after_update.todos)
        end

        # Return only the state changes (todos), not messages
        state_delta = %State{todos: updated_state.todos}
        # Use updated_state.todos (after auto-cleanup) for response message
        {:ok, format_response(updated_state.todos, merge), state_delta}

      {:error, reason} ->
        {:error, "Failed to parse TODOs: #{reason}"}
    end
  end

  defp save_todo_snapshot(nil, _todos), do: :ok

  defp save_todo_snapshot(agent_id, todos) do
    # Mirror the broadcast_todos guard: only emit when AgentServer is running.
    # AgentServer.save_synthetic_message_from is itself a no-op when the
    # server lacks display_message_persistence/conversation_id, but checking
    # here keeps unit tests free of a live AgentServer process.
    case GenServer.whereis(AgentServer.get_name(agent_id)) do
      nil ->
        :ok

      _pid ->
        AgentServer.save_synthetic_message_from(agent_id, %{
          message_type: "system",
          content_type: "todo_snapshot",
          content: build_todo_snapshot_content(todos)
        })
    end
  end

  defp build_todo_snapshot_content(todos) do
    %{
      "todos" =>
        Enum.map(todos, fn todo ->
          %{
            "id" => todo.id,
            "content" => todo.content,
            "status" => to_string(todo.status)
          }
        end),
      "summary" => %{
        "total" => length(todos),
        "pending" => count_status(todos, :pending),
        "in_progress" => count_status(todos, :in_progress),
        "completed" => count_status(todos, :completed),
        "cancelled" => count_status(todos, :cancelled)
      }
    }
  end

  defp count_status(todos, status) do
    Enum.count(todos, fn todo -> todo.status == status end)
  end

  defp get_arg(args, key) when is_map(args) do
    args[key] || args[String.to_atom(key)]
  end

  defp get_boolean_arg(args, key, default) when is_map(args) do
    case get_arg(args, key) do
      nil -> default
      val when is_boolean(val) -> val
      "true" -> true
      "false" -> false
      _other -> default
    end
  end

  defp parse_todos(todos_data) when is_list(todos_data) do
    Todo.list_from_maps(todos_data)
  end

  defp parse_todos(_other), do: {:error, "todos must be an array"}

  defp update_todos(%State{} = state, new_todos, true = _merge) do
    # Merge mode: update existing TODOs by ID
    existing_map =
      state.todos
      |> Enum.map(fn todo ->
        case todo do
          %Todo{} = t -> {t.id, t}
          %{"id" => id} = m -> {id, m}
          %{id: id} = m -> {id, m}
        end
      end)
      |> Map.new()

    new_map =
      new_todos
      |> Enum.map(fn todo -> {todo.id, todo} end)
      |> Map.new()

    merged_map = Map.merge(existing_map, new_map)

    todos_list =
      merged_map
      |> Map.values()
      |> Enum.map(&ensure_todo_struct/1)
      |> Enum.sort_by(& &1.id)

    %{state | todos: todos_list}
  end

  defp update_todos(%State{} = state, new_todos, false = _merge) do
    # Replace mode: use only new TODOs
    todos_list =
      new_todos
      |> Enum.map(&ensure_todo_struct/1)
      |> Enum.sort_by(& &1.id)

    %{state | todos: todos_list}
  end

  defp clear_if_all_completed(%State{todos: []} = state), do: state

  defp clear_if_all_completed(%State{todos: todos} = state) do
    all_completed? = Enum.all?(todos, fn todo -> todo.status == :completed end)

    if all_completed? do
      %{state | todos: []}
    else
      state
    end
  end

  defp ensure_todo_struct(%Todo{} = todo), do: todo

  defp ensure_todo_struct(map) when is_map(map) do
    case Todo.from_map(map) do
      {:ok, todo} -> todo
      {:error, _changeset} -> raise "Invalid TODO data: #{inspect(map)}"
    end
  end

  defp format_response(todos, merge) do
    case todos do
      [] ->
        "TODO list cleared - all tasks completed"

      _other ->
        mode = if merge, do: "merged", else: "replaced"
        count = Enum.count(todos)

        summary =
          todos
          |> Enum.group_by(& &1.status)
          |> Enum.map_join(", ", fn {status, items} -> "#{length(items)} #{status}" end)

        "Successfully #{mode} #{count} TODO(s): #{summary}"
    end
  end

  # Broadcast todos via AgentServer's PubSub infrastructure
  # Uses the public API so middleware doesn't need direct PubSub access
  defp broadcast_todos(nil, _todos), do: :ok

  defp broadcast_todos(agent_id, todos) do
    # Only broadcast if AgentServer is running
    # This check is necessary because execute_write_todos can be called
    # outside of AgentServer context (e.g., in tests or standalone agents)
    case GenServer.whereis(AgentServer.get_name(agent_id)) do
      nil -> :ok
      _pid -> AgentServer.publish_event_from(agent_id, {:todos_updated, todos})
    end
  end
end