defmodule Handoff.SimpleAllocator do
  @moduledoc """
  A simple implementation of the Allocator behavior that provides
  basic allocation strategies for distributing functions across nodes.
  """

  @behaviour Handoff.Allocator

  alias Handoff.Allocator.AllocationError
  alias Handoff.Function

  @doc """
  Allocate functions to nodes based on resource requirements and node capabilities.

  ## Parameters
  - `functions`: List of functions to allocate
  - `caps`: Map of node capabilities in the format %{node() => capabilities_map}

  ## Returns
  A map with function IDs as keys and node assignments as values.
  """
  @impl Handoff.Allocator
  def allocate(functions, caps) do
    # Sort nodes for consistent allocation order, prioritizing Node.self()
    all_nodes = caps |> Map.keys() |> Enum.sort()
    self_node = Node.self()

    nodes =
      if self_node in all_nodes do
        [self_node | List.delete(all_nodes, self_node)]
      else
        # If Node.self() is not in caps, log a warning or handle as per desired behavior.
        # For now, just use the sorted list of all_nodes.
        # IO.inspect("Warning: Node.self() (#{inspect(self_node)}) not found in capabilities map. Defaulting to standard sort.", label: "SimpleAllocator")
        all_nodes
      end

    # Initialize available resources for each node based on initial capabilities
    # This map will be updated as functions are assigned.
    initial_available_resources =
      Enum.reduce(nodes, %{}, fn node, acc ->
        # Assuming caps is %{node_pid => %{cpu: x, memory: y}}
        # If a node isn't in caps or has no resources defined, default to 0 for safety.
        # However, SimpleResourceTracker.register ensures nodes have some capabilities.
        # For this logic, we'll directly use what's in caps.
        # If caps could be sparse or incomplete, more robust fetching would be needed here.
        Map.put(acc, node, Map.get(caps, node, %{cpu: 0, memory: 0}))
      end)

    {collocated_functions, functions} =
      Enum.split_with(functions, fn
        %Function{node: {:collocated, _}} -> true
        _ -> false
      end)

    {functions, collocated_functions, []} =
      merge_collocated_costs(functions, collocated_functions)

    {pinned_functions, dynamic_functions} =
      Enum.split_with(functions, fn %Function{node: node} -> not is_nil(node) end)

    # Process pinned functions first
    {pinned_assignments, available_resources_after_pinning} =
      perform_pinned_allocation(pinned_functions, %{}, initial_available_resources)

    # Process remaining (dynamic) functions
    {dynamic_assignments_map, available_resources_after_dynamic_allocation, _dynamic_nodes} =
      Enum.reduce(
        # Use functions not already pinned
        dynamic_functions,
        # Start with results from pinned
        {pinned_assignments, available_resources_after_pinning, nodes},
        fn
          %Function{id: id, cost: cost},
          {current_assignments, available_resources, current_nodes_list} ->
            if is_nil(cost) || cost == %{} do
              assigned_node = List.first(current_nodes_list)

              {Map.put(current_assignments, id, assigned_node), available_resources,
               current_nodes_list}
            else
              find_node_assignment(
                id,
                current_assignments,
                available_resources,
                cost,
                current_nodes_list
              )
            end
        end
      )

    {final_assignments_map, _} =
      collocated_functions
      |> Enum.map(fn %Function{node: {:collocated, target_id}} = function ->
        target_node = Map.get(dynamic_assignments_map, target_id)
        %{function | node: target_node}
      end)
      |> perform_pinned_allocation(
        dynamic_assignments_map,
        available_resources_after_dynamic_allocation
      )

    final_assignments_map
  end

  defp perform_pinned_allocation(functions, assignments, available_resources) do
    Enum.reduce(
      functions,
      {assignments, available_resources},
      fn %Function{id: id, cost: cost, node: assigned_node}, {acc_assignments, acc_resources} ->
        # Assume assigned_node is valid and its resources are tracked.
        # If the node is not in caps (and thus not in acc_resources initially),
        # this might indicate an issue or a node without declared capacity.
        # For now, we proceed assuming it's a known node.
        # If cost is nil, treat as no resource requirement for subtraction.
        cost = cost || %{}
        node_current_resources = Map.get(acc_resources, assigned_node, %{cpu: 0, memory: 0})

        if not can_allocate?(node_current_resources, cost) do
          raise AllocationError,
                "Insufficient resources on node #{inspect(assigned_node)} for function #{inspect(id)}"
        end

        # For pinned functions, we assign them regardless of can_allocate? outcome,
        # as pinning implies a directive. Resources are subtracted.
        updated_node_res = subtract_resources(node_current_resources, cost)
        new_resources = Map.put(acc_resources, assigned_node, updated_node_res)
        new_assignments = Map.put(acc_assignments, id, assigned_node)
        {new_assignments, new_resources}
      end
    )
  end

  defp find_node_assignment(id, current_assignments, available_resources, cost, nodes) do
    found_node_assignment_tuple =
      Enum.find_value(nodes, fn node ->
        node_current_resources = Map.get(available_resources, node)

        if can_allocate?(node_current_resources, cost) do
          # Node found, update assignments and subtract resources for this node
          new_assignments = Map.put(current_assignments, id, node)
          updated_node_resources = subtract_resources(node_current_resources, cost)

          new_available_resources =
            Map.put(available_resources, node, updated_node_resources)

          # Move the chosen node to the front of the list
          updated_nodes_list = [node | List.delete(nodes, node)]

          # Return value for Enum.find_value: {assignments, resources, updated_nodes}
          {new_assignments, new_available_resources, updated_nodes_list}
        end
      end)

    if found_node_assignment_tuple do
      # Node was found and resources/nodes list were updated
      found_node_assignment_tuple
    else
      # If no node has resources, assign to first node (original fallback)
      # and don't alter available_resources for this function assignment.
      # Move the fallback node to the front of the list.
      assigned_node = List.first(nodes)

      {Map.put(current_assignments, id, assigned_node), available_resources, nodes}
    end
  end

  # Helper function to check if node_resources can satisfy function_cost
  # Assumes resources are maps like %{cpu: x, memory: y}
  defp can_allocate?(node_resources, function_cost)
       when is_map(node_resources) and is_map(function_cost) do
    keys = Enum.uniq(Map.keys(function_cost) ++ Map.keys(node_resources))

    Enum.all?(keys, fn key ->
      available = node_resources[key] || 0
      required = function_cost[key] || 0

      available >= required
    end)
  end

  # Helper function to subtract function_cost from node_resources
  defp subtract_resources(node_resources, function_cost)
       when is_map(function_cost) and is_map(node_resources) do
    keys = Enum.uniq(Map.keys(function_cost) ++ Map.keys(node_resources))

    Map.new(keys, fn key ->
      node_resource = node_resources[key] || 0
      function_cost = function_cost[key] || 0
      result = max(0, node_resource - function_cost)

      {key, result}
    end)
  end

  defp merge_collocated_costs([], collocated_functions) do
    {[], [], collocated_functions}
  end

  defp merge_collocated_costs(functions, []) do
    {functions, [], []}
  end

  defp merge_collocated_costs(functions, collocated_functions) do
    g = build_collocation_graph(functions, collocated_functions)
    collocated_by_target = extract_collocation_components(g)

    updated_functions = merge_costs_into_functions(functions, collocated_by_target)
    remapped_collocations = build_remapping_table(collocated_by_target)

    {remapped, to_collocate} =
      partition_collocated_functions(collocated_functions, remapped_collocations)

    {updated_functions, remapped, to_collocate}
  end

  defp build_collocation_graph(functions, collocated_functions) do
    g = :digraph.new()

    # Add root vertices for all functions
    for %{id: id} <- functions do
      :digraph.add_vertex(g, id, :root)
    end

    # Add collocated functions and their edges
    for %{id: id, node: {:collocated, target_id}, cost: cost} <- collocated_functions do
      add_vertex_if_missing(g, id, cost)
      add_vertex_if_missing(g, target_id, cost)
      :digraph.add_edge(g, id, target_id)
    end

    g
  end

  defp add_vertex_if_missing(graph, vertex_id, cost) do
    if !:digraph.vertex(graph, vertex_id) do
      :digraph.add_vertex(graph, vertex_id, cost)
    end
  end

  defp extract_collocation_components(graph) do
    for component <- :digraph_utils.components(graph), into: %{} do
      {[{root_id, :root}], collocations} =
        component
        |> Enum.map(&:digraph.vertex(graph, &1))
        |> Enum.split_with(fn {_, label} -> label == :root end)

      {root_id, collocations}
    end
  end

  defp merge_costs_into_functions(functions, collocated_by_target) do
    Enum.map(functions, fn %{id: id, cost: cost} = function ->
      case Map.get(collocated_by_target, id) do
        nil -> function
        collocations -> %{function | cost: calculate_merged_cost(cost, collocations)}
      end
    end)
  end

  defp calculate_merged_cost(base_cost, collocations) do
    Enum.reduce(collocations, base_cost || %{}, fn {_id, cost}, acc ->
      merge_cost_if_present(acc, cost)
    end)
  end

  defp merge_cost_if_present(acc, nil), do: acc

  defp merge_cost_if_present(acc, cost) do
    Map.merge(acc, cost, fn _key, v1, v2 -> v1 + v2 end)
  end

  defp build_remapping_table(collocated_by_target) do
    collocated_by_target
    |> Enum.flat_map(fn {id, collocations} ->
      Enum.map(collocations, fn {collocated_id, _cost} ->
        {collocated_id, id}
      end)
    end)
    |> Map.new()
  end

  defp partition_collocated_functions(collocated_functions, remapped_collocations) do
    Enum.reduce(collocated_functions, {[], []}, fn f, {remapped, to_collocate} ->
      case Map.get(remapped_collocations, f.id) do
        nil -> {remapped, [f | to_collocate]}
        target_id -> {[%{f | node: {:collocated, target_id}, cost: nil} | remapped], to_collocate}
      end
    end)
  end
end
