defmodule Til.AstUtils do
  @moduledoc """
  Utility functions for working with Til AST node maps.
  """

  @doc """
  Retrieves a node from the nodes_map by its ID.

  ## Examples

      iex> nodes = %{1 => %{id: 1, name: "node1"}, 2 => %{id: 2, name: "node2"}}
      iex> Til.AstUtils.get_node(nodes, 1)
      %{id: 1, name: "node1"}

      iex> nodes = %{1 => %{id: 1, name: "node1"}}
      iex> Til.AstUtils.get_node(nodes, 3)
      nil
  """
  def get_node(nodes_map, node_id) when is_map(nodes_map) and is_integer(node_id) do
    Map.get(nodes_map, node_id)
  end

  @doc """
  Retrieves the child nodes of a given parent node.
  The parent can be specified by its ID or as a node map.
  Assumes child IDs are stored in the parent's `:children` field.

  ## Examples

      iex> node1 = %{id: 1, children: [2, 3]}
      iex> node2 = %{id: 2, value: "child1"}
      iex> node3 = %{id: 3, value: "child2"}
      iex> nodes = %{1 => node1, 2 => node2, 3 => node3}
      iex> Til.AstUtils.get_child_nodes(nodes, 1)
      [%{id: 2, value: "child1"}, %{id: 3, value: "child2"}]

      iex> Til.AstUtils.get_child_nodes(nodes, node1)
      [%{id: 2, value: "child1"}, %{id: 3, value: "child2"}]

      iex> node4 = %{id: 4} # No children field
      iex> nodes_no_children = %{4 => node4}
      iex> Til.AstUtils.get_child_nodes(nodes_no_children, 4)
      []
  """
  def get_child_nodes(nodes_map, parent_node_id)
      when is_map(nodes_map) and is_integer(parent_node_id) do
    case get_node(nodes_map, parent_node_id) do
      nil -> []
      parent_node -> get_child_nodes_from_node(nodes_map, parent_node)
    end
  end

  def get_child_nodes(nodes_map, parent_node) when is_map(nodes_map) and is_map(parent_node) do
    get_child_nodes_from_node(nodes_map, parent_node)
  end

  defp get_child_nodes_from_node(nodes_map, parent_node) do
    parent_node
    |> Map.get(:children, [])
    |> Enum.map(&get_node(nodes_map, &1))
    # Filter out if a child_id doesn't resolve to a node
    |> Enum.reject(&is_nil(&1))
  end

  @doc """
  Retrieves the parent node of a given child node.
  The child can be specified by its ID or as a node map.
  Assumes parent ID is stored in the child's `:parent_id` field.

  ## Examples

      iex> parent = %{id: 1, name: "parent"}
      iex> child = %{id: 2, parent_id: 1, name: "child"}
      iex> nodes = %{1 => parent, 2 => child}
      iex> Til.AstUtils.get_parent_node(nodes, 2)
      %{id: 1, name: "parent"}

      iex> Til.AstUtils.get_parent_node(nodes, child)
      %{id: 1, name: "parent"}

      iex> root_node = %{id: 3, parent_id: nil}
      iex> nodes_with_root = %{3 => root_node}
      iex> Til.AstUtils.get_parent_node(nodes_with_root, 3)
      nil
  """
  def get_parent_node(nodes_map, child_node_id)
      when is_map(nodes_map) and is_integer(child_node_id) do
    case get_node(nodes_map, child_node_id) do
      nil -> nil
      child_node -> get_parent_node_from_node(nodes_map, child_node)
    end
  end

  def get_parent_node(nodes_map, child_node) when is_map(nodes_map) and is_map(child_node) do
    get_parent_node_from_node(nodes_map, child_node)
  end

  defp get_parent_node_from_node(nodes_map, child_node) do
    case Map.get(child_node, :parent_id) do
      nil -> nil
      parent_id -> get_node(nodes_map, parent_id)
    end
  end

  @doc """
  Generates a string representation of the AST for pretty printing.
  This is a basic implementation and can be expanded for more detail.
  """
  def pretty_print_ast(nodes_map) when is_map(nodes_map) do
    all_node_ids = Map.keys(nodes_map)

    root_nodes =
      nodes_map
      |> Map.values()
      |> Enum.filter(fn node ->
        parent_id = Map.get(node, :parent_id)
        is_nil(parent_id) or not Enum.member?(all_node_ids, parent_id)
      end)
      |> Enum.sort_by(fn node ->
        case Map.get(node, :location) do
          [start_offset | _] when is_integer(start_offset) -> start_offset
          # Fallback to id if location is not as expected or not present
          _ -> node.id
        end
      end)

    Enum.map_join(root_nodes, "\n", fn root_node ->
      do_pretty_print_node(nodes_map, root_node, 0)
    end)
  end

  defp do_pretty_print_node(nodes_map, node, indent_level) do
    prefix = String.duplicate("  ", indent_level)
    node_id = node.id
    ast_type = Map.get(node, :ast_node_type, :unknown)
    raw_string = Map.get(node, :raw_string, "") |> String.replace("\n", "\\n")

    details =
      case ast_type do
        :literal_integer -> "value: #{Map.get(node, :value)}"
        # Consider truncating
        :literal_string -> "value: \"#{Map.get(node, :value)}\""
        :symbol -> "name: #{Map.get(node, :name)}"
        _ -> ""
      end

    error_info =
      if parsing_error = Map.get(node, :parsing_error) do
        " (ERROR: #{parsing_error})"
      else
        ""
      end

    current_node_str =
      "#{prefix}Node #{node_id} [#{ast_type}] raw: \"#{raw_string}\" #{details}#{error_info}"

    children_str =
      node
      |> Map.get(:children, [])
      |> Enum.map(fn child_id ->
        case get_node(nodes_map, child_id) do
          # Should not happen in valid AST
          nil -> "#{prefix}  Child ID #{child_id} (not found)"
          child_node -> do_pretty_print_node(nodes_map, child_node, indent_level + 1)
        end
      end)
      |> Enum.join("\n")

    if String.trim(children_str) == "" do
      current_node_str
    else
      current_node_str <> "\n" <> children_str
    end
  end

  @doc """
  Generates a nested data structure representation of the AST for debugging.
  """
  def build_debug_ast_data(nodes_map) when is_map(nodes_map) do
    all_node_ids = Map.keys(nodes_map)

    root_nodes =
      nodes_map
      |> Map.values()
      |> Enum.filter(fn node ->
        parent_id = Map.get(node, :parent_id)
        is_nil(parent_id) or not Enum.member?(all_node_ids, parent_id)
      end)
      |> Enum.sort_by(fn node ->
        case Map.get(node, :location) do
          [start_offset | _] when is_integer(start_offset) -> start_offset
          _ -> node.id
        end
      end)

    Enum.map(root_nodes, fn root_node ->
      do_build_debug_node_data(nodes_map, root_node)
    end)
  end

  defp do_build_debug_node_data(nodes_map, node) do
    node_id = node.id
    ast_type = Map.get(node, :ast_node_type, :unknown)
    raw_string = Map.get(node, :raw_string, "")

    details =
      case ast_type do
        :literal_integer -> %{value: Map.get(node, :value)}
        :literal_string -> %{value: Map.get(node, :value)}
        :symbol -> %{name: Map.get(node, :name)}
        _ -> %{}
      end

    error_info = Map.get(node, :parsing_error)

    base_node_data = %{
      id: node_id,
      ast_node_type: ast_type,
      raw_string: raw_string,
      details: details
    }

    node_data_with_error =
      if error_info do
        Map.put(base_node_data, :parsing_error, error_info)
      else
        base_node_data
      end

    children_data =
      node
      |> Map.get(:children, [])
      |> Enum.map(fn child_id ->
        case get_node(nodes_map, child_id) do
          nil -> %{error: "Child ID #{child_id} (not found)"}
          child_node -> do_build_debug_node_data(nodes_map, child_node)
        end
      end)

    if Enum.empty?(children_data) do
      node_data_with_error
    else
      Map.put(node_data_with_error, :children, children_data)
    end
  end
end