firebird

defmodule Mix.Compilers.FirebirdWasm do
  @moduledoc """
  Mix compiler integration for Firebird WASM target.

  This compiler scans for Elixir source files containing `@wasm true`
  annotations and compiles them to WebAssembly. It integrates with the
  Mix build system for incremental compilation.

  ## Usage in mix.exs

      def project do
        [
          compilers: Mix.compilers() ++ [:firebird_wasm],
          firebird: [
            wasm_output: "_build/wasm",
            wasm_sources: ["lib/wasm_modules"],
            wat_only: false
          ]
        ]
      end
  """

  @manifest_file "firebird_wasm.manifest"

  @doc """
  Run the Firebird WASM compiler.

  Called by Mix as part of the compilation pipeline when `:firebird_wasm`
  is in the project's compilers list.
  """
  @spec run(keyword()) :: :ok | :noop | {:error, term()}
  def run(_args \\ []) do
    project = Mix.Project.config()
    firebird_opts = Keyword.get(project, :firebird, [])

    output_dir = Keyword.get(firebird_opts, :wasm_output, "_build/wasm")
    source_dirs = Keyword.get(firebird_opts, :wasm_sources, ["lib"])
    wat_only = Keyword.get(firebird_opts, :wat_only, false)

    # Check for wat2wasm unless wat_only
    unless wat_only do
      unless Firebird.Compiler.wat2wasm_available?() do
        Mix.shell().error("wat2wasm not found. Install wabt or use wat_only: true")
        {:error, :wat2wasm_not_found}
      end
    end

    manifest_path = manifest_path()
    old_manifest = read_manifest(manifest_path)

    # Find all source files
    sources = find_wasm_sources(source_dirs)

    # Find files that need recompilation
    stale = stale_sources(sources, old_manifest)

    if Enum.empty?(stale) do
      :noop
    else
      File.mkdir_p!(output_dir)

      Mix.shell().info("Firebird WASM: compiling #{length(stale)} file(s)")

      results =
        Enum.map(stale, fn source_path ->
          compile_source(source_path, output_dir, wat_only)
        end)

      successes = Enum.filter(results, &match?({:ok, _}, &1))
      failures = Enum.filter(results, &match?({:error, _}, &1))

      # Update manifest with successful compilations
      new_entries = Enum.map(successes, fn {:ok, entry} -> entry end)

      # Keep non-stale entries from old manifest
      stale_paths = MapSet.new(stale)
      kept = Enum.reject(old_manifest, fn entry -> entry.source in stale_paths end)

      write_manifest(manifest_path, kept ++ new_entries)

      for {:ok, entry} <- successes do
        Mix.shell().info("  ✅ #{Path.basename(entry.source)} → #{entry.module}")
      end

      for {:error, {path, reason}} <- failures do
        Mix.shell().error("  ❌ #{Path.basename(path)}: #{inspect(reason)}")
      end

      if Enum.any?(failures) do
        {:error, Enum.map(failures, fn {:error, r} -> r end)}
      else
        :ok
      end
    end
  end

  @doc """
  Clean compiled WASM artifacts.
  """
  @spec clean() :: :ok
  def clean do
    manifest_path = manifest_path()
    manifest = read_manifest(manifest_path)

    for entry <- manifest do
      if entry.wat_path, do: File.rm(entry.wat_path)
      if entry.wasm_path, do: File.rm(entry.wasm_path)
    end

    File.rm(manifest_path)
    :ok
  end

  @doc """
  Return the path to the manifest file.
  """
  @spec manifest_path() :: String.t()
  def manifest_path do
    Path.join(Mix.Project.manifest_path(), @manifest_file)
  end

  @doc """
  Find all Elixir source files that contain WASM annotations.
  """
  @spec find_wasm_sources([String.t()]) :: [String.t()]
  def find_wasm_sources(dirs) do
    dirs
    |> Enum.flat_map(fn dir ->
      if File.dir?(dir) do
        Path.wildcard(Path.join(dir, "**/*.{ex,exs}"))
      else
        if File.regular?(dir), do: [dir], else: []
      end
    end)
    |> Enum.filter(&has_wasm_annotation?/1)
    |> Enum.sort()
    |> Enum.uniq()
  end

  @doc """
  Check if a source file contains @wasm annotations.
  """
  @spec has_wasm_annotation?(String.t()) :: boolean()
  def has_wasm_annotation?(path) do
    case File.read(path) do
      {:ok, content} -> String.contains?(content, "@wasm true")
      _ -> false
    end
  end

  # Find sources that need recompilation based on content hash + mtime
  defp stale_sources(sources, manifest) do
    manifest_map =
      Map.new(manifest, fn entry ->
        {entry.source,
         %{compiled_at: entry.compiled_at, content_hash: Map.get(entry, :content_hash)}}
      end)

    Enum.filter(sources, fn source ->
      case Map.get(manifest_map, source) do
        nil ->
          true

        %{content_hash: cached_hash} when is_binary(cached_hash) ->
          # Use content hash if available (more accurate than mtime)
          current_hash = Firebird.Target.ContentHash.hash_file(source)
          current_hash != cached_hash

        %{compiled_at: compiled_at} when not is_nil(compiled_at) ->
          # Fall back to mtime comparison
          case File.stat(source) do
            {:ok, %{mtime: mtime}} ->
              mtime_seconds = :calendar.datetime_to_gregorian_seconds(mtime)
              mtime_seconds > compiled_at

            _ ->
              true
          end

        _ ->
          true
      end
    end)
  end

  defp compile_source(source_path, output_dir, wat_only) do
    case Firebird.Compiler.compile(source_path, output_dir: output_dir, wat_only: wat_only) do
      {:ok, result} ->
        base = Atom.to_string(result.module) |> String.replace(".", "_")

        entry = %{
          source: source_path,
          module: result.module,
          wat_path: Path.join(output_dir, "#{base}.wat"),
          wasm_path: if(result.wasm, do: Path.join(output_dir, "#{base}.wasm"), else: nil),
          compiled_at: :calendar.datetime_to_gregorian_seconds(:calendar.local_time()),
          wat_size: byte_size(result.wat),
          wasm_size: if(result.wasm, do: byte_size(result.wasm), else: 0),
          content_hash: Firebird.Target.ContentHash.hash_file(source_path)
        }

        {:ok, entry}

      {:error, reason} ->
        {:error, {source_path, reason}}
    end
  end

  defp read_manifest(path) do
    case File.read(path) do
      {:ok, content} ->
        content
        |> :erlang.binary_to_term()
        |> case do
          entries when is_list(entries) -> entries
          _ -> []
        end

      _ ->
        []
    end
  rescue
    _ -> []
  end

  defp write_manifest(path, entries) do
    File.mkdir_p!(Path.dirname(path))
    File.write!(path, :erlang.term_to_binary(entries))
  end
end