defmodule Mix.Tasks.Zig.Precompile do
  use Mix.Task

  require Logger

  @moduledoc """
  precompiles a group of zig modules for all supported platforms.

  In order to cross-compile windows-msvc, you will need to have the following environment
  variables set:

  - `MSVC_ROOT`: path to the root of a Visual Studio installation
  - `WINSDK_ROOT`: path to the root of a Windows SDK installation
  - `WINSDK_VERSION`: version of the Windows SDK to use, e.g. `10.0.22621.0`

  TODO: add a way to specify subsets of platforms.
  """

  @shortdoc "precompiles a zig module"

  @msvc_dir System.get_env("MSVC_ROOT", "")
  @winsdk_dir System.get_env("WINSDK_ROOT", "")
  @winsdk_version System.get_env("WINSDK_VER", "")

  if Enum.all?([@msvc_dir, @winsdk_dir, @winsdk_version], &(&1 != "")) or
       match?({_, :nt}, :os.type()) do
    @windows [:msvc, :gnu]
  else
    Logger.warning("not cross-compiling windows-msvc")
    @windows [:gnu]
  end

  @triples [
    aarch64: [freebsd: :none, linux: [:gnu, :musl], macos: :none],
    arm: [linux: [:gnueabi, :gnueabihf, :musleabi, :musleabihf]],
    x86_64: [freebsd: :none, linux: [:gnu, :musl], macos: :none, windows: @windows],
    x86: [linux: [:gnu, :musl], windows: :gnu]
  ]

  defp all_targets do
    for {arch, targets} <- @triples,
        {os, platforms} <- targets,
        platform <- List.wrap(platforms) do
      {arch, os, platform}
    end
  end

  def run([file]) do
    shas =
      all_targets()
      |> Enum.map(fn {arch, os, platform} -> compile(file, arch, os, platform) end)
      |> inspect(pretty: true)
      |> String.split("\n")
      |> Enum.join("\n  ")

    IO.puts("""
    Precompilation complete.  Add the following keyword list to the shasum parameters to in web downloads,
    e.g.:

    ```
    @shasum #{shas}

    use Zig, precompiled: {:web, "https://host-path-to-file/Elixir.MyModule.#VERSION.#TRIPLE.#EXT", @shasum}}
    ```
    """)
  end

  defp compile(file, arch, os, platform) do
    this = self()
    callback = fn file -> send(this, {:result, file}) end

    Application.put_env(:zigler, :precompiling, {arch, os, platform, callback})
    [{module, _}] = Code.compile_file(file)

    try do
      receive do
        {:result, file} ->
          dll_sha =
            file
            |> File.read!()
            |> then(&:crypto.hash(:sha256, &1))
            |> Base.encode16(case: :lower)

          if os == :windows do
            pdb_file = String.replace_suffix(file, ".dll", ".pdb")

            pdb_sha =
              pdb_file
              |> File.read!()
              |> then(&:crypto.hash(:sha256, &1))
              |> Base.encode16(case: :lower)

            {:"#{arch}-#{os}-#{platform}", %{dll: dll_sha, pdb: pdb_sha}}
          else
            {:"#{arch}-#{os}-#{platform}", dll_sha}
          end
      end
    after
      # clean up the build directory
      module
      |> Zig.Builder.staging_directory()
      |> File.rm_rf!()

      Process.sleep(100)
    end
  end
end