extorch_vision

defmodule Mix.Tasks.Compile.Torchvision do
  @moduledoc """
  Custom Mix compiler that builds libtorchvision.so from source.

  Runs automatically during `mix compile` when `:torchvision` is included
  in the project's `:compilers` list. Skips if the library is already built.

  To force a rebuild, use `mix torchvision.build --force`.
  """

  use Mix.Task.Compiler

  @impl true
  def run(_args) do
    try do
      ExTorch.Vision.Build.ensure_built!()
      # Write manifest so the compiler knows it succeeded
      File.mkdir_p!(Path.dirname(manifest_path()))
      File.write!(manifest_path(), "built")
      {:ok, []}
    rescue
      e ->
        {:error,
         [
           %Mix.Task.Compiler.Diagnostic{
             file: "mix.exs",
             severity: :error,
             message: "Failed to build libtorchvision: #{Exception.message(e)}",
             position: 0,
             compiler_name: "torchvision"
           }
         ]}
    end
  end

  @impl true
  def manifests, do: [manifest_path()]

  @impl true
  def clean do
    lib = ExTorch.Vision.Build.lib_path()
    if File.exists?(lib), do: File.rm(lib)
    if File.exists?(manifest_path()), do: File.rm(manifest_path())
    :ok
  end

  defp manifest_path do
    Path.join(Mix.Project.manifest_path(), ".compile.torchvision")
  end
end