defmodule Til.TestHelpers do
  import ExUnit.Assertions

  alias Til.Parser
  # For type-checking related helpers
  alias Til.Typer
  alias Til.AstUtils

  # --- Node Access Helpers ---

  def get_file_node_from_map(nodes_map) do
    file_node =
      Enum.find(Map.values(nodes_map), fn node ->
        # Ensure the map is an AST node before checking its type
        Map.has_key?(node, :ast_node_type) && node.ast_node_type == :file
      end)

    refute is_nil(file_node), "File node not found in #{inspect(nodes_map)}"
    file_node
  end

  def get_node_by_id(nodes_map, node_id) do
    node = Map.get(nodes_map, node_id)
    refute is_nil(node), "Node with ID #{inspect(node_id)} not found in nodes_map."
    node
  end

  def get_first_child_node(nodes_map, parent_node_id \\ nil) do
    parent_node =
      if is_nil(parent_node_id) do
        get_file_node_from_map(nodes_map)
      else
        get_node_by_id(nodes_map, parent_node_id)
      end

    children_ids = Map.get(parent_node, :children, [])

    if Enum.empty?(children_ids) do
      flunk("Parent node #{parent_node.id} has no children, cannot get first child node.")
    end

    get_node_by_id(nodes_map, hd(children_ids))
  end

  def get_nth_child_node(nodes_map, index, parent_node_id \\ nil) do
    parent_node =
      if is_nil(parent_node_id) do
        get_file_node_from_map(nodes_map)
      else
        get_node_by_id(nodes_map, parent_node_id)
      end

    children_ids = Map.get(parent_node, :children, [])

    unless index >= 0 && index < length(children_ids) do
      flunk(
        "Child node at index #{index} not found for parent #{parent_node.id}. Parent has #{length(children_ids)} children. Children IDs: #{inspect(children_ids)}"
      )
    end

    get_node_by_id(nodes_map, Enum.at(children_ids, index))
  end

  # --- Combined Parse/TypeCheck and Get Node Helpers ---
  # These return {node, nodes_map}

  def parse_and_get_first_node(source_string, file_name \\ "test_source") do
    {:ok, parsed_nodes_map} = Parser.parse(source_string, file_name)
    {get_first_child_node(parsed_nodes_map), parsed_nodes_map}
  end

  def parse_and_get_nth_node(source_string, index, file_name \\ "test_source") do
    {:ok, parsed_nodes_map} = Parser.parse(source_string, file_name)
    {get_nth_child_node(parsed_nodes_map, index), parsed_nodes_map}
  end

  def typecheck_and_get_first_node(source_string, file_name \\ "test_source") do
    {:ok, parsed_nodes_map} = Parser.parse(source_string, file_name)
    {:ok, typed_nodes_map} = Typer.type_check(parsed_nodes_map)
    {get_first_child_node(typed_nodes_map), typed_nodes_map}
  end

  def typecheck_and_get_nth_node(source_string, index, file_name \\ "test_source") do
    {:ok, parsed_nodes_map} = Parser.parse(source_string, file_name)
    {:ok, typed_nodes_map} = Typer.type_check(parsed_nodes_map)
    {get_nth_child_node(typed_nodes_map, index), typed_nodes_map}
  end

  # --- Type Assertion Helpers ---

  # Strips :id fields recursively and resolves _id suffixed keys
  # from a type definition structure, using nodes_map for lookups.
  def deep_strip_id(type_definition, nodes_map) do
    cond do
      is_struct(type_definition, MapSet) ->
        type_definition
        |> MapSet.to_list()
        # Recursively call on elements
        |> Enum.map(&deep_strip_id(&1, nodes_map))
        |> MapSet.new()

      is_map(type_definition) ->
        # Remove its own :id.
        map_without_id = Map.delete(type_definition, :id)

        # Recursively process its values, resolving _id keys.
        # Special handling for type_annotation_mismatch error details
        # Check if this map is a type definition itself (has :type_kind) before specific error checks
        if Map.has_key?(type_definition, :type_kind) &&
             type_definition.type_kind == :error &&
             type_definition.reason == :type_annotation_mismatch do
          actual_type_clean =
            case Map.get(type_definition, :actual_type_id) do
              nil -> nil
              id -> deep_strip_id(Map.get(nodes_map, id), nodes_map)
            end

          expected_type_clean =
            case Map.get(type_definition, :expected_type_id) do
              nil -> nil
              id -> deep_strip_id(Map.get(nodes_map, id), nodes_map)
            end

          %{
            type_kind: :error,
            reason: :type_annotation_mismatch,
            actual_type: actual_type_clean,
            expected_type: expected_type_clean
          }
        else
          # Standard processing for other maps
          Enum.reduce(map_without_id, %{}, fn {original_key, original_value}, acc_map ->
            {final_key, final_value} =
              cond do
                # Handle :element_type_id -> :element_type
                original_key == :element_type_id && is_atom(original_value) ->
                  resolved_def = Map.get(nodes_map, original_value)
                  {:element_type, deep_strip_id(resolved_def, nodes_map)}

                # Handle :key_type_id -> :key_type
                original_key == :key_type_id && is_atom(original_value) ->
                  resolved_def = Map.get(nodes_map, original_value)
                  {:key_type, deep_strip_id(resolved_def, nodes_map)}

                # Handle :value_type_id -> :value_type (e.g., in index_signature or %{value_type_id: ..., optional: ...})
                original_key == :value_type_id && is_atom(original_value) ->
                  resolved_def = Map.get(nodes_map, original_value)
                  {:value_type, deep_strip_id(resolved_def, nodes_map)}

                is_map(original_value) ->
                  {original_key, deep_strip_id(original_value, nodes_map)}

                is_list(original_value) ->
                  {original_key, deep_strip_id(original_value, nodes_map)} # Handles lists of type defs

                true ->
                  {original_key, original_value}
              end

            Map.put(acc_map, final_key, final_value)
          end)
        end

      is_list(type_definition) ->
        # Recursively call on elements for lists of type definitions
        Enum.map(type_definition, &deep_strip_id(&1, nodes_map))

      true -> # Literals, atoms, numbers, nil, etc. (leaf nodes in the type structure)
        type_definition
    end
  end

  def assert_node_typed_as(node, typed_nodes_map, expected_type_definition_clean) do
    type_key = node.type_id

    assert is_atom(type_key),
           "Type ID for node #{inspect(node.id)} (#{node.raw_string}) should be an atom key, got: #{inspect(type_key)}"

    actual_type_definition_from_map = Map.get(typed_nodes_map, type_key)

    refute is_nil(actual_type_definition_from_map),
           "Type definition for key #{inspect(type_key)} (from node #{node.id}, raw: '#{node.raw_string}') not found in nodes_map."

    actual_type_definition_clean =
      deep_strip_id(actual_type_definition_from_map, typed_nodes_map)

    assert actual_type_definition_clean == expected_type_definition_clean,
           "Type mismatch for node #{node.id} ('#{node.raw_string}').\nExpected (clean):\n#{inspect(expected_type_definition_clean, pretty: true, limit: :infinity)}\nGot (clean):\n#{inspect(actual_type_definition_clean, pretty: true, limit: :infinity)}"
  end

  def inspect_nodes(node_map) do
    AstUtils.build_debug_ast_data(node_map)
    |> IO.inspect(pretty: true, limit: :infinity)
  end
end