firebird

defmodule Mix.Tasks.Firebird.Compile do
  @moduledoc """
  Compiles Elixir source files to WebAssembly.

  This task takes Elixir source files containing `@wasm true` annotated
  functions and compiles them to WASM binary format.

  ## Usage

      # Compile a single file
      mix firebird.compile lib/math.ex

      # Compile multiple files
      mix firebird.compile lib/math.ex lib/crypto.ex

      # Compile all files in a directory
      mix firebird.compile lib/wasm_modules/

      # Specify output directory
      mix firebird.compile lib/math.ex --output _build/wasm

      # Generate WAT only (no binary)
      mix firebird.compile lib/math.ex --wat-only

      # Show verbose compilation output
      mix firebird.compile lib/math.ex --verbose

  ## Options

  - `--output`, `-o` - Output directory (default: `_build/wasm`)
  - `--wat-only` - Only generate WAT text format, skip binary
  - `--verbose`, `-v` - Show detailed compilation output
  - `--verify` - Load and verify the compiled WASM module

  ## Module Annotations

  Mark functions for WASM compilation with `@wasm true`:

      defmodule MyMath do
        @wasm true
        def add(a, b), do: a + b

        @wasm true
        def fibonacci(0), do: 0
        def fibonacci(1), do: 1
        def fibonacci(n), do: fibonacci(n - 1) + fibonacci(n - 2)
      end

  If no `@wasm` annotations are present, all public functions are compiled.

  ## Supported Elixir Subset

  - Integer arithmetic (+, -, *, div, rem)
  - Comparison operators (==, !=, <, >, <=, >=)
  - Boolean logic (and, or, not)
  - if/else, cond, case expressions
  - Pattern matching on integer literals
  - Multi-clause function definitions
  - Recursive functions
  - Guard clauses

  ## Limitations

  - Only numeric types (i64 integers, f64 floats)
  - No lists, maps, tuples, strings, or atoms
  - No IO, processes, or OTP features
  - No standard library functions
  """

  use Mix.Task

  @shortdoc "Compile Elixir source files to WebAssembly"

  @switches [
    output: :string,
    wat_only: :boolean,
    verbose: :boolean,
    verify: :boolean
  ]

  @aliases [
    o: :output,
    v: :verbose
  ]

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

    output_dir = Keyword.get(opts, :output, "_build/wasm")
    verbose = Keyword.get(opts, :verbose, false)
    wat_only = Keyword.get(opts, :wat_only, false)
    verify = Keyword.get(opts, :verify, false)

    # Check prerequisites
    unless wat_only do
      unless Firebird.Compiler.wat2wasm_available?() do
        Mix.raise("""
        wat2wasm not found! Install the WebAssembly Binary Toolkit (wabt):

          # macOS
          brew install wabt

          # Ubuntu/Debian
          apt-get install wabt

          # Or download from https://github.com/WebAssembly/wabt
        """)
      end
    end

    if Enum.empty?(paths) do
      Mix.raise("No input files specified. Usage: mix firebird.compile <file_or_dir> [...]")
    end

    # Resolve all input files
    files = resolve_files(paths)

    if Enum.empty?(files) do
      Mix.raise("No .ex or .exs files found in the specified paths")
    end

    Mix.shell().info("🔥 Firebird WASM Compiler")
    Mix.shell().info("   Compiling #{length(files)} file(s) → #{output_dir}/")
    Mix.shell().info("")

    File.mkdir_p!(output_dir)

    # Compile each file
    results =
      Enum.map(files, fn file ->
        compile_file(file, output_dir, wat_only, verbose, verify)
      end)

    # Summary
    successes = Enum.count(results, fn {status, _} -> status == :ok end)
    failures = Enum.count(results, fn {status, _} -> status == :error end)

    Mix.shell().info("")
    Mix.shell().info("📊 Results: #{successes} succeeded, #{failures} failed")

    if failures > 0 do
      Mix.raise("Compilation failed for #{failures} file(s)")
    end

    :ok
  end

  defp resolve_files(paths) do
    Enum.flat_map(paths, fn path ->
      cond do
        File.dir?(path) ->
          Path.wildcard(Path.join(path, "**/*.{ex,exs}"))

        File.regular?(path) ->
          [path]

        true ->
          Mix.shell().error("  ⚠ File not found: #{path}")
          []
      end
    end)
    |> Enum.sort()
    |> Enum.uniq()
  end

  defp compile_file(file, output_dir, wat_only, verbose, verify) do
    Mix.shell().info("  📄 #{file}")

    case Firebird.Compiler.compile(file, output_dir: output_dir, wat_only: wat_only) do
      {:ok, result} ->
        module_name = result.module

        if verbose do
          Mix.shell().info("     Module: #{module_name}")
          Mix.shell().info("     WAT size: #{byte_size(result.wat)} bytes")

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

        if verify and result.wasm do
          verify_wasm(result.wasm, module_name)
        end

        Mix.shell().info("     ✅ Compiled #{module_name}")
        {:ok, result}

      {:error, reason} ->
        Mix.shell().error("     ❌ Failed: #{format_error(reason)}")
        {:error, reason}
    end
  end

  defp verify_wasm(wasm_binary, _module_name) do
    case Firebird.load(wasm_binary) do
      {:ok, instance} ->
        exports = Firebird.exports(instance)
        Mix.shell().info("     🔍 Verified: #{length(exports)} export(s): #{inspect(exports)}")
        Firebird.stop(instance)

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

  defp format_error({:parse_error, {line, col}, msg, token}) do
    "Parse error at #{line}:#{col}: #{msg}#{token}"
  end

  defp format_error({:validation_errors, errors}) do
    details =
      Enum.map(errors, fn {func, errs} ->
        "  #{func}: #{Enum.join(errs, ", ")}"
      end)

    "Validation errors:\n#{Enum.join(details, "\n")}"
  end

  defp format_error({:wat2wasm_error, output}) do
    "wat2wasm error: #{output}"
  end

  defp format_error({:file_error, reason, path}) do
    "Cannot read #{path}: #{reason}"
  end

  defp format_error(other) do
    inspect(other)
  end
end