defmodule Tdd.TypeSpec do
  @moduledoc """
  Defines the `TypeSpec` structure and functions for its manipulation.
  Normalization includes alpha-conversion, beta-reduction, and a final
  canonical renaming pass for bound variables.
  """

  # --- Core Types ---
  @type t ::
          :any
          | :none
          | :atom
          | :integer
          | :list
          | :tuple
          | {:literal, term()}
          | {:union, [t()]}
          | {:intersect, [t()]}
          | {:negation, t()}
          | {:tuple, [t()]}
          | {:cons, head :: t(), tail :: t()}
          | {:list_of, element :: t()}
          | {:integer_range, min :: integer() | :neg_inf, max :: integer() | :pos_inf}
          | {:type_var, name :: atom()}
          | {:mu, type_variable_name :: atom(), body :: t()}
          | {:type_lambda, param_names :: [atom()], body :: t()}
          | {:type_apply, constructor_spec :: t(), arg_specs :: [t()]}

  @doc """
  Converts a `TypeSpec` into its canonical (normalized) form.
  Performs structural normalization, alpha-conversion, beta-reduction,
  and a final canonical renaming pass for all bound variables.
  """
  @spec normalize(t()) :: t()
  def normalize(spec) do
    {intermediate_normalized, _counter_after_pass1} = normalize_pass1(spec, %{}, 0)

    {final_spec_before_subtype_redux, _mu_counter, _lambda_counter} =
      canonical_rename_pass(intermediate_normalized, %{}, 0, 0)

    apply_subtype_reduction(final_spec_before_subtype_redux)
  end

  # Final pass for subtype-based reductions on fully canonical specs
  defp apply_subtype_reduction(spec) do
    case spec do
      {:union, members} ->
        recursively_reduced_members = Enum.map(members, &apply_subtype_reduction/1)

        flattened_members =
          Enum.flat_map(recursively_reduced_members, fn
            {:union, sub_members} -> sub_members
            m -> [m]
          end)

        unique_no_none =
          flattened_members
          |> Enum.reject(&(&1 == :none))
          |> Enum.uniq()

        if Enum.member?(unique_no_none, :any) do
          :any
        else
          # Pass `true` for already_normalized flag to is_subtype?
          final_members =
            Enum.reject(unique_no_none, fn member_to_check ->
              Enum.any?(unique_no_none, fn other_member ->
                member_to_check != other_member and
                  is_subtype?(member_to_check, other_member, true)
              end)
            end)

          case Enum.sort(final_members) do
            [] -> :none
            [single] -> single
            list_members -> {:union, list_members}
          end
        end

      {:intersect, members} ->
        recursively_reduced_members = Enum.map(members, &apply_subtype_reduction/1)

        expanded_flattened_members =
          Enum.flat_map(recursively_reduced_members, fn
            {:intersect, sub_members} -> sub_members
            # get_supertypes expects normalized spec, and its output is also normalized
            # Pass flag
            m -> get_supertypes(m, true)
          end)

        unique_no_any =
          expanded_flattened_members
          |> Enum.reject(&(&1 == :any))
          |> Enum.uniq()

        if Enum.member?(unique_no_any, :none) do
          :none
        else
          # Pass `true` for already_normalized flag to is_subtype?
          final_members =
            Enum.reject(unique_no_any, fn member_to_check ->
              Enum.any?(unique_no_any, fn other_member ->
                member_to_check != other_member and
                  is_subtype?(other_member, member_to_check, true)
              end)
            end)

          case Enum.sort(final_members) do
            [] -> :any
            [single] -> single
            list_members -> {:intersect, list_members}
          end
        end

      {:negation, body} ->
        {:negation, apply_subtype_reduction(body)}

      {:tuple, elements} ->
        {:tuple, Enum.map(elements, &apply_subtype_reduction/1)}

      {:cons, head, tail} ->
        {:cons, apply_subtype_reduction(head), apply_subtype_reduction(tail)}

      {:mu, var_name, body} ->
        {:mu, var_name, apply_subtype_reduction(body)}

      {:type_lambda, params, body} ->
        {:type_lambda, params, apply_subtype_reduction(body)}

      {:type_apply, constructor, args} ->
        {:type_apply, apply_subtype_reduction(constructor),
         Enum.map(args, &apply_subtype_reduction/1)}

      atomic_or_literal ->
        atomic_or_literal
    end
  end

  # ------------------------------------------------------------------
  # Pass 1: Structural Normalization, Beta-Reduction, Initial Alpha-Conversion
  # Returns: {normalized_spec, next_counter}
  # ------------------------------------------------------------------
  defp normalize_pass1(spec, env, counter) do
    res_tuple =
      case spec do
        s when is_atom(s) and s in [:any, :none, :atom, :integer, :list, :tuple] ->
          {s, counter}

        {:literal, _val} = lit_spec ->
          {lit_spec, counter}

        {:type_var, name} ->
          {Map.get(env, name, spec), counter}

        {:negation, sub_spec} ->
          normalize_negation_pass1(sub_spec, env, counter)

        {:tuple, elements} ->
          {normalized_elements, next_counter_after_elements} =
            map_fold_counter_for_pass1(elements, env, counter, &normalize_pass1/3)

          {{:tuple, normalized_elements}, next_counter_after_elements}

        {:cons, head, tail} ->
          {normalized_head, counter_after_head} = normalize_pass1(head, env, counter)
          {normalized_tail, counter_after_tail} = normalize_pass1(tail, env, counter_after_head)
          {{:cons, normalized_head, normalized_tail}, counter_after_tail}

        {:integer_range, min, max} ->
          range_spec =
            if is_integer(min) and is_integer(max) and min > max do
              :none
            else
              {:integer_range, min, max}
            end

          {range_spec, counter}

        {:union, members} ->
          normalize_union_pass1(members, env, counter)

        {:intersect, members} ->
          normalize_intersection_pass1(members, env, counter)

        {:list_of, element_spec} ->
          # We transform `list_of(E)` into a `mu` expression.
          # This expression will then be normalized by a recursive call.
          # First, normalize the element's spec.
          {normalized_element, counter_after_element} =
            normalize_pass1(element_spec, env, counter)

          # Create a *temporary, non-canonical* name for the recursive variable.
          # The subsequent `normalize_pass1` call on the `mu` form will perform
          # the proper, canonical renaming.
          temp_rec_var = :"$list_of_rec_var"

          list_body =
            {:union,
             [
               {:literal, []},
               {:cons, normalized_element, {:type_var, temp_rec_var}}
             ]}

          # Now, normalize the full mu-expression. This is the crucial step.
          # It will handle alpha-conversion of `temp_rec_var` and normalization
          # of the body's components.
          normalize_pass1({:mu, temp_rec_var, list_body}, env, counter_after_element)

        {:mu, var_name, body} ->
          # This logic is correct. It creates a fresh canonical name and
          # adds it to the environment for normalizing the body.
          fresh_temp_name = fresh_var_name(:p1_m_var, counter)
          body_env = Map.put(env, var_name, {:type_var, fresh_temp_name})

          {normalized_body, next_counter_after_body} =
            normalize_pass1(body, body_env, counter + 1)

          {{:mu, fresh_temp_name, normalized_body}, next_counter_after_body}

        {:type_lambda, param_names, body} ->
          {reversed_fresh_temp_names, next_counter_after_params, body_env} =
            Enum.reduce(param_names, {[], counter, env}, fn param_name,
                                                            {acc_fresh_names, cnt, current_env} ->
              fresh_name = fresh_var_name(:p1_lambda_var, cnt)

              {[fresh_name | acc_fresh_names], cnt + 1,
               Map.put(current_env, param_name, {:type_var, fresh_name})}
            end)

          fresh_temp_param_names = Enum.reverse(reversed_fresh_temp_names)

          {normalized_body, final_counter} =
            normalize_pass1(body, body_env, next_counter_after_params)

          {{:type_lambda, fresh_temp_param_names, normalized_body}, final_counter}

        {:type_apply, constructor_spec, arg_specs} ->
          {normalized_constructor, counter_after_constructor} =
            normalize_pass1(constructor_spec, env, counter)

          {normalized_arg_specs, counter_after_args} =
            map_fold_counter_for_pass1(
              arg_specs,
              env,
              counter_after_constructor,
              &normalize_pass1/3
            )

          case normalized_constructor do
            {:type_lambda, pass1_formal_params, pass1_lambda_body} ->
              if length(pass1_formal_params) != length(normalized_arg_specs) do
                raise "TypeSpec.normalize_pass1: Arity mismatch in application. Expected #{length(pass1_formal_params)} args, got #{length(normalized_arg_specs)}. Lambda: #{inspect(normalized_constructor)}, Args: #{inspect(normalized_arg_specs)}"
              else
                substitution_map = Map.new(Enum.zip(pass1_formal_params, normalized_arg_specs))

                substituted_body =
                  substitute_vars_pass1(pass1_lambda_body, substitution_map, MapSet.new())

                normalize_pass1(substituted_body, env, counter_after_args)
              end

            _other_constructor ->
              {{:type_apply, normalized_constructor, normalized_arg_specs}, counter_after_args}
          end

        other_spec ->
          raise "TypeSpec.normalize_pass1: Unhandled spec form: #{inspect(other_spec)}"
      end

    res_tuple
  end

  defp map_fold_counter_for_pass1(list, env, initial_counter, fun) do
    Enum.map_reduce(list, initial_counter, fn item, acc_counter ->
      fun.(item, env, acc_counter)
    end)
  end

  defp substitute_vars_pass1(spec, substitutions, bound_in_scope) do
    case spec do
      {:type_var, name} ->
        if MapSet.member?(bound_in_scope, name) do
          spec
        else
          Map.get(substitutions, name, spec)
        end

      {:mu, var_name, body} ->
        newly_bound_scope = MapSet.put(bound_in_scope, var_name)
        active_substitutions = Map.delete(substitutions, var_name)
        {:mu, var_name, substitute_vars_pass1(body, active_substitutions, newly_bound_scope)}

      {:type_lambda, param_names, body} ->
        newly_bound_scope = Enum.reduce(param_names, bound_in_scope, &MapSet.put(&2, &1))
        active_substitutions = Enum.reduce(param_names, substitutions, &Map.delete(&2, &1))

        {:type_lambda, param_names,
         substitute_vars_pass1(body, active_substitutions, newly_bound_scope)}

      {:negation, sub} ->
        {:negation, substitute_vars_pass1(sub, substitutions, bound_in_scope)}

      {:tuple, elements} ->
        {:tuple, Enum.map(elements, &substitute_vars_pass1(&1, substitutions, bound_in_scope))}

      {:cons, h, t} ->
        {:cons, substitute_vars_pass1(h, substitutions, bound_in_scope),
         substitute_vars_pass1(t, substitutions, bound_in_scope)}

      {:list_of, e} ->
        {:list_of, substitute_vars_pass1(e, substitutions, bound_in_scope)}

      {:union, members} ->
        {:union, Enum.map(members, &substitute_vars_pass1(&1, substitutions, bound_in_scope))}

      {:intersect, members} ->
        {:intersect, Enum.map(members, &substitute_vars_pass1(&1, substitutions, bound_in_scope))}

      {:type_apply, con, args} ->
        new_con = substitute_vars_pass1(con, substitutions, bound_in_scope)
        new_args = Enum.map(args, &substitute_vars_pass1(&1, substitutions, bound_in_scope))
        {:type_apply, new_con, new_args}

      _atomic_or_simple_spec ->
        spec
    end
  end

  defp normalize_negation_pass1(sub_spec, env, counter) do
    {normalized_sub, next_counter} = normalize_pass1(sub_spec, env, counter)

    res_spec =
      case normalized_sub do
        {:negation, inner_spec} -> inner_spec
        :any -> :none
        :none -> :any
        _ -> {:negation, normalized_sub}
      end

    {res_spec, next_counter}
  end

  defp normalize_union_pass1(members, env, initial_counter) do
    {list_of_normalized_member_lists, final_counter_after_all_members} =
      Enum.map_reduce(members, initial_counter, fn member_spec, current_processing_counter ->
        {normalized_member_spec_term, counter_after_this_member_normalized} =
          normalize_pass1(member_spec, env, current_processing_counter)

        members_to_add_to_overall_list =
          case normalized_member_spec_term do
            {:union, sub_members} -> sub_members
            _ -> [normalized_member_spec_term]
          end

        {members_to_add_to_overall_list, counter_after_this_member_normalized}
      end)

    normalized_and_flattened = List.flatten(list_of_normalized_member_lists)

    unique_members =
      normalized_and_flattened
      |> Enum.reject(&(&1 == :none))
      |> Enum.uniq()

    if Enum.member?(unique_members, :any) do
      {:any, final_counter_after_all_members}
    else
      sorted_for_pass1 = Enum.sort(unique_members)

      resulting_spec =
        case sorted_for_pass1 do
          [] -> :none
          [single_member] -> single_member
          list_members -> {:union, list_members}
        end

      {resulting_spec, final_counter_after_all_members}
    end
  end

  defp normalize_intersection_pass1(members, env, initial_counter) do
    {list_of_member_groups, final_counter_after_all_members} =
      Enum.map_reduce(members, initial_counter, fn member_spec, current_processing_counter ->
        {normalized_member_spec_term, counter_after_this_member_normalized} =
          normalize_pass1(member_spec, env, current_processing_counter)

        expanded_members =
          case normalized_member_spec_term do
            {:intersect, sub_members} -> sub_members
            _ -> get_supertypes_pass1(normalized_member_spec_term)
          end

        {expanded_members, counter_after_this_member_normalized}
      end)

    normalized_and_flattened_with_supertypes = List.flatten(list_of_member_groups)

    unique_members =
      normalized_and_flattened_with_supertypes
      |> Enum.reject(&(&1 == :any))
      |> Enum.uniq()

    if Enum.member?(unique_members, :none) do
      {:none, final_counter_after_all_members}
    else
      sorted_for_pass1 = Enum.sort(unique_members)

      resulting_spec =
        case sorted_for_pass1 do
          [] -> :any
          [single_member] -> single_member
          list_members -> {:intersect, list_members}
        end

      {resulting_spec, final_counter_after_all_members}
    end
  end

  defp get_supertypes_pass1(spec) do
    supertypes =
      case spec do
        {:literal, val} when is_atom(val) -> [:atom]
        {:literal, val} when is_integer(val) -> [:integer]
        {:literal, val} when is_list(val) -> [:list]
        {:literal, val} when is_tuple(val) -> [:tuple]
        {:mu, _v, _body} -> []
        {:tuple, _} -> [:tuple]
        {:integer_range, _, _} -> [:integer]
        _ -> []
      end

    MapSet.to_list(MapSet.new([spec | supertypes]))
  end

  defp canonical_rename_pass(spec, env, mu_c, lambda_c) do
    case spec do
      {:mu, old_var_name, body} ->
        new_canonical_name = fresh_var_name(:m_var, mu_c)
        body_env = Map.put(env, old_var_name, {:type_var, new_canonical_name})

        {renamed_body, next_mu_c, next_lambda_c} =
          canonical_rename_pass(body, body_env, mu_c + 1, lambda_c)

        {{:mu, new_canonical_name, renamed_body}, next_mu_c, next_lambda_c}

      {:type_lambda, old_param_names, body} ->
        {reversed_new_param_names, next_lambda_c_after_params, body_env} =
          Enum.reduce(old_param_names, {[], lambda_c, env}, fn old_name,
                                                               {acc_new_names, current_lc,
                                                                current_env} ->
            fresh_canonical_name = fresh_var_name(:lambda_var, current_lc)

            {[fresh_canonical_name | acc_new_names], current_lc + 1,
             Map.put(current_env, old_name, {:type_var, fresh_canonical_name})}
          end)

        new_canonical_param_names = Enum.reverse(reversed_new_param_names)

        {renamed_body, final_mu_c, final_lambda_c} =
          canonical_rename_pass(body, body_env, mu_c, next_lambda_c_after_params)

        {{:type_lambda, new_canonical_param_names, renamed_body}, final_mu_c, final_lambda_c}

      {:type_var, name} ->
        {Map.get(env, name, spec), mu_c, lambda_c}

      {:negation, sub_spec} ->
        {renamed_sub, nmc, nlc} = canonical_rename_pass(sub_spec, env, mu_c, lambda_c)
        {{:negation, renamed_sub}, nmc, nlc}

      {:tuple, elements} ->
        {renamed_elements, next_mu_c, next_lambda_c} =
          map_foldl_counters_for_rename(elements, env, mu_c, lambda_c, &canonical_rename_pass/4)

        {{:tuple, renamed_elements}, next_mu_c, next_lambda_c}

      {:cons, head, tail} ->
        {renamed_head, mu_c_after_head, lambda_c_after_head} =
          canonical_rename_pass(head, env, mu_c, lambda_c)

        {renamed_tail, mu_c_after_tail, lambda_c_after_tail} =
          canonical_rename_pass(tail, env, mu_c_after_head, lambda_c_after_head)

        {{:cons, renamed_head, renamed_tail}, mu_c_after_tail, lambda_c_after_tail}

      {:union, members} ->
        sorted_members = Enum.sort(members)

        {renamed_members, next_mu_c, next_lambda_c} =
          map_foldl_counters_for_rename(
            sorted_members,
            env,
            mu_c,
            lambda_c,
            &canonical_rename_pass/4
          )

        {{:union, Enum.sort(renamed_members)}, next_mu_c, next_lambda_c}

      {:intersect, members} ->
        sorted_members = Enum.sort(members)

        {renamed_members, next_mu_c, next_lambda_c} =
          map_foldl_counters_for_rename(
            sorted_members,
            env,
            mu_c,
            lambda_c,
            &canonical_rename_pass/4
          )

        {{:intersect, Enum.sort(renamed_members)}, next_mu_c, next_lambda_c}

      {:type_apply, constructor_spec, arg_specs} ->
        {renamed_constructor, mu_c_after_con, lambda_c_after_con} =
          canonical_rename_pass(constructor_spec, env, mu_c, lambda_c)

        {renamed_args, mu_c_after_args, lambda_c_after_args} =
          map_foldl_counters_for_rename(
            arg_specs,
            env,
            mu_c_after_con,
            lambda_c_after_con,
            &canonical_rename_pass/4
          )

        {{:type_apply, renamed_constructor, renamed_args}, mu_c_after_args, lambda_c_after_args}

      s when is_atom(s) ->
        {s, mu_c, lambda_c}

      {:literal, _} = spec ->
        {spec, mu_c, lambda_c}

      {:integer_range, _, _} = spec ->
        {spec, mu_c, lambda_c}

      {:list_of, _} = spec ->
        raise "TypeSpec.canonical_rename_pass: Unexpected :list_of, should be :mu. Spec: #{inspect(spec)}"

      _other ->
        raise "TypeSpec.canonical_rename_pass: Unhandled spec form: #{inspect(spec)}"
    end
  end

  defp map_foldl_counters_for_rename(list, env, initial_mu_c, initial_lambda_c, fun) do
    {reversed_results, final_mu_c, final_lambda_c} =
      Enum.reduce(list, {[], initial_mu_c, initial_lambda_c}, fn item, {acc_items, mc, lc} ->
        {processed_item, next_mc, next_lc} = fun.(item, env, mc, lc)
        {[processed_item | acc_items], next_mc, next_lc}
      end)

    {Enum.reverse(reversed_results), final_mu_c, final_lambda_c}
  end

  defp fresh_var_name(prefix_atom, counter) do
    :"#{Atom.to_string(prefix_atom)}#{counter}"
  end

  # Public API
  @spec is_subtype?(t(), t()) :: boolean
  def is_subtype?(spec1, spec2), do: is_subtype?(spec1, spec2, false)

  # Internal helper with already_normalized flag
  @spec is_subtype?(t(), t(), boolean) :: boolean
  def is_subtype?(spec1, spec2, already_normalized) do
    cond do
      spec1 == spec2 ->
        true

      spec1 == :none ->
        true

      spec2 == :any ->
        true

      spec1 == :any and spec2 != :any ->
        false

      spec2 == :none and spec1 != :none ->
        false

      true ->
        {norm_s1, norm_s2} =
          if already_normalized do
            {spec1, spec2}
          else
            {normalize(spec1), normalize(spec2)}
          end

        if norm_s1 == norm_s2 do
          true
        else
          do_is_subtype_structural?(norm_s1, norm_s2, MapSet.new())
        end
    end
  end

  defp do_is_subtype_structural?(spec1, spec2, visited) do
    if MapSet.member?(visited, {spec1, spec2}) do
      true
    else
      cond do
        spec1 == :none ->
          true

        spec2 == :any ->
          true

        spec1 == :any and spec2 != :any ->
          false

        spec2 == :none and spec1 != :none ->
          false

        spec1 == spec2 ->
          true

        true ->
          new_visited = MapSet.put(visited, {spec1, spec2})

          case {spec1, spec2} do
            {{:union, members1}, _} ->
              Enum.all?(members1, &do_is_subtype_structural?(&1, spec2, new_visited))

            {_, {:union, members2}} ->
              Enum.any?(members2, &do_is_subtype_structural?(spec1, &1, new_visited))

            {{:intersect, members1}, _} ->
              Enum.any?(members1, &do_is_subtype_structural?(&1, spec2, new_visited))

            {_, {:intersect, members2}} ->
              Enum.all?(members2, &do_is_subtype_structural?(spec1, &1, new_visited))

            {s1, s2}
            when is_atom(s1) and is_atom(s2) and s1 not in [:any, :none] and
                   s2 not in [:any, :none] ->
              s1 == s2

            {{:literal, v1}, {:literal, v2}} ->
              v1 == v2

            {{:literal, val}, :atom} when is_atom(val) ->
              true

            {{:literal, val}, :integer} when is_integer(val) ->
              true

            {{:literal, val}, :list} when is_list(val) ->
              true

            {{:literal, val}, :tuple} when is_tuple(val) ->
              true

            {{:tuple, elems1}, {:tuple, elems2}} when length(elems1) == length(elems2) ->
              Enum.zip_with(elems1, elems2, &do_is_subtype_structural?(&1, &2, new_visited))
              |> Enum.all?()

            {{:tuple, _}, :tuple} ->
              true

            {{:integer_range, _, _}, :integer} ->
              true

            {{:integer_range, min1, max1}, {:integer_range, min2, max2}} ->
              min1_gte_min2 =
                case {min1, min2} do
                  {:neg_inf, _} -> min2 == :neg_inf
                  {_, :neg_inf} -> true
                  {m1_v, m2_v} when is_integer(m1_v) and is_integer(m2_v) -> m1_v >= m2_v
                  _ -> false
                end

              max1_lte_max2 =
                case {max1, max2} do
                  {:pos_inf, _} -> max2 == :pos_inf
                  {_, :pos_inf} -> true
                  {m1_v, m2_v} when is_integer(m1_v) and is_integer(m2_v) -> m1_v <= m2_v
                  _ -> false
                end

              min1_gte_min2 and max1_lte_max2

            {{:literal, val}, {:integer_range, min, max}} when is_integer(val) ->
              (min == :neg_inf or val >= min) and (max == :pos_inf or val <= max)

            {{:mu, v1, b1_body}, {:mu, v2, b2_body}} ->
              # This logic is from the original file, which is correct in principle
              # but was failing due to the TDD layer bug.
              cond do
                is_list_mu_form(b1_body, v1) and is_list_mu_form(b2_body, v2) ->
                  e1 = extract_list_mu_element(b1_body, v1)
                  e2 = extract_list_mu_element(b2_body, v2)
                  do_is_subtype_structural?(e1, e2, new_visited)

                true ->
                  unfolded_b1 = substitute_vars_canonical(b1_body, %{v1 => spec1})
                  do_is_subtype_structural?(unfolded_b1, spec2, new_visited)
              end

            {_non_mu_spec, {:mu, v2, b2_body} = mu_spec2} ->
              unfolded_b2 = substitute_vars_canonical(b2_body, %{v2 => mu_spec2})
              do_is_subtype_structural?(spec1, unfolded_b2, new_visited)

            {{:mu, v1, b1_body} = mu_spec1, _non_mu_spec} ->
              unfolded_b1 = substitute_vars_canonical(b1_body, %{v1 => mu_spec1})
              do_is_subtype_structural?(unfolded_b1, spec2, new_visited)

            {{:negation, n_body1}, {:negation, n_body2}} ->
              do_is_subtype_structural?(n_body2, n_body1, new_visited)

            _ ->
              false
          end
      end
    end
  end

  defp substitute_vars_canonical(spec, substitutions) do
    case spec do
      {:type_var, name} ->
        Map.get(substitutions, name, spec)

      {:mu, var_name, body} ->
        active_substitutions = Map.delete(substitutions, var_name)
        {:mu, var_name, substitute_vars_canonical(body, active_substitutions)}

      {:type_lambda, param_names, body} ->
        active_substitutions = Enum.reduce(param_names, substitutions, &Map.delete(&2, &1))
        {:type_lambda, param_names, substitute_vars_canonical(body, active_substitutions)}

      {:negation, sub} ->
        {:negation, substitute_vars_canonical(sub, substitutions)}

      {:tuple, elements} ->
        {:tuple, Enum.map(elements, &substitute_vars_canonical(&1, substitutions))}

      {:cons, h, t} ->
        {:cons, substitute_vars_canonical(h, substitutions),
         substitute_vars_canonical(t, substitutions)}

      {:list_of, e} ->
        {:list_of, substitute_vars_canonical(e, substitutions)}

      {:union, members} ->
        {:union, Enum.map(members, &substitute_vars_canonical(&1, substitutions))}

      {:intersect, members} ->
        {:intersect, Enum.map(members, &substitute_vars_canonical(&1, substitutions))}

      {:type_apply, con, args} ->
        new_con = substitute_vars_canonical(con, substitutions)
        new_args = Enum.map(args, &substitute_vars_canonical(&1, substitutions))
        {:type_apply, new_con, new_args}

      _atomic_or_simple_spec ->
        spec
    end
  end

  defp is_list_mu_form({:union, members}, rec_var_name) do
    sorted_members = Enum.sort(members)

    match?([{:literal, []}, {:cons, _elem, {:type_var, ^rec_var_name}}], sorted_members) or
      match?([{:cons, _elem, {:type_var, ^rec_var_name}}, {:literal, []}], sorted_members)
  end

  defp is_list_mu_form(_, _), do: false

  defp extract_list_mu_element({:union, members}, rec_var_name) do
    Enum.find_value(members, fn
      {:cons, elem_spec, {:type_var, ^rec_var_name}} -> elem_spec
      _ -> nil
    end) || :any
  end

  # Public API for get_supertypes
  def get_supertypes(spec), do: get_supertypes(spec, false)

  # Internal helper for get_supertypes
  defp get_supertypes(spec_input, already_normalized) do
    fully_normalized_spec = if already_normalized, do: spec_input, else: normalize(spec_input)

    supertypes =
      case fully_normalized_spec do
        {:literal, val} when is_atom(val) -> [:atom]
        {:literal, val} when is_integer(val) -> [:integer]
        {:literal, val} when is_list(val) -> [:list]
        {:literal, val} when is_tuple(val) -> [:tuple]
        {:mu, v, body} -> if is_list_mu_form(body, v), do: [:list], else: []
        {:tuple, _} -> [:tuple]
        {:integer_range, _, _} -> [:integer]
        _ -> []
      end

    MapSet.to_list(MapSet.new([fully_normalized_spec | supertypes]))
  end
