defmodule Tilly.BDD do
  @moduledoc """
  Manages the BDD store, including hash-consing of BDD nodes.
  The BDD store is expected to be part of a `typing_ctx` map under the key `:bdd_store`.
  """

  alias Tilly.BDD.Node

  @false_node_id 0
  @true_node_id 1
  @initial_next_node_id 2
  @universal_ops_module :universal_ops

  @doc """
  Initializes the BDD store within the typing context.
  Pre-interns canonical `false` and `true` BDD nodes.
  """
  def init_bdd_store(typing_ctx) when is_map(typing_ctx) do
    false_structure = Node.mk_false()
    true_structure = Node.mk_true()

    bdd_store = %{
      nodes_by_structure: %{
        {false_structure, @universal_ops_module} => @false_node_id,
        {true_structure, @universal_ops_module} => @true_node_id
      },
      structures_by_id: %{
        @false_node_id => %{structure: false_structure, ops_module: @universal_ops_module},
        @true_node_id => %{structure: true_structure, ops_module: @universal_ops_module}
      },
      next_node_id: @initial_next_node_id,
      ops_cache: %{} # Cache for BDD operations {op_key, id1, id2} -> result_id
    }

    Map.put(typing_ctx, :bdd_store, bdd_store)
  end

  @doc """
  Gets an existing BDD node ID or interns a new one if it's not already in the store.

  Returns a tuple `{new_typing_ctx, node_id}`.
  The `typing_ctx` is updated if a new node is interned.
  """
  def get_or_intern_node(typing_ctx, logical_structure, ops_module_atom) do
    bdd_store = Map.get(typing_ctx, :bdd_store)

    unless bdd_store do
      raise ArgumentError, "BDD store not initialized in typing_ctx. Call init_bdd_store first."
    end

    key = {logical_structure, ops_module_atom}

    case Map.get(bdd_store.nodes_by_structure, key) do
      nil ->
        # Node not found, intern it
        node_id = bdd_store.next_node_id

        new_nodes_by_structure = Map.put(bdd_store.nodes_by_structure, key, node_id)

        node_data = %{structure: logical_structure, ops_module: ops_module_atom}
        new_structures_by_id = Map.put(bdd_store.structures_by_id, node_id, node_data)

        new_next_node_id = node_id + 1

        new_bdd_store =
          %{
            bdd_store
            | nodes_by_structure: new_nodes_by_structure,
              structures_by_id: new_structures_by_id,
              next_node_id: new_next_node_id
          }

        new_typing_ctx = Map.put(typing_ctx, :bdd_store, new_bdd_store)
        {new_typing_ctx, node_id}

      existing_node_id ->
        # Node found
        {typing_ctx, existing_node_id}
    end
  end

  @doc """
  Retrieves the node's structure and ops_module from the BDD store.
  Returns `%{structure: logical_structure_tuple, ops_module: ops_module_atom}` or `nil` if not found.
  """
  def get_node_data(typing_ctx, node_id) do
    with %{bdd_store: %{structures_by_id: structures_by_id}} <- typing_ctx,
         data when not is_nil(data) <- Map.get(structures_by_id, node_id) do
      data
    else
      _ -> nil
    end
  end

  @doc """
  Checks if the given node ID corresponds to the canonical `false` BDD node.
  """
  def is_false_node?(typing_ctx, node_id) do
    # Optimized check for the predefined ID
    if node_id == @false_node_id do
      true
    else
      # Fallback for cases where a node might be structurally false but not have the canonical ID.
      # This should ideally not happen with proper interning of Node.mk_false() via get_or_intern_node.
      case get_node_data(typing_ctx, node_id) do
        %{structure: structure, ops_module: @universal_ops_module} ->
          structure == Node.mk_false()
        _ ->
          false
      end
    end
  end

  @doc """
  Checks if the given node ID corresponds to the canonical `true` BDD node.
  """
  def is_true_node?(typing_ctx, node_id) do
    # Optimized check for the predefined ID
    if node_id == @true_node_id do
      true
    else
      # Fallback for cases where a node might be structurally true but not have the canonical ID.
      case get_node_data(typing_ctx, node_id) do
        %{structure: structure, ops_module: @universal_ops_module} ->
          structure == Node.mk_true()
        _ ->
          false
      end
    end
  end

  @doc """
  Returns the canonical ID for the `false` BDD node.
  """
  def false_node_id(), do: @false_node_id

  @doc """
  Returns the canonical ID for the `true` BDD node.
  """
  def true_node_id(), do: @true_node_id

  @doc """
  Returns the atom used as the `ops_module` for universal nodes like `true` and `false`.
  """
  def universal_ops_module(), do: @universal_ops_module
end