defmodule Mix.Tasks.Lockstep.Shrink do
  @moduledoc """
  Reduce a failing Lockstep trace to a minimal schedule that still
  triggers the same bug.

  ## Usage

      mix lockstep.shrink --trace traces/foo.lockstep \\
                          --run "MyApp.RaceTest.lost_update_body" \\
                          --file test/examples/foo_test.exs

  Loads `--file` (so test modules become available), then drives
  `Lockstep.Shrink.shrink/3` against the trace. Saves a `.shrunk.lockstep`
  alongside the original and prints a one-line summary.

  Use `--verbose` for per-attempt log lines.
  """

  use Mix.Task

  @shortdoc "Shrink a saved Lockstep trace to a minimal repro"

  @impl true
  def run(args) do
    Mix.Task.run("compile")

    {opts, _, _} =
      OptionParser.parse(args,
        strict: [
          trace: :string,
          run: :string,
          file: :string,
          verbose: :boolean,
          max_attempts: :integer,
          iter_timeout: :integer
        ],
        aliases: [t: :trace, r: :run, f: :file, v: :verbose]
      )

    trace_path = Keyword.get(opts, :trace) || Mix.raise("--trace <path> is required")
    run_ref = Keyword.get(opts, :run) || Mix.raise("--run \"Mod.fn\" is required")
    file = Keyword.get(opts, :file)

    if file, do: Code.require_file(file)

    fun = resolve_function(run_ref)

    shrink_opts =
      []
      |> add_opt(opts, :verbose, false)
      |> add_opt(opts, :max_attempts, 500)
      |> add_opt(opts, :iter_timeout, 10_000)

    case Lockstep.Shrink.shrink(fun, trace_path, shrink_opts) do
      {:ok, info} ->
        Mix.shell().info("""

        Lockstep shrink succeeded:
          original schedule decisions: #{info.original_length}
          shrunk schedule decisions:   #{info.new_length}
          attempts:                    #{info.attempts}
          bug signature:               #{inspect(info.signature)}
          shrunk trace:                #{info.new_trace_path}

        Replay it with:
            mix lockstep.replay --trace #{info.new_trace_path} \\
                                --run "#{run_ref}" \\
                                --file #{file || "<your test file>"}
        """)

      {:error, :no_smaller_schedule_found} ->
        Mix.shell().info("""

        Lockstep shrink: no smaller schedule reproduces the same bug.
        The original trace is already minimal under the strict-replay
        shrink algorithm.
        """)

      {:error, :original_does_not_reproduce} ->
        Mix.raise("""

        Lockstep shrink: the original trace did not reproduce a bug
        under strict replay. Likely causes: the test_fun argument
        doesn't match the original test body, or the user code has
        nondeterminism that's making replay diverge.
        """)

      {:error, :original_diverges_under_strict_replay} ->
        Mix.raise("""

        Lockstep shrink: original trace diverges under strict replay.
        The user code is nondeterministic in something other than
        message ordering (raw RNG, ETS, system time, NIFs, real I/O).
        Lockstep can only shrink traces whose nondeterminism is
        entirely in the controlled schedule.
        """)

      {:error, reason} ->
        Mix.raise("Lockstep shrink failed: #{inspect(reason)}")
    end
  end

  # ============================================================

  defp add_opt(acc, opts, key, default) do
    [{key, Keyword.get(opts, key, default)} | acc]
  end

  defp resolve_function(string) do
    case String.split(string, ".") do
      parts when length(parts) >= 2 ->
        {fun_name, mod_parts} = List.pop_at(parts, -1)
        mod = Module.concat(mod_parts)

        unless Code.ensure_loaded?(mod) do
          Mix.raise("module #{inspect(mod)} is not loaded; pass --file to load it")
        end

        fun_atom = String.to_atom(fun_name)

        unless function_exported?(mod, fun_atom, 0) do
          Mix.raise(
            "function #{inspect(mod)}.#{fun_name}/0 is not defined (must be 0-arity and public)"
          )
        end

        fn -> apply(mod, fun_atom, []) end

      _ ->
        Mix.raise("--run must be \"Module.fun\" (e.g. \"MyApp.MyTest.race_body\")")
    end
  end
end