end

defmodule Tdd.Store do
  # NOTE: This module remains unchanged.
  # The original provided code for this module is correct and complete.
  @moduledoc """
  Manages the state of the TDD system's node graph and operation cache.
  """

  # --- State Keys ---
  @nodes_key :tdd_nodes
  @node_by_id_key :tdd_node_by_id
  @next_id_key :tdd_next_id
  @op_cache_key :tdd_op_cache

  # --- Terminal Node IDs ---
  @false_node_id 0
  @true_node_id 1

  # --- Public API ---

  @doc "Initializes the TDD store in the current process."
  def init do
    Process.put(@nodes_key, %{})

    Process.put(@node_by_id_key, %{
      @false_node_id => :false_terminal,
      @true_node_id => :true_terminal
    })

    Process.put(@next_id_key, 2)
    Process.put(@op_cache_key, %{})
    :ok
  end

  @doc "Returns the ID for the TRUE terminal node (the 'any' type)."
  @spec true_node_id() :: non_neg_integer()
  def true_node_id, do: @true_node_id

  @doc "Returns the ID for the FALSE terminal node (the 'none' type)."
  @spec false_node_id() :: non_neg_integer()
  def false_node_id, do: @false_node_id

  @doc "Retrieves the details of a node by its ID."
  @spec get_node(non_neg_integer()) ::
          {:ok,
           {variable :: term(), yes_id :: non_neg_integer(), no_id :: non_neg_integer(),
            dc_id :: non_neg_integer()}}
          | {:ok, :true_terminal | :false_terminal}
          | {:error, :not_found}
  def get_node(id) do
    case Process.get(@node_by_id_key, %{}) do
      %{^id => details} -> {:ok, details}
      # IMPROVEMENT: Handle placeholder details gracefully
      %{^id => {:placeholder, _spec}} -> {:ok, {:placeholder, "Knot-tying placeholder"}}
      %{} -> {:error, :not_found}
    end
  end

  @doc """
  Finds an existing node that matches the structure or creates a new one.
  """
  @spec find_or_create_node(
          variable :: term(),
          yes_id :: non_neg_integer(),
          no_id :: non_neg_integer(),
          dc_id :: non_neg_integer()
        ) :: non_neg_integer()
  def find_or_create_node(variable, yes_id, no_id, dc_id) do
    if yes_id == no_id && yes_id == dc_id do
      yes_id
    else
      node_tuple = {variable, yes_id, no_id, dc_id}
      nodes = Process.get(@nodes_key, %{})

      case Map.get(nodes, node_tuple) do
        id when is_integer(id) ->
          id

        nil ->
          next_id = Process.get(@next_id_key)
          node_by_id = Process.get(@node_by_id_key)

          Process.put(@nodes_key, Map.put(nodes, node_tuple, next_id))
          Process.put(@node_by_id_key, Map.put(node_by_id, next_id, node_tuple))
          Process.put(@next_id_key, next_id + 1)

          next_id
      end
    end
  end

  @doc "Retrieves a result from the operation cache."
  @spec get_op_cache(term()) :: {:ok, term()} | :not_found
  def get_op_cache(cache_key) do
    case Process.get(@op_cache_key, %{}) do
      %{^cache_key => result} -> {:ok, result}
      %{} -> :not_found
    end
  end

  @doc "Puts a result into the operation cache."
  @spec put_op_cache(term(), term()) :: :ok
  def put_op_cache(cache_key, result) do
    cache = Process.get(@op_cache_key, %{})
    Process.put(@op_cache_key, Map.put(cache, cache_key, result))
    :ok
  end

  @doc """
  Creates a unique, temporary placeholder node for a recursive spec.
  Returns the ID of this placeholder.
  """
  @spec create_placeholder(TypeSpec.t()) :: non_neg_integer()
  def create_placeholder(spec) do
    # IMPROVEMENT: Correctly reserve an ID without creating a malformed node.
    next_id = Process.get(@next_id_key)
    node_by_id = Process.get(@node_by_id_key)

    # Associate the ID with a temporary detail for debugging and identification.
    # This does NOT create an entry in the main `nodes` (reverse-lookup) map.
    Process.put(@node_by_id_key, Map.put(node_by_id, next_id, {:placeholder, spec}))
    Process.put(@next_id_key, next_id + 1)

    next_id
  end

  @doc """
  Updates a node's details directly. Used for knot-tying.
  """
  @spec update_node_in_place(
          non_neg_integer(),
          new_details ::
            {:ok,
             {term(), non_neg_integer(), non_neg_integer(), non_neg_integer()}
             | :true_terminal
             | :false_terminal}
        ) :: :ok
  def update_node_in_place(id, {:ok, new_details}) do
    nodes = Process.get(@nodes_key)
    node_by_id = Process.get(@node_by_id_key)

    # This part is now safe. The old_details will be `{:placeholder, ...}`
    # which does not exist as a key in the `nodes` map, so Map.delete is a no-op.
    old_details = Map.get(node_by_id, id)
    nodes = Map.delete(nodes, old_details)

    nodes =
      case new_details do
        {_v, _y, _n, _d} -> Map.put(nodes, new_details, id)
        _ -> nodes
      end

    node_by_id = Map.put(node_by_id, id, new_details)

    Process.put(@nodes_key, nodes)
    Process.put(@node_by_id_key, node_by_id)
    :ok
  end
