fly_deploy

defmodule FlyDeploy.ReloadScript do
  @moduledoc false
  # Downloads and applies upgrades on individual machines

  require Logger

  @doc """
  Normal hot upgrade on a running system.
  Uses suspend/resume for safe process upgrades.
  """
  def hot_upgrade(tarball_url, app) do
    IO.puts("Downloading tarball from #{tarball_url}...")
    {:ok, _} = Application.ensure_all_started(:req)

    # Use AWS SigV4 for authenticated download
    tmp_file_path = Path.join(System.tmp_dir!(), "fly_deploy_upgrade.tar.gz")

    response =
      Req.get!(tarball_url,
        into: File.stream!(tmp_file_path, [:write, :binary]),
        raw: true,
        receive_timeout: 60_000,
        connect_options: [timeout: 60_000],
        aws_sigv4: [
          access_key_id: System.fetch_env!("AWS_ACCESS_KEY_ID"),
          secret_access_key: System.fetch_env!("AWS_SECRET_ACCESS_KEY"),
          service: "s3",
          region: "auto"
        ]
      )

    Logger.info("[FlyDeploy.ReloadScript] Download response: #{response.status}")

    # Check downloaded tarball
    download_size = File.stat!(tmp_file_path).size
    IO.puts("  Downloaded: #{download_size} bytes")

    IO.puts("Extracting...")
    File.mkdir_p!("/tmp/upgrade")
    :erl_tar.extract(~c"#{tmp_file_path}", [:compressed, {:cwd, ~c"/tmp/upgrade"}])

    IO.puts("Copying beam files to currently loaded paths...")
    # Instead of copying the whole lib directory (which has version-specific paths),
    # we need to copy beam files directly to where they're currently loaded from
    # Find all beam files in the tarball
    beam_files = Path.wildcard("/tmp/upgrade/lib/**/ebin/*.beam")
    IO.puts("  Found #{length(beam_files)} beam files to copy")

    Enum.each(beam_files, fn tarball_beam_path ->
      # Extract just the beam filename
      beam_filename = Path.basename(tarball_beam_path)
      module_name = beam_filename |> String.replace_suffix(".beam", "") |> String.to_atom()

      # Find where this module is currently loaded from
      case :code.which(module_name) do
        path when is_list(path) ->
          loaded_path = List.to_string(path)
          # Copy the new beam file to overwrite the currently loaded one
          File.cp!(tarball_beam_path, loaded_path)

        _ ->
          # Module not currently loaded, skip it
          :ok
      end
    end)

    IO.puts("  ✓ Copied beam files to loaded paths")

    IO.puts("Performing safe hot upgrade...")

    # Perform the 4-phase upgrade
    result = safe_upgrade_application(app)

    IO.puts("  Modules reloaded: #{result.modules_reloaded}")

    IO.puts(
      "  Processes upgraded: #{result.processes_succeeded} succeeded, #{result.processes_failed} failed, #{result.processes_skipped} skipped"
    )

    IO.puts("✅ Hot reload complete!")
  end

  @doc """
  Replay a hot upgrade on application startup.

  This is called when a machine restarts and needs to reapply a hot upgrade
  that was previously deployed. Unlike hot_upgrade/2, this:
  1. Pre-loads all changed modules AFTER copying beam files
  2. Doesn't need to suspend/resume processes (they don't exist yet)
  """
  def replay_upgrade_startup(tarball_url, _app) do
    IO.puts("Replaying hot upgrade from #{tarball_url} on startup...")
    {:ok, _} = Application.ensure_all_started(:req)

    # Download tarball
    response =
      Req.get!(tarball_url,
        into: File.stream!("/tmp/upgrade.tar.gz", [:write, :binary]),
        raw: true,
        receive_timeout: 60_000,
        connect_options: [timeout: 60_000],
        aws_sigv4: [
          access_key_id: System.fetch_env!("AWS_ACCESS_KEY_ID"),
          secret_access_key: System.fetch_env!("AWS_SECRET_ACCESS_KEY"),
          service: "s3",
          region: "auto"
        ]
      )

    Logger.info("[FlyDeploy.ReloadScript] Download response: #{response.status}")
    download_size = File.stat!("/tmp/upgrade.tar.gz").size
    IO.puts("  Downloaded: #{download_size} bytes")

    # Extract tarball
    IO.puts("Extracting...")
    File.mkdir_p!("/tmp/upgrade")
    :erl_tar.extract(~c"/tmp/upgrade.tar.gz", [:compressed, {:cwd, ~c"/tmp/upgrade"}])

    # Copy beam files to loaded paths
    IO.puts("Copying beam files to currently loaded paths...")
    beam_files = Path.wildcard("/tmp/upgrade/lib/**/ebin/*.beam")
    IO.puts("  Found #{length(beam_files)} beam files to copy")

    Enum.each(beam_files, fn tarball_beam_path ->
      beam_filename = Path.basename(tarball_beam_path)
      module_name = beam_filename |> String.replace_suffix(".beam", "") |> String.to_atom()

      case :code.which(module_name) do
        path when is_list(path) ->
          loaded_path = List.to_string(path)
          File.cp!(tarball_beam_path, loaded_path)

        _ ->
          :ok
      end
    end)

    IO.puts("  ✓ Copied beam files to loaded paths")

    # Use :c.lm() to load modified modules (purges old code and loads new)
    IO.puts("Loading modified modules...")
    modified_modules = :c.lm()

    IO.puts(
      "  ✓ Loaded #{length(modified_modules)} modified modules: #{inspect(modified_modules)}"
    )

    IO.puts("✅ Startup hot upgrade replay complete!")
  end

  # Safely upgrades all application processes with new code.
  #
  # Returns a map with upgrade statistics:
  # - `modules_reloaded` - Number of modules that were reloaded
  # - `processes_succeeded` - Number of processes successfully upgraded
  # - `processes_failed` - Number of processes that failed to upgrade
  # - `processes_skipped` - Number of processes skipped (not GenServer/proc_lib)
  defp safe_upgrade_application(app) do
    Logger.info("[FlyDeploy.ReloadScript] Starting safe upgrade for #{app}")

    # Detect changed modules
    changed_modules = :code.modified_modules()
    Logger.info("[FlyDeploy.ReloadScript] Changed modules: #{inspect(changed_modules)}")

    # Find all processes that need upgrading BEFORE loading new code
    processes = find_processes_to_upgrade(changed_modules)
    Logger.info("[FlyDeploy.ReloadScript] Found #{length(processes)} processes to upgrade")

    # Phase 1: Suspend ALL processes
    Logger.info("[FlyDeploy.ReloadScript] Phase 1: Suspending all processes...")

    suspended_processes =
      Enum.map(processes, fn {pid, module} ->
        try do
          :sys.suspend(pid)
          Logger.info("[FlyDeploy.ReloadScript] Suspended process #{inspect(pid)} (#{module})")
          {:ok, pid, module}
        rescue
          e ->
            Logger.error(
              "[FlyDeploy.ReloadScript] Failed to suspend process #{inspect(pid)} (#{module}): #{Exception.message(e)}"
            )

            {:error, pid, module}
        end
      end)

    successfully_suspended =
      suspended_processes
      |> Enum.filter(fn {status, _, _} -> status == :ok end)

    Logger.info(
      "[FlyDeploy.ReloadScript] Suspended #{length(successfully_suspended)} processes successfully"
    )

    # Phase 2: Reload ALL changed modules (while all processes are suspended)
    Logger.info("[FlyDeploy.ReloadScript] Phase 2: Reloading all modules...")

    Enum.each(changed_modules, fn module ->
      Logger.info("[FlyDeploy.ReloadScript] Reloading module: #{module}")
      :code.purge(module)
      {:module, ^module} = :code.load_file(module)
    end)

    # Phase 3: Upgrade ALL processes (call code_change on each)
    Logger.info("[FlyDeploy.ReloadScript] Phase 3: Upgrading all processes...")

    upgrade_results =
      Enum.map(successfully_suspended, fn {:ok, pid, module} ->
        try do
          :sys.change_code(pid, module, :undefined, [])
          Logger.info("[FlyDeploy.ReloadScript] Upgraded process #{inspect(pid)} (#{module})")
          {:ok, pid, module}
        rescue
          e ->
            Logger.error(
              "[FlyDeploy.ReloadScript] Failed to upgrade process #{inspect(pid)} (#{module}): #{Exception.message(e)}"
            )

            {:error, pid, module}
        end
      end)

    # Phase 4: Resume ALL processes (even failed ones, to avoid leaving them suspended)
    Logger.info("[FlyDeploy.ReloadScript] Phase 4: Resuming all processes...")

    Enum.each(successfully_suspended, fn {:ok, pid, module} ->
      try do
        :sys.resume(pid)
        Logger.info("[FlyDeploy.ReloadScript] Resumed process #{inspect(pid)} (#{module})")
      rescue
        e ->
          Logger.error(
            "[FlyDeploy.ReloadScript] Failed to resume process #{inspect(pid)} (#{module}): #{Exception.message(e)}"
          )
      end
    end)

    # Phase 5: Trigger LiveView reloads for upgraded LiveView modules (if any)
    trigger_liveview_reloads(changed_modules)

    # Calculate stats
    succeeded = Enum.count(upgrade_results, fn {status, _, _} -> status == :ok end)
    failed = Enum.count(upgrade_results, fn {status, _, _} -> status == :error end)

    stats = %{
      modules_reloaded: length(changed_modules),
      processes_succeeded: succeeded,
      processes_failed: failed,
      processes_skipped: length(suspended_processes) - length(successfully_suspended)
    }

    Logger.info("[FlyDeploy.ReloadScript] Upgrade complete: #{inspect(stats)}")
    stats
  end

  defp find_processes_to_upgrade(changed_modules) do
    Process.list()
    |> Enum.map(fn pid ->
      case Process.info(pid, [:dictionary, :initial_call]) do
        [dictionary: dict, initial_call: _] ->
          # Check if process has a $initial_call set by proc_lib
          case Keyword.get(dict, :"$initial_call") do
            {module, _func, _arity} ->
              if module in changed_modules do
                {pid, module}
              else
                nil
              end

            _ ->
              nil
          end

        _ ->
          nil
      end
    end)
    |> Enum.reject(&is_nil/1)
  end

  defp trigger_liveview_reloads(changed_modules) do
    # Filter to only LiveView modules
    liveview_modules = Enum.filter(changed_modules, &liveview_module?/1)

    unless Enum.empty?(liveview_modules) do
      Logger.info("[FlyDeploy.ReloadScript] Phase 5: Triggering LiveView reloads...")

      Logger.info(
        "[FlyDeploy.ReloadScript] Found #{length(liveview_modules)} LiveView modules to reload"
      )

      # Find all LiveView process PIDs
      liveview_pids = find_liveview_processes()

      # For each upgraded LiveView module, send reload message to all LiveView processes
      Enum.each(liveview_modules, fn module ->
        # Get the actual source file path from the module's compile info
        source_path = get_source_path(module)

        # Send message directly to each LiveView process
        Enum.each(liveview_pids, fn pid ->
          send(pid, {:phoenix_live_reload, "fly_deploy", source_path})
        end)

        Logger.info(
          "[FlyDeploy.ReloadScript] Sent LiveView reload for #{module} to #{length(liveview_pids)} LiveView processes (#{source_path})"
        )
      end)
    end
  end

  defp find_liveview_processes do
    # Find all processes running Phoenix.LiveView
    Process.list()
    |> Enum.filter(fn pid ->
      case Process.info(pid, [:dictionary]) do
        [dictionary: dict] ->
          # LiveView processes have specific keys in their dictionary
          Keyword.has_key?(dict, :"$initial_call") and
            match?({Phoenix.LiveView.Channel, _, _}, Keyword.get(dict, :"$initial_call"))

        _ ->
          false
      end
    end)
  end

  defp get_source_path(module) do
    # Get source file path from module compile info
    case module.module_info(:compile) |> Keyword.get(:source) do
      path when is_list(path) -> List.to_string(path)
      path when is_binary(path) -> path
      _ -> inspect(module)
    end
  end

  defp liveview_module?(module) do
    # Check if module uses Phoenix.LiveView
    try do
      behaviours =
        module.module_info(:attributes)
        |> Keyword.get(:behaviour, [])

      Phoenix.LiveView in behaviours or Phoenix.LiveComponent in behaviours
    rescue
      _ -> false
    end
  end
end