firebird

defmodule Mix.Tasks.Firebird.Target.Link do
  @moduledoc """
  Link multiple WASM modules into a single combined binary.

  ## Usage

      # Link all @wasm modules in lib/
      mix firebird.target.link

      # Link specific files
      mix firebird.target.link --files lib/math.ex,lib/utils.ex

      # Link with module name
      mix firebird.target.link --name MyApp

      # Prefix exports with module name to avoid collisions
      mix firebird.target.link --prefix

      # Output to custom directory
      mix firebird.target.link --output _build/wasm

  ## Export Naming

  By default, all public functions are exported with their original names.
  If two modules have functions with the same name, use `--prefix` to
  namespace them:

      # Without prefix: add, multiply
      # With prefix: Math_add, Utils_multiply
  """

  use Mix.Task

  @shortdoc "Link multiple Elixir modules into a single WASM binary"

  @switches [
    output: :string,
    files: :string,
    dirs: :string,
    name: :string,
    prefix: :boolean,
    wat_only: :boolean,
    verbose: :boolean,
    optimize: :boolean,
    tco: :boolean,
    inline: :boolean
  ]

  @aliases [
    o: :output,
    f: :files,
    d: :dirs,
    n: :name,
    v: :verbose
  ]

  @impl Mix.Task
  def run(args) do
    {opts, _rest, _invalid} = OptionParser.parse(args, switches: @switches, aliases: @aliases)

    link(opts)
  end

  @doc """
  Link multiple modules into a single WASM binary.
  """
  @spec link(keyword()) :: :ok | {:error, term()}
  def link(opts) do
    output_dir = Keyword.get(opts, :output, "_build/wasm")
    name = Keyword.get(opts, :name, "Combined")
    prefix = Keyword.get(opts, :prefix, false)
    wat_only = Keyword.get(opts, :wat_only, false)
    verbose = Keyword.get(opts, :verbose, false)

    # Check prerequisites
    unless wat_only do
      unless Firebird.Compiler.wat2wasm_available?() do
        Mix.raise("wat2wasm not found! Install wabt or use --wat-only")
      end
    end

    # Discover sources
    sources = discover_sources(opts)

    if Enum.empty?(sources) do
      Mix.shell().info("🔗 No WASM sources found to link")
      :ok
    else
      Mix.shell().info("🔗 Firebird WASM Linker")
      Mix.shell().info("   Sources: #{length(sources)} file(s)")
      Mix.shell().info("   Output:  #{output_dir}/#{name}.wasm")
      Mix.shell().info("   Prefix:  #{prefix}")
      Mix.shell().info("")

      combiner_opts = [
        name: name,
        prefix: prefix,
        output_dir: output_dir,
        wat_only: wat_only,
        optimize: Keyword.get(opts, :optimize, false),
        tco: Keyword.get(opts, :tco, false),
        inline: Keyword.get(opts, :inline, false)
      ]

      case Firebird.Compiler.Combiner.combine(sources, combiner_opts) do
        {:ok, result} ->
          exports = result[:exports] || []

          if verbose do
            Mix.shell().info("   Exports:")
            Enum.each(exports, fn e -> Mix.shell().info("     • #{e}") end)
            Mix.shell().info("")
          end

          Mix.shell().info("   ✅ Linked #{length(exports)} export(s) into #{name}")

          if result.wat do
            Mix.shell().info("   📄 WAT: #{byte_size(result.wat)} bytes")
          end

          if result.wasm do
            Mix.shell().info("   📦 WASM: #{byte_size(result.wasm)} bytes")
          end

          :ok

        {:error, reason} ->
          Mix.shell().error("   ❌ Link failed: #{inspect(reason)}")
          {:error, reason}
      end
    end
  end

  defp discover_sources(opts) do
    files =
      case Keyword.get(opts, :files) do
        nil -> []
        files_str -> String.split(files_str, ",", trim: true)
      end

    dirs =
      case Keyword.get(opts, :dirs) do
        nil ->
          if Enum.empty?(files) do
            ["lib"]
          else
            []
          end

        dirs_str ->
          String.split(dirs_str, ",", trim: true)
      end

    explicit_files = Enum.filter(files, &File.regular?/1)

    discovered =
      dirs
      |> Enum.flat_map(fn dir ->
        if File.dir?(dir) do
          Path.wildcard(Path.join(dir, "**/*.{ex,exs}"))
        else
          []
        end
      end)
      |> Enum.filter(&has_wasm_content?/1)

    (explicit_files ++ discovered) |> Enum.sort() |> Enum.uniq()
  end

  defp has_wasm_content?(path) do
    case File.read(path) do
      {:ok, content} ->
        String.contains?(content, "@wasm true") or String.contains?(path, "wasm_modules")

      _ ->
        false
    end
  end
end