end

defmodule Tdd.Variable do
  @moduledoc """
  Defines the canonical structure for all Tdd predicate variables.
  REFAC: This module is unchanged, but its functions for recursive types will
  now be called with TDD IDs instead of TypeSpecs by the Tdd.Compiler.
  """

  # --- Category 0: Primary Type Discriminators ---
  @spec v_is_atom() :: term()
  def v_is_atom, do: {0, :is_atom, nil, nil}

  @spec v_is_integer() :: term()
  def v_is_integer, do: {0, :is_integer, nil, nil}

  @spec v_is_list() :: term()
  def v_is_list, do: {0, :is_list, nil, nil}

  @spec v_is_tuple() :: term()
  def v_is_tuple, do: {0, :is_tuple, nil, nil}

  # --- Category 1: Atom Properties ---
  @spec v_atom_eq(atom()) :: term()
  def v_atom_eq(atom_val) when is_atom(atom_val), do: {1, :value, atom_val, nil}

  # --- Category 2: Integer Properties ---
  @spec v_int_lt(integer()) :: term()
  def v_int_lt(n) when is_integer(n), do: {2, :alt, n, nil}

  @spec v_int_eq(integer()) :: term()
  def v_int_eq(n) when is_integer(n), do: {2, :beq, n, nil}

  @spec v_int_gt(integer()) :: term()
  def v_int_gt(n) when is_integer(n), do: {2, :cgt, n, nil}

  # --- Category 4: Tuple Properties ---
  @spec v_tuple_size_eq(non_neg_integer()) :: term()
  def v_tuple_size_eq(size) when is_integer(size) and size >= 0, do: {4, :a_size, size, nil}

  @doc "Applies a predicate to a tuple element. The predicate is now represented by its TDD ID."
  @spec v_tuple_elem_pred(non_neg_integer(), sub_problem_tdd_id :: non_neg_integer()) :: term()
  def v_tuple_elem_pred(index, sub_problem_tdd_id)
      when is_integer(index) and index >= 0 and is_integer(sub_problem_tdd_id) do
    # REFAC: The nested term is now a TDD ID, not a spec or a variable.
    {4, :b_element, index, sub_problem_tdd_id}
  end

  # --- Category 5: List Properties ---
  @doc "Predicate: The list is the empty list `[]`."
  @spec v_list_is_empty() :: term()
  def v_list_is_empty, do: {5, :b_is_empty, nil, nil}

  @doc "Applies a predicate to the head. The predicate is now represented by its TDD ID."
  @spec v_list_head_pred(sub_problem_tdd_id :: non_neg_integer()) :: term()
  def v_list_head_pred(sub_problem_tdd_id) when is_integer(sub_problem_tdd_id),
    do: {5, :c_head, sub_problem_tdd_id, nil}

  @doc "Applies a predicate to the tail. The predicate is now represented by its TDD ID."
  @spec v_list_tail_pred(sub_problem_tdd_id :: non_neg_integer()) :: term()
  def v_list_tail_pred(sub_problem_tdd_id) when is_integer(sub_problem_tdd_id),
    do: {5, :d_tail, sub_problem_tdd_id, nil}
