firebird

defmodule Mix.Tasks.Firebird.Gen do
  @moduledoc """
  Generate an Elixir module wrapper from a WASM file.

  Inspects the WASM file's exports and generates a typed Elixir module
  with function wrappers, documentation, and supervision support.

  ## Usage

      mix firebird.gen path/to/module.wasm
      mix firebird.gen path/to/module.wasm --module MyApp.Math
      mix firebird.gen path/to/module.wasm --module MyApp.Math --out lib/my_app/math.ex
      mix firebird.gen path/to/module.wasm --style pool

  ## Options

  - `--module` - Module name (default: inferred from filename)
  - `--out` - Output file path (default: inferred from module name)
  - `--style` - Generation style: `module` (default), `pool`, or `basic`
  - `--test` - Also generate a test file with smoke tests
  - `--wasi` - Enable WASI support
  - `--force` - Overwrite existing files
  - `--print` - Print to stdout instead of writing to file
  """

  use Mix.Task

  @shortdoc "Generate Elixir module from a WASM file"

  @impl Mix.Task
  def run(args) do
    {opts, positional, _} =
      OptionParser.parse(args,
        switches: [
          module: :string,
          out: :string,
          style: :string,
          wasi: :boolean,
          force: :boolean,
          print: :boolean,
          test: :boolean
        ],
        aliases: [m: :module, o: :out, s: :style, f: :force, p: :print, t: :test]
      )

    wasm_path =
      case positional do
        [path | _] -> path
        [] -> Mix.raise("Usage: mix firebird.gen <wasm_file> [--module ModuleName]")
      end

    unless File.exists?(wasm_path) do
      Mix.raise(
        "WASM file not found: #{wasm_path}\n\nMake sure the file exists and the path is correct."
      )
    end

    # Start the application so we can use Wasmex
    Mix.Task.run("app.start", ["--no-start"])
    Application.ensure_all_started(:wasmex)

    module_name = opts[:module] || infer_module_name(wasm_path)
    style = String.to_atom(opts[:style] || "module")
    wasi = opts[:wasi] || false

    Mix.shell().info("🔥 Inspecting #{wasm_path}...")

    case Firebird.Inspector.inspect_file(wasm_path, wasi: wasi) do
      {:ok, info} ->
        code = generate_code(module_name, wasm_path, info, style, wasi)

        if opts[:print] do
          Mix.shell().info(code)
        else
          out_path = opts[:out] || infer_output_path(module_name)
          write_module(out_path, code, opts[:force])

          if opts[:test] do
            test_code = generate_test(module_name, wasm_path, info, wasi)
            test_path = infer_test_path(module_name)
            write_module(test_path, test_code, opts[:force])
          end
        end

        print_summary(module_name, info, style)

      {:error, reason} ->
        Mix.raise("Failed to inspect WASM file: #{inspect(reason)}")
    end
  end

  defp infer_module_name(path) do
    path
    |> Path.basename(".wasm")
    |> Macro.camelize()
    |> then(&"MyApp.#{&1}")
  end

  defp infer_output_path(module_name) do
    module_name
    |> String.split(".")
    |> Enum.map(&Macro.underscore/1)
    |> then(fn parts -> Path.join(["lib" | parts]) <> ".ex" end)
  end

  defp write_module(path, code, force) do
    if File.exists?(path) && !force do
      Mix.raise(
        "File already exists: #{path}\n\nUse --force to overwrite, or --out to specify a different path."
      )
    end

    File.mkdir_p!(Path.dirname(path))
    File.write!(path, code)
    Mix.shell().info("  ✅ Created #{path}")
  end

  defp generate_code(module_name, wasm_path, info, :module, wasi) do
    public_fns = Enum.reject(info.functions, &String.starts_with?(&1.name, "_"))
    opts_line = if wasi, do: ",\n    opts: [wasi: true]", else: ""

    fn_lines =
      public_fns
      |> Enum.map(fn f ->
        param_doc =
          if f.params == [], do: "", else: " (#{Enum.join(Enum.map(f.params, &inspect/1), ", ")})"

        doc_line = "  @doc \"Call `#{f.name}`#{param_doc} → #{inspect(f.results)}\""
        "#{doc_line}\n  wasm_fn :#{f.name}, args: #{f.arity}"
      end)
      |> Enum.join("\n\n")

    """
    defmodule #{module_name} do
      @moduledoc \"\"\"
      Auto-generated Elixir wrapper for #{Path.basename(wasm_path)}.

      Generated by `mix firebird.gen` on #{Date.utc_today()}.

      ## Usage

          # Add to your supervision tree:
          children = [#{module_name}]

          # Then call functions:
    #{public_fns |> Enum.take(3) |> Enum.map(fn f -> "      #{module_name}.#{f.name}!(#{Enum.map_join(1..max(f.arity, 1), ", ", fn i -> "arg#{i}" end)})" end) |> Enum.join("\n")}
      \"\"\"

      use Firebird,
        wasm: #{inspect(wasm_path)}#{opts_line}

    #{fn_lines}
    end
    """
  end

  defp generate_code(module_name, wasm_path, info, :pool, wasi) do
    public_fns = Enum.reject(info.functions, &String.starts_with?(&1.name, "_"))

    pool_name =
      module_name |> String.split(".") |> List.last() |> Macro.underscore() |> String.to_atom()

    wasi_opt = if wasi, do: ", wasi: true", else: ""

    fn_lines =
      public_fns
      |> Enum.map(fn f ->
        args = Enum.map_join(1..max(f.arity, 1), ", ", fn i -> "arg#{i}" end)

        arg_vars =
          if f.arity > 0, do: Enum.map_join(1..f.arity, ", ", fn i -> "arg#{i}" end), else: ""

        arg_list = if f.arity > 0, do: "[#{arg_vars}]", else: "[]"

        """
          @doc "Call `#{f.name}` on the pool."
          def #{f.name}(#{args}) do
            Firebird.Pool.call(@pool_name, :#{f.name}, #{arg_list})
          end

          def #{f.name}!(#{args}) do
            Firebird.Pool.call!(@pool_name, :#{f.name}, #{arg_list})
          end
        """
      end)
      |> Enum.join("\n")

    """
    defmodule #{module_name} do
      @moduledoc \"\"\"
      Pool-backed wrapper for #{Path.basename(wasm_path)}.

      Generated by `mix firebird.gen` on #{Date.utc_today()}.

      ## Setup

          # In your supervision tree:
          children = [#{module_name}]

      ## Usage

      #{public_fns |> Enum.take(3) |> Enum.map(fn f -> "    #{module_name}.#{f.name}!(#{Enum.map_join(1..max(f.arity, 1), ", ", fn i -> "arg#{i}" end)})" end) |> Enum.join("\n")}
      \"\"\"

      @pool_name :#{pool_name}_pool
      @wasm_path #{inspect(wasm_path)}

      def child_spec(opts) do
        size = Keyword.get(opts, :size, System.schedulers_online())

        %{
          id: __MODULE__,
          start: {Firebird.Pool, :start_link,
            [wasm: @wasm_path, size: size, name: @pool_name#{wasi_opt}]},
          type: :worker,
          restart: :permanent
        }
      end

      def start_link(opts \\\\ []) do
        size = Keyword.get(opts, :size, System.schedulers_online())
        Firebird.Pool.start_link(wasm: @wasm_path, size: size, name: @pool_name#{wasi_opt})
      end

      def status, do: Firebird.Pool.status(@pool_name)

    #{fn_lines}
    end
    """
  end

  defp generate_code(module_name, wasm_path, info, :basic, wasi) do
    public_fns = Enum.reject(info.functions, &String.starts_with?(&1.name, "_"))
    wasi_opt = if wasi, do: ", wasi: true", else: ""

    fn_lines =
      public_fns
      |> Enum.map(fn f ->
        args =
          if f.arity > 0, do: Enum.map_join(1..f.arity, ", ", fn i -> "arg#{i}" end), else: ""

        arg_list = if f.arity > 0, do: "[#{args}]", else: "[]"

        """
          @doc "Call `#{f.name}` — one-shot (load, call, cleanup)."
          def #{f.name}(#{args}) do
            Firebird.run(@wasm_path, :#{f.name}, #{arg_list}#{wasi_opt})
          end

          def #{f.name}!(#{args}) do
            Firebird.run!(@wasm_path, :#{f.name}, #{arg_list}#{wasi_opt})
          end
        """
      end)
      |> Enum.join("\n")

    """
    defmodule #{module_name} do
      @moduledoc \"\"\"
      Stateless wrapper for #{Path.basename(wasm_path)}.

      Each call loads a fresh WASM instance. Simple but not pooled.
      For high throughput, use `mix firebird.gen --style pool`.

      Generated by `mix firebird.gen` on #{Date.utc_today()}.
      \"\"\"

      @wasm_path #{inspect(wasm_path)}

    #{fn_lines}
    end
    """
  end

  defp generate_code(module_name, wasm_path, info, style, wasi) do
    Mix.shell().info("  ⚠️  Unknown style '#{style}', falling back to :module")
    generate_code(module_name, wasm_path, info, :module, wasi)
  end

  defp infer_test_path(module_name) do
    module_name
    |> String.split(".")
    |> Enum.map(&Macro.underscore/1)
    |> then(fn parts -> Path.join(["test" | parts]) <> "_test.exs" end)
  end

  defp generate_test(module_name, wasm_path, info, wasi) do
    public_fns = Enum.reject(info.functions, &String.starts_with?(&1.name, "_"))
    wasi_opt = if wasi, do: ", wasi: true", else: ""

    test_cases =
      public_fns
      |> Enum.take(5)
      |> Enum.map(fn f ->
        zeros = List.duplicate(0, f.arity)

        """
          test "#{f.name} is callable", %{wasm: wasm} do
            assert {:ok, result} = Firebird.call(wasm, :#{f.name}, #{inspect(zeros)})
            assert is_list(result)
          end
        """
      end)
      |> Enum.join("\n")

    exports_assertions =
      public_fns
      |> Enum.map(fn f -> "    assert :#{f.name} in exports" end)
      |> Enum.join("\n")

    """
    defmodule #{module_name}Test do
      use ExUnit.Case
      import Firebird.TestHelpers

      setup_wasm #{inspect(wasm_path)}#{wasi_opt}

      test "exports expected functions", %{wasm: wasm} do
        exports = Firebird.exports(wasm)
    #{exports_assertions}
      end

    #{test_cases}
    end
    """
  end

  defp print_summary(module_name, info, style) do
    fn_count = length(Enum.reject(info.functions, &String.starts_with?(&1.name, "_")))

    Mix.shell().info("""

    🔥 Generated #{module_name} (#{style} style)

       Functions: #{fn_count}
       WASM size: #{format_bytes(info.size_bytes)}

       Next steps:
       1. Add #{module_name} to your supervision tree
       2. Call functions: #{module_name}.function_name!(args)
    """)
  end

  defp format_bytes(bytes) when bytes < 1024, do: "#{bytes} B"
  defp format_bytes(bytes) when bytes < 1_048_576, do: "#{Float.round(bytes / 1024, 1)} KB"
  defp format_bytes(bytes), do: "#{Float.round(bytes / 1_048_576, 1)} MB"
end