defmodule Tdd.Debug.TracerData do
  @moduledoc """
  Provides functions to sanitize Elixir terms for tracing and JSON serialization.
  """

  @max_sanitize_depth 5

  def sanitize_values(list) when is_list(list) do
    Enum.map(list, &sanitize_value/1)
  end

  def sanitize_value(val) do
    do_sanitize(val, @max_sanitize_depth)
  end

  defp do_sanitize(val, _depth_limit)
       when is_integer(val) or is_float(val) or is_atom(val) or is_boolean(val),
       do: val

  # Assume strings are fine
  defp do_sanitize(val, _depth_limit) when is_binary(val), do: val
  defp do_sanitize(nil, _depth_limit), do: nil

  defp do_sanitize(val, _depth_limit) when is_pid(val), do: inspect(val)
  defp do_sanitize(val, _depth_limit) when is_port(val), do: inspect(val)
  defp do_sanitize(val, _depth_limit) when is_reference(val), do: inspect(val)

  defp do_sanitize(val, _depth_limit) when is_function(val) do
    try do
      arity = Function.info(val, :arity) |> elem(1)

      case Function.info(val, :name) do
        {:name, name} when is_atom(name) and name != :"" ->
          "<Function #{Atom.to_string(name)}/#{arity}>"

        _ ->
          "<Function/#{arity}>"
      end
    catch
      _, _ -> "<Function>"
    end
  end

  defp do_sanitize(val, depth_limit) when is_list(val) do
    if depth_limit <= 0 do
      "[...trimmed_list...]"
    else
      Enum.map(val, &do_sanitize(&1, depth_limit - 1))
    end
  end

  defp do_sanitize(val, depth_limit) when is_map(val) do
    if depth_limit <= 0 do
      "%{...trimmed_map...}"
    else
      # Regular map
      if Map.has_key?(val, :__struct__) do
        module = val.__struct__
        # Sanitize only exposed fields, not internal :__meta__ or the :__struct__ key itself
        data_to_sanitize = val |> Map.delete(:__struct__) |> Map.delete(:__meta__)

        sanitized_fields =
          data_to_sanitize
          |> Enum.map(fn {k, v} -> {k, do_sanitize(v, depth_limit - 1)} end)
          |> Enum.into(%{})

        %{type: :struct, module: module, fields: sanitized_fields}
      else
        val
        |> Enum.map(fn {k, v} ->
          {do_sanitize(k, depth_limit - 1), do_sanitize(v, depth_limit - 1)}
        end)
        |> Enum.into(%{})
      end
    end
  end

  defp do_sanitize(val, depth_limit) when is_tuple(val) do
    if depth_limit <= 0 do
      "{...trimmed_tuple...}"
    else
      val |> Tuple.to_list() |> Enum.map(&do_sanitize(&1, depth_limit - 1)) |> List.to_tuple()
    end
  end

  # Fallback for other types (e.g., DateTime, Date, Time, NaiveDateTime)
  defp do_sanitize(val, _depth_limit) do
    # For common structs that have `to_string` or are Calendar types
    cond do
      is_struct(val, DateTime) or is_struct(val, Date) or is_struct(val, Time) or
          is_struct(val, NaiveDateTime) ->
        # Jason can handle these if they are ISO8601 strings
        inspect(val)

      true ->
        # Generic struct or unknown type
        if is_struct(val) do
          # A simpler representation
          "Struct<#{inspect(val.__struct__)}>"
        else
          # Fallback, converts to string
          inspect(val)
        end
    end
  end

  def sanitize_error(exception_instance, stacktrace) do
    %{
      type: :error,
      class: Atom.to_string(exception_instance.__struct__),
      message: Exception.message(exception_instance),
      # Sanitize stacktrace to keep it manageable
      stacktrace: Enum.map(stacktrace, &Exception.format_stacktrace_entry/1) |> Enum.take(15)
    }
  end
end