end

defmodule Tdd.Predicate.Info do
  # NOTE: This module remains largely unchanged. The traits for recursive variables
  # correctly identify them by structure, independent of what's inside.
  @moduledoc "A knowledge base for the properties of TDD predicate variables."
  alias Tdd.Variable

  @doc "Returns a map of traits for a given predicate variable."
  @spec get_traits(term()) :: map() | nil

  def get_traits({0, :is_atom, _, _}) do
    %{
      type: :primary,
      category: :atom,
      implies: [
        {Variable.v_is_integer(), false},
        {Variable.v_is_list(), false},
        {Variable.v_is_tuple(), false}
      ]
    }
  end

  def get_traits({0, :is_integer, _, _}) do
    %{
      type: :primary,
      category: :integer,
      implies: [
        {Variable.v_is_atom(), false},
        {Variable.v_is_list(), false},
        {Variable.v_is_tuple(), false}
      ]
    }
  end

  def get_traits({0, :is_list, _, _}) do
    %{
      type: :primary,
      category: :list,
      implies: [
        {Variable.v_is_atom(), false},
        {Variable.v_is_integer(), false},
        {Variable.v_is_tuple(), false}
      ]
    }
  end

  def get_traits({0, :is_tuple, _, _}) do
    %{
      type: :primary,
      category: :tuple,
      implies: [
        {Variable.v_is_atom(), false},
        {Variable.v_is_integer(), false},
        {Variable.v_is_list(), false}
      ]
    }
  end

  def get_traits({1, :value, _val, _}) do
    %{type: :atom_value, category: :atom, implies: [{Variable.v_is_atom(), true}]}
  end

  def get_traits({2, :alt, _, _}),
    do: %{type: :integer_prop, category: :integer, implies: [{Variable.v_is_integer(), true}]}

  def get_traits({2, :beq, _, _}),
    do: %{type: :integer_prop, category: :integer, implies: [{Variable.v_is_integer(), true}]}

  def get_traits({2, :cgt, _, _}),
    do: %{type: :integer_prop, category: :integer, implies: [{Variable.v_is_integer(), true}]}

  def get_traits({4, :a_size, _, _}) do
    %{type: :tuple_prop, category: :tuple, implies: [{Variable.v_is_tuple(), true}]}
  end

  # REFAC: The trait recognizes the structure. The content `_tdd_id` is opaque here.
  def get_traits({4, :b_element, index, _tdd_id}) do
    %{
      type: :tuple_recursive,
      category: :tuple,
      sub_key: {:elem, index},
      implies: [{Variable.v_is_tuple(), true}]
    }
  end

  def get_traits({5, :b_is_empty, _, _}) do
    %{type: :list_prop, category: :list, implies: [{Variable.v_is_list(), true}]}
  end

  # REFAC: The trait recognizes the structure. The content `_tdd_id` is opaque here.
  def get_traits({5, :c_head, _tdd_id, _}) do
    %{
      type: :list_recursive,
      category: :list,
      sub_key: :head,
      implies: [{Variable.v_is_list(), true}, {Variable.v_list_is_empty(), false}]
    }
  end

  # REFAC: The trait recognizes the structure. The content `_tdd_id` is opaque here.
  def get_traits({5, :d_tail, _tdd_id, _}) do
    %{
      type: :list_recursive,
      category: :list,
      sub_key: :tail,
      implies: [{Variable.v_is_list(), true}, {Variable.v_list_is_empty(), false}]
    }
  end

  def get_traits(_), do: nil
