towerops/lib/towerops/agents.ex

defmodule Towerops.Agents do
  @moduledoc """
  The Agents context for managing remote SNMP polling agents.

  This context provides functions for creating and managing agent authentication tokens,
  assigning devices to agents, and verifying agent credentials.
  """

  import Ecto.Query, warn: false

  alias Towerops.Agents.AgentAssignment
  alias Towerops.Agents.AgentToken
  alias Towerops.Devices.Device
  alias Towerops.Repo

  ## Token management

  @doc """
  Creates a new agent token for an organization.

  Returns {:ok, agent_token, token_string} where token_string is the plain token
  that should be shown to the user ONCE and never stored.

  ## Examples

      iex> create_agent_token(org_id, "Datacenter Agent")
      {:ok, %AgentToken{}, "abc123..."}

      iex> create_agent_token(nil, "")
      {:error, %Ecto.Changeset{}}

  """
  @spec create_agent_token(String.t(), String.t()) :: {:ok, AgentToken.t(), String.t()} | {:error, Ecto.Changeset.t()}
  def create_agent_token(organization_id, name) do
    {token, changeset} = AgentToken.build_token(organization_id, name)

    case Repo.insert(changeset) do
      {:ok, agent_token} -> {:ok, agent_token, token}
      {:error, changeset} -> {:error, changeset}
    end
  end

  @doc """
  Creates a new cloud poller agent token (superadmin only).

  Cloud pollers are application-wide agents that are not tied to any organization
  and can poll devices for any organization.

  Returns {:ok, agent_token, token_string} where token_string is the plain token
  that should be shown to the user ONCE and never stored.

  ## Examples

      iex> create_cloud_poller("Cloud Poller 1")
      {:ok, %AgentToken{is_cloud_poller: true}, "abc123..."}

  """
  @spec create_cloud_poller(String.t()) :: {:ok, AgentToken.t(), String.t()} | {:error, Ecto.Changeset.t()}
  def create_cloud_poller(name) do
    {token, changeset} = AgentToken.build_token(nil, name, is_cloud_poller: true)

    case Repo.insert(changeset) do
      {:ok, agent_token} -> {:ok, agent_token, token}
      {:error, changeset} -> {:error, changeset}
    end
  end

  @doc """
  Lists all agent tokens for an organization.

  Returns a list of agent tokens ordered by creation date (newest first).
  Does not include cloud pollers.

  ## Examples

      iex> list_organization_agent_tokens(org_id)
      [%AgentToken{}, ...]

  """
  @spec list_organization_agent_tokens(String.t()) :: [AgentToken.t()]
  def list_organization_agent_tokens(organization_id) do
    AgentToken
    |> where([t], t.organization_id == ^organization_id and t.is_cloud_poller == false)
    |> order_by([t], desc: t.inserted_at)
    |> Repo.all()
  end

  @doc """
  Lists all cloud poller agent tokens (superadmin only).

  Returns a list of cloud poller agent tokens ordered by creation date (newest first).

  ## Examples

      iex> list_cloud_pollers()
      [%AgentToken{is_cloud_poller: true}, ...]

  """
  @spec list_cloud_pollers() :: [AgentToken.t()]
  def list_cloud_pollers do
    AgentToken
    |> where([t], t.is_cloud_poller == true)
    |> order_by([t], desc: t.inserted_at)
    |> Repo.all()
  end

  @doc """
  Counts the number of devices directly assigned to an agent.

  This only counts devices with explicit assignments, not devices
  that inherit the agent from site or organization defaults.

  ## Examples

      iex> count_assigned_devices(agent_token_id)
      5

  """
  @spec count_assigned_devices(Ecto.UUID.t()) :: non_neg_integer()
  def count_assigned_devices(agent_token_id) do
    AgentAssignment
    |> where([a], a.agent_token_id == ^agent_token_id and a.enabled == true)
    |> Repo.aggregate(:count)
  end

  @doc """
  Gets a single agent token by ID.

  Raises `Ecto.NoResultsError` if the token does not exist.

  ## Examples

      iex> get_agent_token!(id)
      %AgentToken{}

      iex> get_agent_token!("nonexistent")
      ** (Ecto.NoResultsError)

  """
  @spec get_agent_token!(Ecto.UUID.t()) :: AgentToken.t()
  def get_agent_token!(id), do: Repo.get!(AgentToken, id)

  @doc """
  Updates the heartbeat timestamp and metadata for an agent token.

  This is called automatically when an agent makes an API request.

  ## Examples

      iex> update_agent_token_heartbeat(token_id, "192.168.1.1", %{version: "0.1.0"})
      {1, nil}

  """
  @spec update_agent_token_heartbeat(Ecto.UUID.t(), String.t() | nil, map() | nil) ::
          {non_neg_integer(), nil | [term()]}
  def update_agent_token_heartbeat(agent_token_id, ip_address, metadata \\ nil) do
    now = DateTime.truncate(DateTime.utc_now(), :second)

    # Only update metadata if explicitly provided (not nil and not empty)
    updates =
      case metadata do
        nil -> [last_seen_at: now, last_ip: ip_address]
        meta when map_size(meta) > 0 -> [last_seen_at: now, last_ip: ip_address, metadata: meta]
        _ -> [last_seen_at: now, last_ip: ip_address]
      end

    AgentToken
    |> where([t], t.id == ^agent_token_id)
    |> Repo.update_all(set: updates)
  end

  @doc """
  Verifies an agent token and returns the agent token record if valid.

  Returns {:ok, agent_token} if the token is valid and enabled,
  otherwise {:error, :invalid_token}.

  ## Examples

      iex> verify_agent_token("valid_token")
      {:ok, %AgentToken{}}

      iex> verify_agent_token("invalid_token")
      {:error, :invalid_token}

  """
  @spec verify_agent_token(String.t()) :: {:ok, AgentToken.t()} | {:error, :invalid_token}
  def verify_agent_token(token) do
    with {:ok, token_string} <- AgentToken.verify_token(token),
         %AgentToken{} = agent_token <-
           Repo.get_by(AgentToken, token: token_string, enabled: true) do
      {:ok, agent_token}
    else
      _ -> {:error, :invalid_token}
    end
  end

  @doc """
  Updates an agent token with the given attributes.

  Currently only supports updating the name and allow_remote_debug fields.
  Other fields like the token itself cannot be changed (use regenerate_token instead).

  ## Examples

      iex> update_agent_token(agent_token, %{name: "New Name"})
      {:ok, %AgentToken{name: "New Name"}}

      iex> update_agent_token(agent_token, %{name: ""})
      {:error, %Ecto.Changeset{}}

  """
  @spec update_agent_token(AgentToken.t(), map()) :: {:ok, AgentToken.t()} | {:error, Ecto.Changeset.t()}
  def update_agent_token(%AgentToken{} = agent_token, attrs) do
    agent_token
    |> AgentToken.update_changeset(attrs)
    |> Repo.update()
  end

  @doc """
  Toggles remote debugging for an agent.

  Creates an audit log entry for accountability.

  ## Examples

      iex> toggle_agent_debug(agent_token, true, user)
      {:ok, %AgentToken{allow_remote_debug: true}}

      iex> toggle_agent_debug(agent_token, false, user)
      {:ok, %AgentToken{allow_remote_debug: false}}

  """
  @spec toggle_agent_debug(AgentToken.t(), boolean(), Towerops.Accounts.User.t()) ::
          {:ok, AgentToken.t()} | {:error, Ecto.Changeset.t()}
  def toggle_agent_debug(%AgentToken{} = agent_token, enabled, user) do
    case update_agent_token(agent_token, %{allow_remote_debug: enabled}) do
      {:ok, updated_agent} ->
        # Create audit log
        Towerops.Admin.create_audit_log(%{
          superuser_id: user.id,
          action: if(enabled, do: "agent_debug_enable", else: "agent_debug_disable"),
          metadata: %{
            agent_token_id: agent_token.id,
            agent_name: agent_token.name,
            organization_id: agent_token.organization_id
          }
        })

        {:ok, updated_agent}

      error ->
        error
    end
  end

  @doc """
  Revokes an agent token by setting enabled to false.

  ## Examples

      iex> revoke_agent_token(id)
      {:ok, %AgentToken{enabled: false}}

  """
  @spec revoke_agent_token(Ecto.UUID.t()) :: {:ok, AgentToken.t()} | {:error, Ecto.Changeset.t()}
  def revoke_agent_token(id) do
    AgentToken
    |> Repo.get!(id)
    |> Ecto.Changeset.change(enabled: false)
    |> Repo.update()
  end

  @doc """
  Deletes an agent token and cleans up all references.

  This function:
  1. Removes all direct device assignments for this agent
  2. Removes this agent from any sites using it as default
  3. Removes this agent from any organizations using it as default
  4. Deletes the agent token

  Devices that were assigned to this agent will fall back to:
  - Site-level agent (if configured)
  - Organization-level default agent (if configured)
  - Cloud polling (if no fallback exists)

  ## Examples

      iex> delete_agent_token(id)
      {:ok, %AgentToken{}}

  """
  @spec delete_agent_token(Ecto.UUID.t()) :: {:ok, AgentToken.t()} | {:error, any()}
  def delete_agent_token(id) do
    Repo.transaction(fn ->
      agent_token = Repo.get!(AgentToken, id)

      # 1. Remove all direct device assignments
      Repo.delete_all(from(a in AgentAssignment, where: a.agent_token_id == ^id))

      # 2. Remove this agent from any sites using it as default
      Repo.update_all(from(s in Towerops.Sites.Site, where: s.agent_token_id == ^id), set: [agent_token_id: nil])

      # 3. Remove this agent from any organizations using it as default
      Repo.update_all(from(o in Towerops.Organizations.Organization, where: o.default_agent_token_id == ^id),
        set: [default_agent_token_id: nil]
      )

      # 4. Delete the agent token
      Repo.delete!(agent_token)
    end)
  end

  ## Assignment management

  @doc """
  Assigns device to an agent.

  Creates an agent assignment that links device to an agent token.
  Devices can only be assigned to one agent at a time due to unique constraint.

  ## Examples

      iex> assign_device_to_agent(agent_id, device_id)
      {:ok, %AgentAssignment{}}

      iex> assign_device_to_agent(agent_id, already_assigned_device_id)
      {:error, %Ecto.Changeset{}}

  """
  @spec assign_device_to_agent(Ecto.UUID.t(), Ecto.UUID.t()) ::
          {:ok, AgentAssignment.t()} | {:error, Ecto.Changeset.t()}
  def assign_device_to_agent(agent_token_id, device_id) do
    result =
      %AgentAssignment{}
      |> AgentAssignment.changeset(%{
        agent_token_id: agent_token_id,
        device_id: device_id
      })
      |> Repo.insert()

    case result do
      {:ok, assignment} ->
        _ = broadcast_assignment_change(agent_token_id, :device_assigned)
        {:ok, assignment}

      error ->
        error
    end
  end

  @doc """
  Removes the agent assignment for a piece of device.

  ## Examples

      iex> unassign_device(device_id)
      {1, nil}

  """
  @spec unassign_device(Ecto.UUID.t()) :: {non_neg_integer(), nil | [term()]}
  def unassign_device(device_id) do
    # Get the current assignment to know which agent to notify
    current_assignment = get_device_assignment(device_id)

    result =
      AgentAssignment
      |> where([a], a.device_id == ^device_id)
      |> Repo.delete_all()

    # Notify the agent that was previously assigned
    _ =
      if current_assignment do
        broadcast_assignment_change(current_assignment.agent_token_id, :device_unassigned)
      end

    result
  end

  @doc """
  Lists all devices assigned to an agent.

  Returns device records with preloaded site and SNMP device associations.

  ## Examples

      iex> list_agent_devices(agent_token_id)
      [%Device{site: %Site{}, snmp_device: %Device{sensors: [...], interfaces: [...]}}]

  """
  @spec list_agent_devices(Ecto.UUID.t()) :: [Device.t()]
  def list_agent_devices(agent_token_id) do
    Repo.all(
      from(e in Device,
        join: a in AgentAssignment,
        on: a.device_id == e.id,
        where: a.agent_token_id == ^agent_token_id and a.enabled == true,
        preload: [:site, snmp_device: [:sensors, :interfaces]]
      )
    )
  end

  @doc """
  Lists all devices that should be polled by an agent.

  This includes devices that is:
  1. Directly assigned to this agent
  2. At a site that has this agent as default
  3. In an organization that has this agent as default (and not overridden by site or device)
  4. This agent is the global default cloud poller (and device has no assignment at any level)

  Returns device records with preloaded associations, filtered to only SNMP-enabled devices.

  ## Examples

      iex> list_agent_polling_targets(agent_token_id)
      [%Device{snmp_enabled: true, site: %Site{}, snmp_device: %Device{sensors: [...], interfaces: [...]}}]

  """
  @spec list_agent_polling_targets(Ecto.UUID.t()) :: [Device.t()]
  def list_agent_polling_targets(agent_token_id) do
    global_default = Towerops.Settings.get_global_default_cloud_poller()
    is_global_default = agent_token_id == global_default

    query =
      from(e in Device,
        left_join: s in assoc(e, :site),
        join: o in assoc(e, :organization),
        left_join: aa in AgentAssignment,
        on: aa.device_id == e.id and aa.enabled == true,
        where: e.snmp_enabled == true,
        preload: [
          :agent_assignments,
          :organization,
          site: :organization,
          snmp_device: [:sensors, :interfaces]
        ]
      )

    query
    |> where_agent_matches(agent_token_id, is_global_default)
    |> Repo.all()
  end

  # Adds where clause to match devices assigned to this agent
  # Supports devices assigned directly to org (no site) or to a site
  # credo:disable-for-next-line Credo.Check.Refactor.CyclomaticComplexity
  defp where_agent_matches(query, agent_token_id, is_global_default) do
    from([e, s, o, aa] in query,
      where:
        aa.agent_token_id == ^agent_token_id or
          (is_nil(aa.agent_token_id) and not is_nil(s.id) and
             s.agent_token_id == ^agent_token_id) or
          (is_nil(aa.agent_token_id) and
             (is_nil(s.id) or is_nil(s.agent_token_id)) and
             o.default_agent_token_id == ^agent_token_id) or
          (^is_global_default and is_nil(aa.agent_token_id) and
             (is_nil(s.id) or is_nil(s.agent_token_id)) and
             is_nil(o.default_agent_token_id))
    )
  end

  @doc """
  Gets the agent assignment for a piece of device.

  Returns the assignment record if it exists, nil otherwise.

  ## Examples

      iex> get_device_assignment(device_id)
      %AgentAssignment{}

      iex> get_device_assignment(unassigned_device_id)
      nil

  """
  @spec get_device_assignment(Ecto.UUID.t()) :: AgentAssignment.t() | nil
  def get_device_assignment(device_id) do
    Repo.get_by(AgentAssignment, device_id: device_id)
  end

  @doc """
  Updates or creates an device assignment.

  If agent_token_id is nil, removes any existing assignment.
  Otherwise, creates or updates the assignment.

  ## Examples

      iex> update_device_assignment(device_id, agent_token_id)
      {:ok, %AgentAssignment{}}

      iex> update_device_assignment(device_id, nil)
      {:ok, nil}

  """
  @spec update_device_assignment(Ecto.UUID.t(), Ecto.UUID.t() | nil) ::
          {:ok, AgentAssignment.t() | nil} | {:error, Ecto.Changeset.t()}
  def update_device_assignment(device_id, nil) do
    # Remove assignment if setting to nil
    unassign_device(device_id)
    {:ok, nil}
  end

  def update_device_assignment(device_id, agent_token_id) do
    case get_device_assignment(device_id) do
      nil ->
        # Create new assignment
        assign_device_to_agent(agent_token_id, device_id)

      existing ->
        update_existing_assignment(existing, agent_token_id)
    end
  end

  defp update_existing_assignment(existing, agent_token_id) do
    old_agent_id = existing.agent_token_id

    result =
      existing
      |> AgentAssignment.changeset(%{agent_token_id: agent_token_id})
      |> Repo.update()

    case result do
      {:ok, assignment} ->
        # Notify both the old and new agents if they're different
        _ =
          if old_agent_id != agent_token_id do
            _ = broadcast_assignment_change(old_agent_id, :device_unassigned)
            broadcast_assignment_change(agent_token_id, :device_assigned)
          end

        {:ok, assignment}

      error ->
        error
    end
  end

  @doc """
  Gets the effective agent token for a piece of device.

  Resolves the agent token using the following hierarchy:
  1. Device's direct agent assignment (highest priority)
  2. Site's agent_token_id
  3. Organization's default_agent_token_id
  4. Global default cloud poller (application-wide fallback)
  5. Cloud polling (nil - lowest priority)

  Returns the agent_token_id if found, nil otherwise.

  The device must be preloaded with site: [organization: :default_agent_token].

  ## Examples

      iex> device = Repo.preload(device, site: [organization: :default_agent_token])
      iex> get_effective_agent_token(device)
      "agent-token-uuid"

      iex> get_effective_agent_token(equipment_without_any_assignment)
      nil

  """
  @spec get_effective_agent_token(Device.t()) :: Ecto.UUID.t() | nil
  def get_effective_agent_token(%Device{} = device) do
    # 1. Check device assignment
    case get_device_assignment(device.id) do
      %AgentAssignment{agent_token_id: agent_token_id} when not is_nil(agent_token_id) ->
        agent_token_id

      _ ->
        get_fallback_agent_token(device.site)
    end
  end

  # Checks site, organization, and global default in sequence
  defp get_fallback_agent_token(site) do
    case site do
      %{agent_token_id: agent_token_id} when not is_nil(agent_token_id) ->
        agent_token_id

      %{organization: %{default_agent_token_id: agent_token_id}}
      when not is_nil(agent_token_id) ->
        agent_token_id

      _ ->
        Towerops.Settings.get_global_default_cloud_poller()
    end
  end

  @doc """
  Gets information about where the effective agent token comes from.

  Returns a tuple of {agent_token_id, source} where source is one of:
  - :device - Direct device assignment
  - :site - Inherited from site
  - :organization - Inherited from organization
  - :global - Global default cloud poller
  - :none - No agent assigned

  ## Examples

      iex> get_effective_agent_token_with_source(device)
      {"token-id", :device}

      iex> get_effective_agent_token_with_source(equipment_no_assignment)
      {nil, :none}

  """
  @spec get_effective_agent_token_with_source(Device.t()) ::
          {Ecto.UUID.t() | nil, :device | :site | :organization | :global | :none}
  def get_effective_agent_token_with_source(%Device{} = device) do
    # 1. Check device assignment
    case get_device_assignment(device.id) do
      %AgentAssignment{agent_token_id: agent_token_id} when not is_nil(agent_token_id) ->
        {agent_token_id, :device}

      _ ->
        get_fallback_agent_token_with_source(device.site)
    end
  end

  # Checks site, organization, and global default in sequence with source tracking
  defp get_fallback_agent_token_with_source(site) do
    case site do
      %{agent_token_id: agent_token_id} when not is_nil(agent_token_id) ->
        {agent_token_id, :site}

      %{organization: %{default_agent_token_id: agent_token_id}}
      when not is_nil(agent_token_id) ->
        {agent_token_id, :organization}

      _ ->
        case Towerops.Settings.get_global_default_cloud_poller() do
          agent_token_id when not is_nil(agent_token_id) ->
            {agent_token_id, :global}

          _ ->
            {nil, :none}
        end
    end
  end

  @doc """
  Checks if a device should be polled by Phoenix directly.

  Returns false if the device is assigned to a non-cloud-poller agent
  (i.e., a customer/local agent). Phoenix should never poll devices that
  are assigned to local agents - those devices should only be polled by
  their assigned agent.

  Returns true if:
  - No agent is assigned to the device
  - The device is assigned to a cloud poller (Phoenix-hosted agent)

  ## Examples

      iex> should_phoenix_poll_device?(device_with_local_agent)
      false

      iex> should_phoenix_poll_device?(device_with_cloud_poller)
      true

      iex> should_phoenix_poll_device?(device_no_agent)
      true

  """
  @spec should_phoenix_poll_device?(Device.t()) :: boolean()
  def should_phoenix_poll_device?(%Device{} = device) do
    case get_effective_agent_token(device) do
      nil ->
        # No agent assigned - Phoenix can poll
        true

      agent_token_id ->
        # Check if the agent is a cloud poller
        agent_token = get_agent_token!(agent_token_id)
        agent_token.is_cloud_poller
    end
  end

  ## PubSub notifications

  @doc """
  Broadcasts an assignment change to any connected agent channels.

  This allows agents to immediately receive updated job lists when
  devices are assigned or unassigned, without waiting for reconnection.
  """
  @spec broadcast_assignment_change(Ecto.UUID.t() | nil, atom()) :: :ok
  def broadcast_assignment_change(agent_token_id, event) when not is_nil(agent_token_id) do
    Phoenix.PubSub.broadcast(
      Towerops.PubSub,
      "agent:#{agent_token_id}:assignments",
      {:assignments_changed, event}
    )
  end

  def broadcast_assignment_change(nil, _event), do: :ok
end