syntropy

defmodule Syntropy.PositionComputer do
  @moduledoc """
  Periodic lattice position refresh using the shared contract interval.
  """

  use GenServer

  alias Syntropy.{
    Agent,
    ClusterInfo,
    ClusterInventory,
    Contract,
    EventRecorder,
    LatticeMath,
    LatticeSupervisor
  }

  @type contribution_type :: :solo_task | :parallel_task | :synthesis_task
  @type connection_kind :: :thought | :synthesis
  @type ordering_signature :: {[String.t()], non_neg_integer(), non_neg_integer()}
  @type agent_ref :: ClusterInfo.agent_ref()

  @rpc_timeout 5_000
  @type state :: %{
          contributions: %{optional(String.t()) => [String.t()]},
          connection_config: %{
            required(String.t()) => float()
          },
          history_limit: pos_integer(),
          last_ordering_signature: ordering_signature() | nil,
          weights: %{required(String.t()) => float()}
        }

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

  @spec recompute_now(keyword()) :: :ok
  def recompute_now(opts \\ []) do
    GenServer.call(__MODULE__, {:recompute, opts})
  end

  @spec record_contribution(String.t() | [String.t()], contribution_type() | String.t()) :: :ok
  def record_contribution(agent_ids, contribution_type) do
    agent_ids
    |> List.wrap()
    |> Enum.map(&ClusterInfo.local_agent_ref/1)
    |> record_contribution_refs(contribution_type)
  end

  @spec record_connections([String.t()], connection_kind()) :: :ok
  def record_connections(agent_ids, kind) do
    agent_ids
    |> Enum.uniq()
    |> Enum.map(&ClusterInfo.local_agent_ref/1)
    |> record_connections_refs(kind)
  end

  @spec record_contribution_refs([agent_ref()], contribution_type() | String.t()) :: :ok
  def record_contribution_refs(agent_refs, contribution_type) do
    agent_refs
    |> Enum.map(&normalize_ref/1)
    |> Enum.uniq_by(& &1.runtime_id)
    |> Enum.group_by(&node_name_for_ref/1, & &1.id)
    |> Enum.each(fn
      {nil, _agent_ids} ->
        :ok

      {node_name, agent_ids} ->
        if node_name == ClusterInfo.node_name() do
          record_local_contribution(agent_ids, contribution_type)
        else
          :rpc.call(
            node_name_to_atom(node_name),
            __MODULE__,
            :record_local_contribution,
            [agent_ids, contribution_type],
            @rpc_timeout
          )
        end
    end)

    :ok
  end

  @spec record_local_contribution([String.t()], contribution_type() | String.t()) :: :ok
  def record_local_contribution(agent_ids, contribution_type) do
    GenServer.call(__MODULE__, {:record_contribution, agent_ids, contribution_type})
  end

  @spec record_connections_refs([agent_ref()], connection_kind()) :: :ok
  def record_connections_refs(agent_refs, kind) do
    agent_refs
    |> Enum.map(&normalize_ref/1)
    |> Enum.uniq_by(& &1.runtime_id)
    |> pairwise_refs()
    |> Enum.each(&record_pair_ref(&1, kind))

    :ok
  end

  @spec recompute_refs([agent_ref()], keyword()) :: :ok
  def recompute_refs(agent_refs, opts \\ []) do
    agent_refs
    |> Enum.map(&normalize_ref/1)
    |> Enum.uniq_by(&node_name_for_ref/1)
    |> Enum.each(fn
      %{node_name: nil} ->
        :ok

      %{node_name: node_name} ->
        if node_name == ClusterInfo.node_name() do
          recompute_now(opts)
        else
          :rpc.call(
            node_name_to_atom(node_name),
            __MODULE__,
            :recompute_now,
            [opts],
            @rpc_timeout
          )
        end
    end)

    :ok
  end

  @spec reset() :: :ok
  def reset do
    GenServer.call(__MODULE__, :reset)
  end

  @impl true
  def init(_opts) do
    scoring_contract = Contract.runtime_scoring_contract()
    topology_contract = Contract.runtime_topology_contract()
    schedule_recompute()

    {:ok,
     %{
       contributions: %{},
       connection_config: topology_contract["connections"],
       history_limit: get_in(scoring_contract, ["contributions", "history_limit"]),
       last_ordering_signature: nil,
       weights: get_in(scoring_contract, ["contributions", "weights"])
     }}
  end

  @impl true
  def handle_call({:recompute, opts}, _from, state) do
    next_state =
      recompute_and_record(
        Keyword.get(opts, :task_id),
        state,
        Keyword.get(opts, :apply_ema, true)
      )

    {:reply, :ok, next_state}
  end

  @impl true
  def handle_call({:record_contribution, agent_ids, contribution_type}, _from, state) do
    next_state =
      Enum.reduce(agent_ids, state, fn agent_id, acc ->
        update_in(acc, [:contributions], fn contributions ->
          history =
            contributions
            |> Map.get(agent_id, [])
            |> then(&[normalize_contribution_type(contribution_type) | &1])
            |> Enum.take(acc.history_limit)

          Map.put(contributions, agent_id, history)
        end)
      end)

    {:reply, :ok, next_state}
  end

  @impl true
  def handle_call({:record_connections, agent_ids, kind}, _from, state) do
    agent_ids
    |> pairwise()
    |> Enum.each(&record_pair(&1, kind, state.connection_config))

    {:reply, :ok, state}
  end

  @impl true
  def handle_call(:reset, _from, state) do
    {:reply, :ok, %{state | contributions: %{}, last_ordering_signature: nil}}
  end

  @impl true
  def handle_info(:recompute, state) do
    next_state = recompute_and_record(nil, state, true)
    schedule_recompute()
    {:noreply, next_state}
  end

  defp recompute_and_record(task_id, state, apply_ema) do
    agents_before_refresh = LatticeSupervisor.list_agents()

    if agents_before_refresh == [] do
      %{state | last_ordering_signature: nil}
    else
      Enum.each(agents_before_refresh, fn agent ->
        Agent.refresh(agent.id,
          synthesis_score: contribution_score(state, agent.id),
          apply_ema: apply_ema
        )
      end)

      agents = LatticeSupervisor.list_agents()
      comparable_pairs_count = length(LatticeMath.comparable_pairs(agents))
      incomparable_pairs_count = length(LatticeMath.incomparable_pairs(agents))
      ranked_agent_ids = Enum.map(agents, &Map.get(&1, :runtime_id, &1.id))

      signature = {ranked_agent_ids, comparable_pairs_count, incomparable_pairs_count}

      if should_record_ordering_event?(task_id, state.last_ordering_signature, signature) do
        EventRecorder.record("ordering_recomputed", %{
          task_id: task_id,
          ranked_agent_ids: ranked_agent_ids,
          comparable_pairs_count: comparable_pairs_count,
          incomparable_pairs_count: incomparable_pairs_count
        })
      end

      %{state | last_ordering_signature: signature}
    end
  end

  defp contribution_score(state, agent_id) do
    state.contributions
    |> Map.get(agent_id, [])
    |> Enum.reduce(0.0, fn contribution_type, acc ->
      acc + Map.fetch!(state.weights, contribution_type)
    end)
    |> min(1.0)
  end

  defp normalize_contribution_type(contribution_type) when is_atom(contribution_type) do
    Atom.to_string(contribution_type)
  end

  defp normalize_contribution_type(contribution_type) when is_binary(contribution_type) do
    contribution_type
  end

  defp record_pair([left_id, right_id], _kind, _connection_config) when left_id == right_id,
    do: :ok

  defp record_pair([left_id, right_id], kind, connection_config) do
    with {:ok, left} <- fetch_agent(left_id),
         {:ok, right} <- fetch_agent(right_id),
         false <- left.temporary or right.temporary do
      left_connections = normalize_connections(left.connections)
      right_connections = normalize_connections(right.connections)

      current_left = Map.get(left_connections, right_id, %{weight: 0.0, interactions: 0})
      current_right = Map.get(right_connections, left_id, %{weight: 0.0, interactions: 0})
      current_weight = max(current_left.weight, current_right.weight)
      current_interactions = max(current_left.interactions, current_right.interactions)
      target_weight = min(1.0, current_weight + delta_for(kind, connection_config))

      smoothed_weight =
        clamp(
          connection_config["ema_alpha"] * target_weight +
            (1.0 - connection_config["ema_alpha"]) * current_weight
        )

      updated_connection = %{
        weight: smoothed_weight,
        interactions: current_interactions + 1
      }

      Agent.update_connections(left_id, Map.put(left_connections, right_id, updated_connection))
      Agent.update_connections(right_id, Map.put(right_connections, left_id, updated_connection))
    else
      _ -> :ok
    end
  end

  defp pairwise(agent_ids) do
    for {left_id, index} <- Enum.with_index(agent_ids),
        right_id <- Enum.drop(agent_ids, index + 1) do
      [left_id, right_id]
    end
  end

  defp fetch_agent(agent_id) do
    {:ok, Agent.get_state(agent_id)}
  catch
    :exit, _reason -> {:error, :missing_agent}
  end

  defp record_pair_ref([left_ref, right_ref], _kind)
       when left_ref.runtime_id == right_ref.runtime_id,
       do: :ok

  defp record_pair_ref([left_ref, right_ref], kind) do
    connection_config = Contract.runtime_topology_contract()["connections"]

    with %{} = left <- ClusterInventory.fetch_agent(left_ref),
         %{} = right <- ClusterInventory.fetch_agent(right_ref),
         false <- left.temporary or right.temporary do
      left_connections = normalize_connections(left.connections)
      right_connections = normalize_connections(right.connections)
      left_target_id = connection_target_id(left, right)
      right_target_id = connection_target_id(right, left)

      current_left = Map.get(left_connections, left_target_id, %{weight: 0.0, interactions: 0})
      current_right = Map.get(right_connections, right_target_id, %{weight: 0.0, interactions: 0})
      current_weight = max(current_left.weight, current_right.weight)
      current_interactions = max(current_left.interactions, current_right.interactions)
      target_weight = min(1.0, current_weight + delta_for(kind, connection_config))

      smoothed_weight =
        clamp(
          connection_config["ema_alpha"] * target_weight +
            (1.0 - connection_config["ema_alpha"]) * current_weight
        )

      updated_connection = %{
        weight: smoothed_weight,
        interactions: current_interactions + 1
      }

      update_connections(left_ref, Map.put(left_connections, left_target_id, updated_connection))

      update_connections(
        right_ref,
        Map.put(right_connections, right_target_id, updated_connection)
      )
    else
      _ -> :ok
    end
  end

  defp pairwise_refs(agent_refs) do
    for {left_ref, index} <- Enum.with_index(agent_refs),
        right_ref <- Enum.drop(agent_refs, index + 1) do
      [left_ref, right_ref]
    end
  end

  defp update_connections(%{id: agent_id} = ref, connections) do
    case node_name_for_ref(ref) do
      nil ->
        :ok

      node_name ->
        if node_name == ClusterInfo.node_name() do
          Agent.update_connections(agent_id, connections)
        else
          :rpc.call(
            node_name_to_atom(node_name),
            Agent,
            :update_connections,
            [agent_id, connections],
            @rpc_timeout
          )
        end

        :ok
    end
  end

  defp normalize_ref(
         %{id: _id, runtime_id: _runtime_id, node_id: _node_id, node_name: _node_name} = ref
       ),
       do: ref

  defp normalize_ref(%{id: _id, runtime_id: _runtime_id, node_id: _node_id} = ref) do
    Map.put(ref, :node_name, node_name_for_ref(ref))
  end

  defp normalize_ref(%{id: local_id, node_id: owner_node_id}) do
    %{
      id: local_id,
      runtime_id: ClusterInfo.runtime_id(local_id, owner_node_id),
      node_id: owner_node_id,
      node_name: node_name_for_ref(%{id: local_id, node_id: owner_node_id})
    }
  end

  defp normalize_ref(runtime_id) when is_binary(runtime_id) do
    runtime_id
    |> ClusterInfo.parse_runtime_id()
    |> case do
      {:ok, ref} -> normalize_ref(ref)
      :error -> raise ArgumentError, "invalid runtime agent reference: #{inspect(runtime_id)}"
    end
  end

  defp node_name_for_ref(%{node_name: node_name}) when is_binary(node_name), do: node_name

  defp node_name_for_ref(%{node_id: node_id}) do
    if node_id == ClusterInfo.node_id() do
      ClusterInfo.node_name()
    else
      cluster_node_name(node_id)
    end
  end

  defp node_name_to_atom(node_name) do
    Enum.find([Node.self() | Node.list()], &(Atom.to_string(&1) == node_name)) ||
      String.to_atom(node_name)
  end

  defp connection_target_id(agent, other) do
    if Map.get(agent, :node_id, ClusterInfo.node_id()) ==
         Map.get(other, :node_id, ClusterInfo.node_id()) do
      other.id
    else
      other.runtime_id
    end
  end

  defp cluster_node_name(node_id) do
    ClusterInventory.cluster_export(include_local_active_tasks: false)
    |> Map.fetch!(:nodes)
    |> Enum.find_value(fn
      %{node_id: ^node_id, node_name: node_name} -> node_name
      _other -> nil
    end)
  end

  defp delta_for(:thought, connection_config), do: connection_config["thought_delta"]
  defp delta_for(:synthesis, connection_config), do: connection_config["synthesis_delta"]

  defp normalize_connections(connections) do
    Map.new(connections, fn
      {target_id, %{weight: weight} = value} ->
        {target_id,
         %{weight: clamp(weight), interactions: max(Map.get(value, :interactions, 0), 0)}}

      {target_id, %{"weight" => weight} = value} ->
        {target_id,
         %{weight: clamp(weight), interactions: value |> Map.get("interactions", 0) |> max(0)}}

      {target_id, weight} when is_number(weight) ->
        {target_id, %{weight: clamp(weight), interactions: 0}}
    end)
  end

  defp schedule_recompute do
    interval = Contract.constants().recompute_interval_seconds * 1_000
    Process.send_after(self(), :recompute, interval)
  end

  defp should_record_ordering_event?(task_id, _previous_signature, _signature)
       when is_binary(task_id),
       do: true

  defp should_record_ordering_event?(_task_id, previous_signature, signature),
    do: previous_signature != signature

  defp clamp(value), do: min(max(value, 0.0), 1.0)
end