end

defmodule Tdd.Consistency.Engine do
  @moduledoc """
  A rule-based engine for checking the semantic consistency of a set of assumptions.
  REFAC: This module is largely unchanged, but we now make `remap_sub_problem_vars`
  and `unwrap_var` public so they can be shared with Tdd.Algo.
  """
  alias Tdd.Predicate.Info
  alias Tdd.Variable

  @doc "Checks if a map of assumptions is logically consistent."
  @spec check(map()) :: :consistent | :contradiction
  def check(assumptions), do: do_check(assumptions)

  @doc "Expands a map of assumptions with all their logical implications."
  @spec expand(map()) :: {:ok, map()} | {:error, :contradiction}
  def expand(assumptions), do: expand_with_implications(assumptions)

  # --- The Core Recursive Checker ---
  defp do_check(assumptions) do
    with {:ok, expanded} <- expand_with_implications(assumptions),
         :ok <- check_flat_consistency(expanded) do
      sub_problems =
        expanded
        |> Enum.group_by(fn {var, _val} -> (Info.get_traits(var) || %{})[:sub_key] end)
        |> Map.drop([nil])

      if map_size(sub_problems) == 0 do
        :consistent
      else
        Enum.find_value(sub_problems, :consistent, fn {_sub_key, sub_assumptions_list} ->
          remapped_assumptions = remap_sub_problem_vars(sub_assumptions_list)

          case do_check(remapped_assumptions) do
            :consistent -> nil
            :contradiction -> :contradiction
          end
        end)
      end
    else
      {:error, _reason} -> :contradiction
    end
  end

  # --- Recursive Checking Helpers ---

  @doc "Converts a list of scoped assumptions into a map of base assumptions for a sub-problem."
  @spec remap_sub_problem_vars([{term(), boolean()}]) :: map()
  def remap_sub_problem_vars(assumptions_list) do
    Map.new(assumptions_list, fn {var, val} ->
      {unwrap_var(var), val}
    end)
  end

  @doc "Extracts the inner content from a recursive variable."
  @spec unwrap_var(term()) :: term()
  def unwrap_var(var) do
    case var do
      # REFAC: These variables now contain TDD IDs, but this function just extracts
      # whatever is inside. The consumer (`handle_recursive_subproblem`) will know it's an ID.
      {4, :b_element, _index, inner_content} -> inner_content
      {5, :c_head, inner_content, _} -> inner_content
      {5, :d_tail, inner_content, _} -> inner_content
      other -> other
    end
  end

  # --- Implication Expansion (Unchanged) ---
  defp expand_with_implications(assumptions) do
    expand_loop(assumptions, assumptions)
  end

  defp expand_loop(new_assumptions, all_assumptions) do
    implications =
      Enum.flat_map(new_assumptions, fn
        {var, true} -> Map.get(Info.get_traits(var) || %{}, :implies, [])
        _ -> []
      end)

    case Enum.reduce(implications, {:ok, %{}}, fn {implied_var, implied_val}, acc ->
           reduce_implication({implied_var, implied_val}, all_assumptions, acc)
         end) do
      {:error, :contradiction} = err ->
        err

      {:ok, newly_added} when map_size(newly_added) == 0 ->
        {:ok, all_assumptions}

      {:ok, newly_added} ->
        expand_loop(newly_added, Map.merge(all_assumptions, newly_added))
    end
  end

  defp reduce_implication({var, val}, all_assumptions, {:ok, new_acc}) do
    case Map.get(all_assumptions, var) do
      nil -> {:ok, Map.put(new_acc, var, val)}
      ^val -> {:ok, new_acc}
      _other_val -> {:error, :contradiction}
    end
  end

  defp reduce_implication(_implication, _all_assumptions, error_acc), do: error_acc

  # --- Flat Consistency Checks (Unchanged) ---
  defp check_flat_consistency(assumptions) do
    with :ok <- check_primary_type_exclusivity(assumptions),
         :ok <- check_atom_consistency(assumptions),
         :ok <- check_list_consistency(assumptions),
         :ok <- check_integer_consistency(assumptions),
         :ok <- check_tuple_consistency(assumptions) do
      :ok
    else
      :error -> {:error, :consistency_error}
    end
  end

  defp check_primary_type_exclusivity(assumptions) do
    primary_types = [
      Variable.v_is_atom(),
      Variable.v_is_integer(),
      Variable.v_is_list(),
      Variable.v_is_tuple()
    ]

    true_primary_types = Enum.count(primary_types, &(Map.get(assumptions, &1) == true))

    if true_primary_types > 1, do: :error, else: :ok
  end

  defp check_atom_consistency(assumptions) do
    true_atom_values =
      Enum.reduce(assumptions, MapSet.new(), fn
        {{1, :value, atom_val, _}, true}, acc -> MapSet.put(acc, atom_val)
        _, acc -> acc
      end)

    if MapSet.size(true_atom_values) > 1, do: :error, else: :ok
  end

  defp check_tuple_consistency(assumptions) do
    true_tuple_sizes =
      Enum.reduce(assumptions, MapSet.new(), fn
        {{4, :a_size, size, _}, true}, acc -> MapSet.put(acc, size)
        _, acc -> acc
      end)

    if MapSet.size(true_tuple_sizes) > 1, do: :error, else: :ok
  end

  defp check_list_consistency(assumptions) do
    is_empty = Map.get(assumptions, Variable.v_list_is_empty()) == true
    has_head_prop = Enum.any?(assumptions, &match?({{5, :c_head, _, _}, true}, &1))
    has_tail_prop = Enum.any?(assumptions, &match?({{5, :d_tail, _, _}, true}, &1))

    if is_empty and (has_head_prop or has_tail_prop), do: :error, else: :ok
  end

  defp check_integer_consistency(assumptions) do
    initial_range = {:neg_inf, :pos_inf}

    result =
      Enum.reduce_while(assumptions, initial_range, fn assumption, {min, max} ->
        case assumption do
          {{2, :alt, n, _}, true} -> narrow_range(min, safe_min(max, n - 1))
          {{2, :alt, n, _}, false} -> narrow_range(safe_max(min, n), max)
          {{2, :beq, n, _}, true} -> narrow_range(safe_max(min, n), safe_min(max, n))
          {{2, :beq, n, _}, false} when min == n and max == n -> {:halt, :invalid}
          {{2, :cgt, n, _}, true} -> narrow_range(safe_max(min, n + 1), max)
          {{2, :cgt, n, _}, false} -> narrow_range(min, safe_min(max, n))
          _ -> {:cont, {min, max}}
        end
      end)

    case result,
      do: (
        :invalid -> :error
        _ -> :ok
      )
  end

  defp narrow_range(min, max) do
    is_invalid =
      case {min, max} do
        {:neg_inf, _} -> false
        {_, :pos_inf} -> false
        {m, n} when is_integer(m) and is_integer(n) -> m > n
        _ -> false
      end

    if is_invalid, do: {:halt, :invalid}, else: {:cont, {min, max}}
  end

  defp safe_max(:neg_inf, x), do: x
  defp safe_max(x, :neg_inf), do: x
  defp safe_max(:pos_inf, _), do: :pos_inf
  defp safe_max(_, :pos_inf), do: :pos_inf
  defp safe_max(a, b), do: :erlang.max(a, b)

  defp safe_min(:pos_inf, x), do: x
  defp safe_min(x, :pos_inf), do: x
  defp safe_min(:neg_inf, _), do: :neg_inf
  defp safe_min(_, :neg_inf), do: :neg_inf
  defp safe_min(a, b), do: :erlang.min(a, b)