defmodule Tdd.Debug do
  @moduledoc """
  Provides macros to wrap `def` and `defp` for function call/return tracing.
  Logs arguments and return values using `IO.inspect`.
  Builds an in-memory call tree data structure per traced process.
  """

  # --- Agent for Tracing State ---
  @agent_name Tdd.Debug.StateAgent

  def init_agent_if_needed do
    case Process.whereis(@agent_name) do
      nil -> Agent.start_link(fn -> MapSet.new() end, name: @agent_name)
      _pid -> :ok
    end

    :ok
  end

  def add_traced_pid(pid) when is_pid(pid) do
    init_agent_if_needed()
    Agent.update(@agent_name, &MapSet.put(&1, pid))
  end

  def remove_traced_pid(pid) when is_pid(pid) do
    case Process.whereis(@agent_name) do
      nil -> :ok
      agent_pid -> Agent.cast(agent_pid, fn state -> MapSet.delete(state, pid) end)
    end
  end

  def is_pid_traced?(pid) when is_pid(pid) do
    case Process.whereis(@agent_name) do
      nil ->
        false

      agent_pid ->
        try do
          Agent.get(agent_pid, &MapSet.member?(&1, pid), :infinity)
        rescue
          # Agent might have died
          _e in [Exit, ArgumentError] -> false
        end
    end
  end

  # --- Tracing Data Storage ---
  @doc """
  Initializes the tracing data structures in the process dictionary.
  """
  def init_trace_storage do
    Process.put(:tdd_debug_call_stack, [])
    Process.put(:tdd_debug_session_roots, [])
    :ok
  end

  @doc """
  Retrieves the collected trace data (a list of root call tree nodes)
  for the current process and clears it from the process dictionary.
  Children within each node and the root list itself are reversed to be in chronological order.
  """
  def get_and_clear_trace_data do
    data = Process.get(:tdd_debug_session_roots, [])
    # Reverse children recursively and then reverse the list of roots
    processed_data = Enum.map(data, &reverse_children_recursively/1) |> Enum.reverse()

    Process.delete(:tdd_debug_session_roots)
    Process.delete(:tdd_debug_call_stack)
    processed_data
  end

  defp reverse_children_recursively(node) do
    if !node.children or node.children == [] do
      node
    else
      reversed_and_processed_children =
        Enum.map(node.children, &reverse_children_recursively/1)
        # Children were added prepending, so reverse for chronological
        |> Enum.reverse()

      %{node | children: reversed_and_processed_children}
    end
  end

  # --- Tracing Control Functions ---
  @doc "Enables function call tracing for the current process."
  def enable_tracing do
    pid_to_trace = self()
    add_traced_pid(pid_to_trace)
    # Initialize storage for call trees
    init_trace_storage()

    ref = Process.monitor(pid_to_trace)

    Process.spawn(
      fn ->
        receive do
          {:DOWN, ^ref, :process, ^pid_to_trace, _reason} ->
            remove_traced_pid(pid_to_trace)
        after
          # 1 hour safety timeout
          3_600_000 ->
            remove_traced_pid(pid_to_trace)
        end
      end,
      # Changed from [:monitor] to [] as monitor is implicit with Process.monitor
      [:monitor]
    )

    :ok
  end

  @doc "Disables function call tracing for the current process."
  def disable_tracing do
    remove_traced_pid(self())
    # Note: Does not clear trace data, get_and_clear_trace_data() does that.
    :ok
  end

  @doc """
  Runs the given 0-arity function with tracing enabled.
  Returns a tuple: `{{:ok, result} | {:error, {exception, stacktrace}}, trace_data}`.
  Trace data is a list of call tree root nodes.
  """
  def run(fun) when is_function(fun, 0) do
    enable_tracing()

    outcome =
      try do
        {:ok, fun.()}
      rescue
        kind -> {:error, {kind, __STACKTRACE__}}
      end

    trace_data = get_and_clear_trace_data()
    disable_tracing()
    # trace_data
    #       |> IO.inspect()
    binary = JSON.encode!(trace_data)
    File.write("trace.json", binary)
    {outcome, trace_data}
  end

  # --- Process Dictionary for Call Depth (used for printing indent) ---
  defp get_depth, do: Process.get(:tdd_debug_depth, 0)

  def increment_depth do
    new_depth = get_depth() + 1
    Process.put(:tdd_debug_depth, new_depth)
    new_depth
  end

  def decrement_depth do
    new_depth = max(0, get_depth() - 1)
    Process.put(:tdd_debug_depth, new_depth)
    new_depth
  end

  # --- Core Macro Logic ---
  defmacro __using__(_opts) do
    quote do
      import Kernel, except: [def: 1, def: 2, defp: 1, defp: 2]
      # Import this module's public functions/macros
      import Tdd.Debug
      # Ensure this module is compiled for macros
      require Tdd.Debug
    end
  end

  @doc false
  defmacro def(call, clauses \\ Keyword.new()) do
    generate_traced_function(:def, call, clauses, __CALLER__)
  end

  @doc false
  defmacro defp(call, clauses \\ Keyword.new()) do
    generate_traced_function(:defp, call, clauses, __CALLER__)
  end

  defp is_simple_variable_ast?(ast_node) do
    case ast_node do
      {var_name, _meta, context_module}
      when is_atom(var_name) and (context_module == nil or is_atom(context_module)) ->
        # Ensure it's not the underscore atom itself
        var_name != :_

      _ ->
        false
    end
  end

  defp generate_traced_function(type, call_ast, clauses, caller_env) do
    require Macro

    {function_name_ast, meta_call, original_args_patterns_ast_nullable} = call_ast
    original_args_patterns_ast_list = original_args_patterns_ast_nullable || []

    original_body_ast =
      if Keyword.keyword?(clauses) do
        Keyword.get(clauses, :do, clauses)
      else
        clauses
      end || quote(do: nil)

    # --- Argument Handling for Logging (Existing Logic) ---
    mapped_and_generated_vars_tuples =
      Enum.map(Enum.with_index(original_args_patterns_ast_list), fn {original_pattern_ast, index} ->
        td_arg_var = Macro.var(String.to_atom("__td_arg_#{index}__"), nil)

        {final_pattern_for_head, rhs_for_td_arg_assignment} =
          case original_pattern_ast do
            {:_, _, _} = underscore_ast ->
              {underscore_ast, quote(do: :__td_ignored_argument__)}

            {:=, _meta_assign, [lhs_of_assign, _rhs_of_assign]} = assignment_pattern_ast ->
              if is_simple_variable_ast?(lhs_of_assign) do
                {assignment_pattern_ast, lhs_of_assign}
              else
                captured_val_var = Macro.unique_var(String.to_atom("tdc_assign_#{index}"), Elixir)

                new_head_pattern =
                  quote do
                    unquote(captured_val_var) = unquote(assignment_pattern_ast)
                  end

                {new_head_pattern, captured_val_var}
              end

            {:\\, _meta_default, [pattern_before_default, _default_value_ast]} =
                default_arg_pattern_ast ->
              cond do
                is_simple_variable_ast?(pattern_before_default) ->
                  {default_arg_pattern_ast, pattern_before_default}

                match?({:=, _, [lhs_inner_assign, _]}, pattern_before_default) and
                    is_simple_variable_ast?(pattern_before_default |> elem(2) |> Enum.at(0)) ->
                  {:=, _, [lhs_inner_assign, _]} = pattern_before_default
                  {default_arg_pattern_ast, lhs_inner_assign}

                true ->
                  captured_val_var = Macro.unique_var(String.to_atom("tdc_def_#{index}"), Elixir)

                  new_head_pattern =
                    quote do
                      unquote(captured_val_var) = unquote(default_arg_pattern_ast)
                    end

                  {new_head_pattern, captured_val_var}
              end

            ast_node ->
              if is_simple_variable_ast?(ast_node) do
                {ast_node, ast_node}
              else
                captured_val_var = Macro.unique_var(String.to_atom("tdc_pat_#{index}"), Elixir)

                new_head_pattern =
                  quote do
                    unquote(captured_val_var) = unquote(ast_node)
                  end

                {new_head_pattern, captured_val_var}
              end
          end

        assignment_ast =
          quote do
            unquote(td_arg_var) = unquote(rhs_for_td_arg_assignment)
          end

        {final_pattern_for_head, assignment_ast, td_arg_var}
      end)

    {new_args_patterns_for_head_list, assignments_for_logging_vars_ast_list,
     generated_vars_to_log_asts} =
      if mapped_and_generated_vars_tuples == [] do
        {[], [], []}
      else
        # Transpose list of 3-element tuples into 3 lists
        list_of_lists = Enum.map(mapped_and_generated_vars_tuples, &Tuple.to_list(&1))

        [
          Enum.map(list_of_lists, &List.first(&1)),
          Enum.map(list_of_lists, &Enum.at(&1, 1)),
          Enum.map(list_of_lists, &Enum.at(&1, 2))
        ]
        |> then(fn [a, b, c] -> {a, b, c} end)
      end

    new_call_ast = {function_name_ast, meta_call, new_args_patterns_for_head_list}

    # --- Variables for Runtime Info ---
    # These vars will hold string versions of module/function name and arity at runtime
    # Hygienic vars to avoid collisions.
    module_name_runtime_var = Macro.var(:__td_module_name__, __MODULE__)
    printable_fn_name_runtime_var = Macro.var(:__td_printable_fn_name__, __MODULE__)
    arity_runtime_var = Macro.var(:__td_arity__, __MODULE__)

    # Arity calculation at macro expansion time
    arity_value = length(original_args_patterns_ast_list)

    traced_body_inner_ast =
      quote do
        # --- Resolve Module/Function/Arity Info at Runtime ---
        unquote(module_name_runtime_var) = Atom.to_string(unquote(caller_env.module))
        unquote(arity_runtime_var) = unquote(arity_value)

        # Resolve var if func name is from a var
        runtime_resolved_fn_name_val = unquote(function_name_ast)

        unquote(printable_fn_name_runtime_var) =
          if is_atom(runtime_resolved_fn_name_val) do
            Atom.to_string(runtime_resolved_fn_name_val)
          else
            # For unquote(var) or operators
            Macro.to_string(runtime_resolved_fn_name_val)
          end

        # --- Main Tracing Logic ---
        if Tdd.Debug.is_pid_traced?(self()) do
          # --- On Entry: Prepare Data & Log ---
          # Assign __td_arg_N__ vars
          unquote_splicing(assignments_for_logging_vars_ast_list)
          runtime_arg_values_for_node_and_log = [unquote_splicing(generated_vars_to_log_asts)]

          # sanitized_args_for_node = Tdd.Debug.TracerData.sanitize_values(runtime_arg_values_for_node_and_log)
          sanitized_args_for_node = inspect(runtime_arg_values_for_node_and_log)
          # For print indent & node depth
          current_call_print_depth = Tdd.Debug.increment_depth()

          # Create placeholder node, push to call_stack
          new_node_details = %{
            id: System.unique_integer([:positive, :monotonic]),
            # module: unquote(module_name_runtime_var),
            # function: unquote(printable_fn_name_runtime_var),
            function:
              "#{unquote(module_name_runtime_var)}.#{unquote(printable_fn_name_runtime_var)}",
            # arity: unquote(arity_runtime_var),
            args: sanitized_args_for_node,
            depth: current_call_print_depth,
            # timestamp_enter_monotonic: System.monotonic_time(),
            # Will be populated in reverse chronological order
            children: [],
            # Placeholder
            result: :__td_not_returned_yet__
          }

          Process.put(:tdd_debug_call_stack, [
            new_node_details | Process.get(:tdd_debug_call_stack, [])
          ])

          # Logging Call (existing logic)
          indent = String.duplicate("  ", current_call_print_depth - 1)

          IO.puts(
            "#{indent}CALL: #{unquote(module_name_runtime_var)}.#{unquote(printable_fn_name_runtime_var)}"
          )

          IO.puts("#{indent}  ARGS: #{inspect(runtime_arg_values_for_node_and_log)}")

          try do
            # --- Execute Original Body ---
            result_val = unquote(original_body_ast)

            # --- On Normal Exit: Finalize Node & Log ---
            [completed_node_placeholder | parent_stack_nodes] = Process.get(:tdd_debug_call_stack)
            ts_exit_monotonic = System.monotonic_time()

            # duration_ms = System.convert_time_unit(ts_exit_monotonic - completed_node_placeholder.timestamp_enter_monotonic, :native, :millisecond)
            sanitized_result_val = Tdd.Debug.TracerData.sanitize_value(result_val)

            finalized_node =
              completed_node_placeholder
              # |> Map.put(:timestamp_exit_monotonic, ts_exit_monotonic)
              # |> Map.put(:duration_ms, duration_ms)
              |> Map.put(:result, %{type: :ok, value: sanitized_result_val})

            # For print indent
            _ = Tdd.Debug.decrement_depth()

            # Update call stack: Add finalized_node to parent's children or session_roots
            new_call_stack_after_success =
              case parent_stack_nodes do
                [parent_on_stack | grand_parent_stack_nodes] ->
                  updated_parent = %{
                    parent_on_stack
                    | children: [finalized_node | parent_on_stack.children]
                  }

                  [updated_parent | grand_parent_stack_nodes]

                # This was a root call
                [] ->
                  Process.put(:tdd_debug_session_roots, [
                    finalized_node | Process.get(:tdd_debug_session_roots, [])
                  ])

                  # New call_stack is empty for this branch
                  []
              end

            Process.put(:tdd_debug_call_stack, new_call_stack_after_success)

            # Logging Return
            IO.puts(
              "#{indent}RETURN from #{unquote(module_name_runtime_var)}.#{unquote(printable_fn_name_runtime_var)}: #{inspect(result_val)}"
            )

            # Return actual result
            result_val
          rescue
            exception_class_err ->
              # --- On Error Exit: Finalize Node & Log ---
              error_instance_err = exception_class_err
              stacktrace_err = __STACKTRACE__

              [erroring_node_placeholder | parent_stack_nodes_err] =
                Process.get(:tdd_debug_call_stack)

              ts_exit_monotonic_err = System.monotonic_time()

              duration_ms_err =
                System.convert_time_unit(
                  ts_exit_monotonic_err - erroring_node_placeholder.timestamp_enter_monotonic,
                  :native,
                  :millisecond
                )

              sanitized_error_info =
                Tdd.Debug.TracerData.sanitize_error(error_instance_err, stacktrace_err)

              finalized_error_node =
                erroring_node_placeholder
                |> Map.put(:timestamp_exit_monotonic, ts_exit_monotonic_err)
                |> Map.put(:duration_ms, duration_ms_err)
                |> Map.put(:result, sanitized_error_info)

              # For print indent
              _ = Tdd.Debug.decrement_depth()

              # Update call stack: Add finalized_error_node to parent's children or session_roots
              new_call_stack_after_error =
                case parent_stack_nodes_err do
                  [parent_on_stack_err | grand_parent_stack_nodes_err] ->
                    updated_parent_err = %{
                      parent_on_stack_err
                      | children: [finalized_error_node | parent_on_stack_err.children]
                    }

                    [updated_parent_err | grand_parent_stack_nodes_err]

                  # This was a root call that errored
                  [] ->
                    Process.put(:tdd_debug_session_roots, [
                      finalized_error_node | Process.get(:tdd_debug_session_roots, [])
                    ])

                    # New call_stack is empty for this branch
                    []
                end

              Process.put(:tdd_debug_call_stack, new_call_stack_after_error)

              # Logging Error
              IO.puts(
                "#{indent}ERROR in #{unquote(module_name_runtime_var)}.#{unquote(printable_fn_name_runtime_var)}: #{inspect(error_instance_err)}"
              )

              # Reraise the original error
              reraise error_instance_err, stacktrace_err
          end
        else
          # --- Not Traced: Execute Original Body Directly ---
          unquote(original_body_ast)
        end
      end

    # --- Final Definition ---
    final_definition_ast =
      quote location: :keep do
        Kernel.unquote(type)(
          unquote(new_call_ast),
          do: unquote(traced_body_inner_ast)
        )
      end

    final_definition_ast
  end

  # --- TDD Graph Printing (Kept as it was, assuming it's for a different purpose) ---
  @doc "Prints a formatted representation of a TDD graph structure."
  def print_tdd_graph(id, label \\ "") do
    IO.puts("---#{label} TDD Graph (ID: #{id}) ---")
    do_print_tdd_node(id, 0, MapSet.new(), Tdd.Store)
    IO.puts("------------------------")
  end

  defp do_print_tdd_node(id, indent_level, visited, store_module) do
    prefix = String.duplicate("  ", indent_level)

    if MapSet.member?(visited, id) do
      IO.puts("#{prefix}ID #{id} -> [Seen, recursive link]")
      :ok
    else
      new_visited = MapSet.put(visited, id)
      # Assumes store_module.get_node/1 exists
      case store_module.get_node(id) do
        {:ok, :true_terminal} ->
          IO.puts("#{prefix}ID #{id} -> TRUE")

        {:ok, :false_terminal} ->
          IO.puts("#{prefix}ID #{id} -> FALSE")

        {:ok, {var, y_id, n_id, dc_id}} ->
          IO.puts("#{prefix}ID #{id}: IF #{inspect(var)}")

          case var do
            {_, :c_head, nested_id, _} when nested_id != nil ->
              print_tdd_graph(nested_id, "Var in #{id}")

            _ ->
              ""
          end

          IO.puts("#{prefix}  ├─ Yes (to ID #{y_id}):")
          do_print_tdd_node(y_id, indent_level + 2, new_visited, store_module)
          IO.puts("#{prefix}  ├─ No  (to ID #{n_id}):")
          do_print_tdd_node(n_id, indent_level + 2, new_visited, store_module)
          IO.puts("#{prefix}  └─ DC  (to ID #{dc_id}):")
          do_print_tdd_node(dc_id, indent_level + 2, new_visited, store_module)

        {:error, reason} ->
          IO.puts("#{prefix}ID #{id}: ERROR - #{reason}")
      end
    end
  end

  @doc "Logs a value with a label, using IO.inspect."
  def log(value, label, opts \\ []) do
    # You can add a flag here to disable all logging globally
    if true do
      # Default options for better visibility
      default_opts = [width: 120, pretty: true]
      final_opts = Keyword.merge(default_opts, opts)
      IO.inspect(value, [{:label, "[DEBUG] #{label}"} | final_opts])
    end

    # Return the original value to allow piping
    value
  end

  @doc "Gets and increments a depth counter for tracing recursion."
  defp get_depth() do
    depth = Process.get(:debug_depth, 0)
    Process.put(:debug_depth, depth + 1)
    String.duplicate("  ", depth)
  end

  @doc "Decrements the depth counter."
  defp dec_depth() do
    depth = Process.get(:debug_depth, 1) |> Kernel.-(1) |> max(0)
    Process.put(:debug_depth, depth)
  end

  @doc "Logs a message with indentation for tracing recursion."
  def log_entry(label) do
    prefix = get_depth()
    IO.inspect(prefix, label: "PREFIX")
    IO.puts("#{prefix}[DEBUG] >> #{label}")
  end

  @doc "Logs a return value with indentation."
  def log_exit(value, label) do
    dec_depth()
    prefix = String.duplicate("  ", Process.get(:debug_depth, 0))
    IO.inspect(value, label: prefix <> "[DEBUG] << #{label}")
    # Return the value
    value
  end
end