end

defmodule Tdd.Algo do
  @moduledoc """
  Implements the core, stateless algorithms for TDD manipulation.
  """
  use Tdd.Debug
  alias Tdd.Store
  alias Tdd.Consistency.Engine
  alias Tdd.Debug

  # --- Binary Operation: Apply ---
  # This function is correct and does not need to be changed.
  @spec apply(atom, (atom, atom -> atom), non_neg_integer, non_neg_integer) :: non_neg_integer
  def apply(op_name, op_lambda, u1_id, u2_id) do
    cache_key = {:apply, op_name, Enum.sort([u1_id, u2_id])}

    case Store.get_op_cache(cache_key) do
      {:ok, result_id} ->
        result_id

      :not_found ->
        result_id = do_apply(op_name, op_lambda, u1_id, u2_id)
        Store.put_op_cache(cache_key, result_id)
        result_id
    end
    |> simplify()
  end

  # This function is correct and does not need to be changed.
  defp do_apply(op_name, op_lambda, u1_id, u2_id) do
    with {:ok, u1_details} <- Store.get_node(u1_id),
         {:ok, u2_details} <- Store.get_node(u2_id) do
      cond do
        (u1_details == :true_terminal or u1_details == :false_terminal) and
            (u2_details == :true_terminal or u2_details == :false_terminal) ->
          if op_lambda.(u1_details, u2_details) == :true_terminal,
            do: Store.true_node_id(),
            else: Store.false_node_id()

        # FIX: Add cases to handle placeholder nodes co-inductively.
        # Treat the placeholder as `any`.
        match?({:ok, {:placeholder, _}}, {:ok, u1_details}) ->
          # op(placeholder, u2) -> op(any, u2)
          op_lambda.(:true_terminal, u2_details)
          |> case do
            :true_terminal -> Store.true_node_id()
            :false_terminal -> Store.false_node_id()
            # This happens if op(any, u2) = u2 (e.g., intersection)
            _other -> u2_id
          end

        match?({:ok, {:placeholder, _}}, {:ok, u2_details}) ->
          # op(u1, placeholder) -> op(u1, any)
          op_lambda.(u1_details, :true_terminal)
          |> case do
            :true_terminal -> Store.true_node_id()
            :false_terminal -> Store.false_node_id()
            # This happens if op(u1, any) = u1 (e.g., intersection)
            _other -> u1_id
          end

        u1_details == :true_terminal or u1_details == :false_terminal ->
          {var2, y2, n2, d2} = u2_details

          Store.find_or_create_node(
            var2,
            apply(op_name, op_lambda, u1_id, y2),
            apply(op_name, op_lambda, u1_id, n2),
            apply(op_name, op_lambda, u1_id, d2)
          )

        u2_details == :true_terminal or u2_details == :false_terminal ->
          {var1, y1, n1, d1} = u1_details

          Store.find_or_create_node(
            var1,
            apply(op_name, op_lambda, y1, u2_id),
            apply(op_name, op_lambda, n1, u2_id),
            apply(op_name, op_lambda, d1, u2_id)
          )

        true ->
          {var1, y1, n1, d1} = u1_details
          {var2, y2, n2, d2} = u2_details
          top_var = Enum.min([var1, var2])

          res_y =
            apply(
              op_name,
              op_lambda,
              if(var1 == top_var, do: y1, else: u1_id),
              if(var2 == top_var, do: y2, else: u2_id)
            )

          res_n =
            apply(
              op_name,
              op_lambda,
              if(var1 == top_var, do: n1, else: u1_id),
              if(var2 == top_var, do: n2, else: u2_id)
            )

          res_d =
            apply(
              op_name,
              op_lambda,
              if(var1 == top_var, do: d1, else: u1_id),
              if(var2 == top_var, do: d2, else: u2_id)
            )

          Store.find_or_create_node(top_var, res_y, res_n, res_d)
      end
    end
  end

  # --- Unary Operation: Negation ---
  # This function is correct and does not need to be changed.
  @spec negate(non_neg_integer) :: non_neg_integer
  def negate(tdd_id) do
    cache_key = {:negate, tdd_id}

    case Store.get_op_cache(cache_key) do
      {:ok, result_id} ->
        result_id

      :not_found ->
        result_id =
          case Store.get_node(tdd_id) do
            {:ok, :true_terminal} ->
              Store.false_node_id()

            {:ok, :false_terminal} ->
              Store.true_node_id()

            {:ok, {var, y, n, d}} ->
              Store.find_or_create_node(var, negate(y), negate(n), negate(d))

            # FIX: Handle placeholder nodes encountered during co-inductive simplification.
            # A placeholder is treated as `any` for the check, so its negation is `none`.
            {:ok, {:placeholder, _spec}} ->
              Store.false_node_id()
          end

        Store.put_op_cache(cache_key, result_id)
        result_id
    end
  end

  # --- Unary Operation: Semantic Simplification ---
  # This function is correct and does not need to be changed.
  @spec simplify(non_neg_integer(), map()) :: non_neg_integer
  def simplify(tdd_id, assumptions \\ %{}) do
    sorted_assumptions = Enum.sort(assumptions)
    cache_key = {:simplify, tdd_id, sorted_assumptions}
    require Logger

    case Store.get_op_cache(cache_key) do
      {:ok, result_id} ->
        Logger.info("result #{inspect([result_id])}")

        if result_id == 1 do
          Logger.info("ASD #{inspect([tdd_id, assumptions])}")
        end

        result_id

      :not_found ->
        result_id = do_simplify(tdd_id, sorted_assumptions, MapSet.new())
        Store.put_op_cache(cache_key, result_id)
        result_id
    end
  end

  # This function is correct and does not need to be changed.
  defp do_simplify(tdd_id, sorted_assumptions, context) do
    current_state = {tdd_id, sorted_assumptions}

    if MapSet.member?(context, current_state) do
      Store.true_node_id()
    else
      new_context = MapSet.put(context, current_state)
      assumptions = Map.new(sorted_assumptions)

      if Engine.check(assumptions) == :contradiction do
        Store.false_node_id()
      else
        case Store.get_node(tdd_id) do
          {:ok, :true_terminal} ->
            Store.true_node_id()

          {:ok, :false_terminal} ->
            Store.false_node_id()

          {:ok, {var, y, n, d}} ->
            # Dispatch to the handler for recursive variables.
            case var do
              {5, :c_head, constraint_id, _} ->
                handle_recursive_subproblem(
                  :simplify,
                  :head,
                  constraint_id,
                  {var, y, n, d},
                  sorted_assumptions,
                  new_context
                )

              {5, :d_tail, constraint_id, _} ->
                handle_recursive_subproblem(
                  :simplify,
                  :tail,
                  constraint_id,
                  {var, y, n, d},
                  sorted_assumptions,
                  new_context
                )

              {4, :b_element, index, constraint_id} ->
                handle_recursive_subproblem(
                  :simplify,
                  {:elem, index},
                  constraint_id,
                  {var, y, n, d},
                  sorted_assumptions,
                  new_context
                )

              _ ->
                # The rest of the logic for standard variables is unchanged.
                case Map.get(assumptions, var) do
                  true ->
                    do_simplify(y, sorted_assumptions, new_context)

                  false ->
                    do_simplify(n, sorted_assumptions, new_context)

                  :dc ->
                    do_simplify(d, sorted_assumptions, new_context)

                  nil ->
                    assumptions_imply_true =
                      Engine.check(Map.put(assumptions, var, false)) == :contradiction

                    assumptions_imply_false =
                      Engine.check(Map.put(assumptions, var, true)) == :contradiction

                    cond do
                      assumptions_imply_true and assumptions_imply_false ->
                        Store.false_node_id()

                      assumptions_imply_true ->
                        do_simplify(y, Enum.sort(Map.put(assumptions, var, true)), new_context)

                      assumptions_imply_false ->
                        do_simplify(n, Enum.sort(Map.put(assumptions, var, false)), new_context)

                      true ->
                        s_y =
                          do_simplify(y, Enum.sort(Map.put(assumptions, var, true)), new_context)

                        s_n =
                          do_simplify(n, Enum.sort(Map.put(assumptions, var, false)), new_context)

                        s_d =
                          do_simplify(d, Enum.sort(Map.put(assumptions, var, :dc)), new_context)

                        Store.find_or_create_node(var, s_y, s_n, s_d)
                    end
                end
            end
        end
      end
    end
  end

  # --- Unary Operation: Substitute ---
  # FIX: The implementation of substitute needs to change.

  @spec substitute(non_neg_integer(), non_neg_integer(), non_neg_integer()) :: non_neg_integer()
  def substitute(root_id, from_id, to_id) do
    if root_id == from_id, do: to_id, else: do_substitute(root_id, from_id, to_id)
  end

  # This helper inspects and replaces TDD IDs embedded in predicate variables.
  defp substitute_in_var(var, from_id, to_id) do
    case var do
      {4, :b_element, index, ^from_id} -> {4, :b_element, index, to_id}
      {5, :c_head, ^from_id, nil} -> {5, :c_head, to_id, nil}
      {5, :d_tail, ^from_id, nil} -> {5, :d_tail, to_id, nil}
      _other -> var
    end
  end

  defp do_substitute(root_id, from_id, to_id) do
    cache_key = {:substitute, root_id, from_id, to_id}

    case Store.get_op_cache(cache_key) do
      {:ok, result_id} ->
        result_id

      :not_found ->
        result_id =
          case Store.get_node(root_id) do
            {:ok, :true_terminal} ->
              Store.true_node_id()

            {:ok, :false_terminal} ->
              Store.false_node_id()

            {:ok, {var, y, n, d}} ->
              # FIX: Substitute within the variable term itself.
              new_var = substitute_in_var(var, from_id, to_id)
              new_y = substitute(y, from_id, to_id)
              new_n = substitute(n, from_id, to_id)
              new_d = substitute(d, from_id, to_id)
              Store.find_or_create_node(new_var, new_y, new_n, new_d)

            {:error, reason} ->
              raise "substitute encountered an error getting node #{root_id}: #{reason}"
          end

        Store.put_op_cache(cache_key, result_id)
        result_id
    end
  end

  # --- Coinductive Emptiness Check ---
  # This function is correct and does not need to be changed.
  @spec check_emptiness(non_neg_integer()) :: non_neg_integer()
  def check_emptiness(tdd_id) do
    cache_key = {:check_emptiness, tdd_id}

    case Store.get_op_cache(cache_key) do
      {:ok, id} ->
        id

      :not_found ->
        assumptions_list = []
        result_id = do_check_emptiness(tdd_id, assumptions_list, MapSet.new())
        Store.put_op_cache(cache_key, result_id)
        result_id
    end
  end

  # This function is correct and does not need to be changed.
  defp do_check_emptiness(tdd_id, sorted_assumptions, context) do
    current_state = {tdd_id, sorted_assumptions}

    if MapSet.member?(context, current_state) do
      Store.false_node_id()
    else
      new_context = MapSet.put(context, current_state)
      assumptions = Map.new(sorted_assumptions)

      if Engine.check(assumptions) == :contradiction do
        Store.false_node_id()
      else
        IO.inspect(tdd_id, label: "Fetching node in do_check_emptiness")
        # Debug.print_tdd_graph(tdd_id)

        case Store.get_node(tdd_id) do
          {:ok, :true_terminal} ->
            Store.true_node_id()

          {:ok, :false_terminal} ->
            Store.false_node_id()

          {:ok, {var, y, n, d}} ->
            # Dispatch to the handler for recursive variables.
            case var do
              {5, :c_head, constraint_id, _} ->
                handle_recursive_subproblem(
                  :check_emptiness,
                  :head,
                  constraint_id,
                  {var, y, n, d},
                  sorted_assumptions,
                  new_context
                )

              {5, :d_tail, constraint_id, _} ->
                handle_recursive_subproblem(
                  :check_emptiness,
                  :tail,
                  constraint_id,
                  {var, y, n, d},
                  sorted_assumptions,
                  new_context
                )

              {4, :b_element, index, constraint_id} ->
                handle_recursive_subproblem(
                  :check_emptiness,
                  {:elem, index},
                  constraint_id,
                  {var, y, n, d},
                  sorted_assumptions,
                  new_context
                )

              _ ->
                # The rest of the logic is the same as the do_simplify counterpart
                case Map.get(assumptions, var) do
                  true ->
                    do_check_emptiness(y, sorted_assumptions, new_context)

                  false ->
                    do_check_emptiness(n, sorted_assumptions, new_context)

                  :dc ->
                    do_check_emptiness(d, sorted_assumptions, new_context)

                  nil ->
                    assumptions_imply_true =
                      Engine.check(Map.put(assumptions, var, false)) == :contradiction

                    assumptions_imply_false =
                      Engine.check(Map.put(assumptions, var, true)) == :contradiction

                    cond do
                      assumptions_imply_true and assumptions_imply_false ->
                        Store.false_node_id()

                      assumptions_imply_true ->
                        do_check_emptiness(
                          y,
                          Enum.sort(Map.put(assumptions, var, true)),
                          new_context
                        )

                      assumptions_imply_false ->
                        do_check_emptiness(
                          n,
                          Enum.sort(Map.put(assumptions, var, false)),
                          new_context
                        )

                      true ->
                        s_y =
                          do_check_emptiness(
                            y,
                            Enum.sort(Map.put(assumptions, var, true)),
                            new_context
                          )

                        s_n =
                          do_check_emptiness(
                            n,
                            Enum.sort(Map.put(assumptions, var, false)),
                            new_context
                          )

                        s_d =
                          do_check_emptiness(
                            d,
                            Enum.sort(Map.put(assumptions, var, :dc)),
                            new_context
                          )

                        Store.find_or_create_node(var, s_y, s_n, s_d)
                    end
                end
            end
        end
      end
    end
  end

  # This function, containing our previous fix, is correct and does not need to be changed.
  defp handle_recursive_subproblem(
         algo_type,
         sub_key,
         # This is a TDD ID for the constraint on the sub-problem.
         constraint_id,
         node_details,
         sorted_assumptions,
         # This is the coinductive context (a MapSet).
         context
       ) do
    {var, y, n, d} = node_details
    assumptions = Map.new(sorted_assumptions)

    # 1. Build the TDD for the sub-problem's effective type by intersecting all
    #    its constraints from the current assumption set.
    op_intersect = fn
      :false_terminal, _ -> :false_terminal
      _, :false_terminal -> :false_terminal
      t, :true_terminal -> t
      :true_terminal, t -> t
    end

    sub_problem_constraints =
      Enum.filter(assumptions, fn {v, _} ->
        (Tdd.Predicate.Info.get_traits(v) || %{})[:sub_key] == sub_key
      end)

    sub_problem_tdd_id =
      Enum.reduce(sub_problem_constraints, Store.true_node_id(), fn {var, val}, acc_id ->
        constraint_for_this_assumption = Engine.unwrap_var(var)

        id_to_intersect =
          if val, do: constraint_for_this_assumption, else: negate(constraint_for_this_assumption)

        apply(:intersect, op_intersect, acc_id, id_to_intersect)
      end)

    # 2. Check for the three logical outcomes:
    #    - Does the path imply the constraint is satisfied?
    #    - Does the path imply the constraint is violated?
    #    - Or are both outcomes still possible?

    # Implies satisfied: `sub_problem_tdd_id <: constraint_id`
    # This is equivalent to `(sub_problem_tdd_id & !constraint_id)` being empty.
    neg_constraint_id = negate(constraint_id)

    intersect_sub_with_neg_constraint =
      apply(:intersect, op_intersect, sub_problem_tdd_id, neg_constraint_id)

    implies_satisfied =
      check_emptiness(intersect_sub_with_neg_constraint) == Store.false_node_id()

    # Implies violated: `sub_problem_tdd_id` and `constraint_id` are disjoint.
    # This is equivalent to `(sub_problem_tdd_id & constraint_id)` being empty.
    intersect_sub_with_constraint =
      apply(:intersect, op_intersect, sub_problem_tdd_id, constraint_id)

    implies_violated = check_emptiness(intersect_sub_with_constraint) == Store.false_node_id()

    # 3. Branch based on the logical outcome.
    cond do
      implies_satisfied and implies_violated ->
        # The sub-problem itself must be empty/impossible under the current assumptions.
        # This whole path is a contradiction.
        Store.false_node_id()

      implies_satisfied ->
        # The constraint is guaranteed by the path. Follow the 'yes' branch.
        # The assumptions already imply this, so no change to them is needed.
        case algo_type do
          :simplify -> do_simplify(y, sorted_assumptions, context)
          :check_emptiness -> do_check_emptiness(y, sorted_assumptions, context)
        end

      implies_violated ->
        # The constraint is impossible given the path. Follow the 'no' branch.
        # We can add this new fact `{var, false}` to strengthen the assumptions.
        new_assumptions = Map.put(assumptions, var, false) |> Enum.sort()

        case algo_type do
          :simplify -> do_simplify(n, new_assumptions, context)
          :check_emptiness -> do_check_emptiness(n, new_assumptions, context)
        end

      true ->
        # Neither outcome is guaranteed. Both are possible.
        # We must explore both branches and combine the results.
        new_assumptions_for_no_branch = Map.put(assumptions, var, false) |> Enum.sort()

        case algo_type do
          :check_emptiness ->
            # Is there ANY path to true? Explore both and take their union.
            res_y = do_check_emptiness(y, sorted_assumptions, context)
            res_n = do_check_emptiness(n, new_assumptions_for_no_branch, context)

            # Define local terminal logic for union
            op_union = fn
              :true_terminal, _ -> :true_terminal
              _, :true_terminal -> :true_terminal
              t, :false_terminal -> t
              :false_terminal, t -> t
            end

            apply(:sum, op_union, res_y, res_n)

          :simplify ->
            # Simplify both sub-trees and rebuild the node, as we cannot simplify this node away.
            res_y = do_simplify(y, sorted_assumptions, context)
            res_n = do_simplify(n, new_assumptions_for_no_branch, context)
            # 'dc' branch is independent of 'var'
            res_d = do_simplify(d, sorted_assumptions, context)
            Store.find_or_create_node(var, res_y, res_n, res_d)
        end
    end
  end
end

defmodule Tdd.Compiler do
  @moduledoc """
  Compiles a `TypeSpec` into a canonical TDD ID.
  REFAC: This module now embeds TDD IDs into recursive predicate variables
  instead of raw TypeSpecs, a key part of the architectural decoupling.
  """
  alias Tdd.TypeSpec
  alias Tdd.Variable
  alias Tdd.Store
  alias Tdd.Algo
  alias Tdd.Debug

  @doc "The main public entry point. Takes a spec and returns its TDD ID."
  @spec spec_to_id(TypeSpec.t()) :: non_neg_integer()
  def spec_to_id(spec) do
    # FIX: State initialization is moved to the public API boundary (`is_subtype/2`),
    # making this function a pure compiler component.
    normalized_spec = TypeSpec.normalize(spec)
    compile_normalized_spec(normalized_spec, %{})
  end

  # This context-aware compilation function from the previous fix remains correct.
  defp compile_normalized_spec(normalized_spec, context) do
    sorted_context = Enum.sort(context)
    cache_key = {:compile_spec_with_context, normalized_spec, sorted_context}

    case Store.get_op_cache(cache_key) do
      {:ok, id} ->
        id

      :not_found ->
        id_to_cache =
          case normalized_spec do
            {:type_var, var_name} ->
              Map.get(context, var_name) ||
                raise "Tdd.Compiler: Unbound type variable: #{inspect(var_name)}"

            {:mu, var_name, body_spec} ->
              placeholder_id = Store.create_placeholder({:mu_knot_tying_for, var_name, body_spec})
              new_context = Map.put(context, var_name, placeholder_id)
              body_id = compile_normalized_spec(body_spec, new_context)
              {:ok, body_details} = Store.get_node(body_id)
              Store.update_node_in_place(placeholder_id, {:ok, body_details})
              Algo.simplify(placeholder_id)

            other_spec ->
              raw_id = do_structural_compile(other_spec, context)
              Algo.simplify(raw_id)
          end

        Store.put_op_cache(cache_key, id_to_cache)
        id_to_cache
    end
  end

  defp do_structural_compile(structural_spec, context) do
    case structural_spec do
      :any ->
        Store.true_node_id()

      :none ->
        Store.false_node_id()

      :atom ->
        create_base_type_tdd(Variable.v_is_atom())

      :integer ->
        create_base_type_tdd(Variable.v_is_integer())

      :list ->
        create_base_type_tdd(Variable.v_is_list())

      :tuple ->
        create_base_type_tdd(Variable.v_is_tuple())

      {:literal, val} when is_atom(val) ->
        compile_value_equality(:atom, Variable.v_atom_eq(val), context)

      {:literal, val} when is_integer(val) ->
        compile_value_equality(:integer, Variable.v_int_eq(val), context)

      {:literal, []} ->
        compile_value_equality(:list, Variable.v_list_is_empty(), context)

      {:integer_range, min, max} ->
        compile_integer_range(min, max, context)

      {:union, specs} ->
        Enum.map(specs, &compile_normalized_spec(&1, context))
        |> Enum.reduce(Store.false_node_id(), fn id, acc ->
          Algo.apply(:sum, &op_union_terminals/2, id, acc)
        end)

      {:intersect, specs} ->
        Enum.map(specs, &compile_normalized_spec(&1, context))
        |> Enum.reduce(Store.true_node_id(), fn id, acc ->
          Algo.apply(:intersect, &op_intersect_terminals/2, id, acc)
        end)

      {:negation, sub_spec} ->
        Algo.negate(compile_normalized_spec(sub_spec, context))

      {:cons, head_spec, tail_spec} ->
        id_list = compile_normalized_spec(:list, context)
        id_is_empty = create_base_type_tdd(Variable.v_list_is_empty())
        id_not_is_empty = Algo.negate(id_is_empty)

        non_empty_list_id =
          Algo.apply(:intersect, &op_intersect_terminals/2, id_list, id_not_is_empty)

        head_id = compile_normalized_spec(head_spec, context)
        tail_id = compile_normalized_spec(tail_spec, context)

        head_checker_var = Variable.v_list_head_pred(head_id)
        head_checker_tdd = create_base_type_tdd(head_checker_var)

        tail_checker_var = Variable.v_list_tail_pred(tail_id)
        tail_checker_tdd = create_base_type_tdd(tail_checker_var)

        [non_empty_list_id, head_checker_tdd, tail_checker_tdd]
        |> Enum.reduce(Store.true_node_id(), fn id, acc ->
          Algo.apply(:intersect, &op_intersect_terminals/2, id, acc)
        end)

      {:tuple, elements_specs} ->
        size = length(elements_specs)
        base_id = compile_normalized_spec(:tuple, context)
        size_tdd = create_base_type_tdd(Variable.v_tuple_size_eq(size))
        initial_id = Algo.apply(:intersect, &op_intersect_terminals/2, base_id, size_tdd)

        elements_specs
        |> Enum.with_index()
        |> Enum.reduce(initial_id, fn {elem_spec, index}, acc_id ->
          elem_id = compile_normalized_spec(elem_spec, context)
          elem_checker_var = Variable.v_tuple_elem_pred(index, elem_id)
          elem_checker_tdd = create_base_type_tdd(elem_checker_var)

          Algo.apply(:intersect, &op_intersect_terminals/2, acc_id, elem_checker_tdd)
        end)

      {:type_lambda, _, _} ->
        raise "Tdd.Compiler: Cannot compile :type_lambda directly. Spec should be ground. Spec: #{inspect(structural_spec)}"

      {:type_apply, _, _} ->
        raise "Tdd.Compiler: Cannot compile :type_apply directly. Spec should be ground and fully beta-reduced. Spec: #{inspect(structural_spec)}"

      _ ->
        raise "Tdd.Compiler.do_structural_compile: Unhandled structural spec form: #{inspect(structural_spec)}"
    end
  end

  defp create_base_type_tdd(var),
    do:
      Store.find_or_create_node(
        var,
        Store.true_node_id(),
        Store.false_node_id(),
        Store.false_node_id()
      )

  defp compile_value_equality(base_type_spec, value_var, context) do
    eq_node = create_base_type_tdd(value_var)
    base_node_id = compile_normalized_spec(base_type_spec, context)
    Algo.apply(:intersect, &op_intersect_terminals/2, base_node_id, eq_node)
  end

  defp compile_integer_range(min, max, context) do
    base_id = compile_normalized_spec(:integer, context)

    lt_min_tdd = if min != :neg_inf, do: create_base_type_tdd(Variable.v_int_lt(min))

    gte_min_tdd =
      if lt_min_tdd, do: Algo.negate(lt_min_tdd), else: compile_normalized_spec(:any, context)

    id_with_min = Algo.apply(:intersect, &op_intersect_terminals/2, base_id, gte_min_tdd)

    if max == :pos_inf do
      id_with_min
    else
      lt_max_plus_1_tdd = create_base_type_tdd(Variable.v_int_lt(max + 1))
      Algo.apply(:intersect, &op_intersect_terminals/2, id_with_min, lt_max_plus_1_tdd)
    end
  end

  # --- Terminal Logic Helpers ---
  defp op_union_terminals(:true_terminal, _), do: :true_terminal
  defp op_union_terminals(_, :true_terminal), do: :true_terminal
  defp op_union_terminals(t, :false_terminal), do: t
  defp op_union_terminals(:false_terminal, t), do: t
  defp op_intersect_terminals(:false_terminal, _), do: :false_terminal
  defp op_intersect_terminals(_, :false_terminal), do: :false_terminal
  defp op_intersect_terminals(t, :true_terminal), do: t
  defp op_intersect_terminals(:true_terminal, t), do: t

  # --- Public Subtyping Check ---
  @doc "Checks if spec1 is a subtype of spec2 using TDDs."
  @spec is_subtype(TypeSpec.t(), TypeSpec.t()) :: boolean
  def is_subtype(spec1, spec2) do
    # FIX: Each subtyping check is an independent, top-level operation.
    # It must initialize the store to guarantee a clean slate and prevent
    # state from leaking between unrelated checks.
    Store.init()

    id1 = spec_to_id(spec1)
    Debug.print_tdd_graph(id1, "IS_SUBTYPE id1")
    id2 = spec_to_id(spec2)
    Debug.print_tdd_graph(id2, "IS_SUBTYPE id2")
    neg_id2 = Algo.negate(id2)
    Debug.print_tdd_graph(neg_id2, "IS_SUBTYPE neg_id2")
    intersect_id = Algo.apply(:intersect, &op_intersect_terminals/2, id1, neg_id2)
    Debug.print_tdd_graph(intersect_id, "IS_SUBTYPE intersect")

    # The crash was because `intersect_id` was not an integer, which should not
    # happen if `apply` works correctly on valid IDs. By fixing the state
    # initialization, `id1` and `id2` will now be valid in the same store context.
    final_id = Algo.check_emptiness(intersect_id)
    final_id == Store.false_node_id()
  end
end