defmodule Mix.Tasks.Compile.Wasm do
  @moduledoc """
  Custom Mix compiler for automatic WASM compilation.

  Integrates Firebird's Elixir-to-WASM compiler into the standard
  `mix compile` pipeline. When added to your project's compilers list,
  it automatically compiles `@wasm true` annotated modules to WASM
  as part of the normal build process.

  ## Setup

  Add `:wasm` to your project's compilers in `mix.exs`:

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

  Or conditionally, only when targeting WASM:

      def project do
        [
          compilers: Mix.compilers() ++ wasm_compilers(),
          ...
        ]
      end

      defp wasm_compilers do
        if System.get_env("MIX_TARGET") == "wasm", do: [:wasm], else: []
      end

  ## Configuration

  All configuration is read from the `:firebird` key in `mix.exs`:

  - `:wasm_output` - Output directory (default: `"_build/wasm"`)
  - `:wasm_sources` - Source directories to scan (default: `["lib"]`)
  - `:optimize` - Enable IR optimizations (default: `false`)
  - `:tco` - Enable tail call optimization (default: `false`)
  - `:inline` - Enable function inlining (default: `false`)
  - `:wat_only` - Only generate WAT, skip binary (default: `false`)
  - `:verify` - Verify compiled modules (default: `false`)

  ## Incremental Compilation

  The compiler tracks source file content hashes and only recompiles
  files whose contents have actually changed. Touching a file without
  changing its content will NOT trigger recompilation.

  ## Environment Variables

  - `MIX_TARGET=wasm` - Commonly used to indicate WASM build
  - `FIREBIRD_OPTIMIZE=1` - Override optimize setting
  - `FIREBIRD_TCO=1` - Override TCO setting
  - `FIREBIRD_WAT_ONLY=1` - Override WAT-only setting
  - `FIREBIRD_VERBOSE=1` - Enable verbose output
  """

  use Mix.Task.Compiler

  @manifest_file "compile.wasm.manifest"

  @impl Mix.Task.Compiler
  def run(_argv) do
    config = Firebird.Target.MixTarget.config_from_env()

    sources = Firebird.Target.discover_sources(config)

    if Enum.empty?(sources) do
      {:noop, []}
    else
      manifest = read_manifest()
      stale = find_stale_sources(sources, manifest)

      if Enum.empty?(stale) do
        {:noop, []}
      else
        File.mkdir_p!(config.output_dir)

        if config.verbose do
          Mix.shell().info("Compiling #{length(stale)} WASM file(s)")
        end

        compiler_opts = [
          output_dir: config.output_dir,
          optimize: config.optimize,
          tco: config.tco,
          inline: config.inline,
          wat_only: config.wat_only,
          source_map: config.source_map
        ]

        {results, diagnostics} = compile_sources(stale, compiler_opts, config.verbose)

        # Update manifest
        new_entries = build_manifest_entries(results)
        stale_set = MapSet.new(stale)
        kept = Enum.reject(manifest, fn entry -> entry.source in stale_set end)
        write_manifest(kept ++ new_entries)

        if Enum.any?(diagnostics, &(&1.severity == :error)) do
          {:error, diagnostics}
        else
          {:ok, diagnostics}
        end
      end
    end
  end

  @impl Mix.Task.Compiler
  def clean do
    manifest = read_manifest()

    Enum.each(manifest, fn entry ->
      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

  @impl Mix.Task.Compiler
  def manifests do
    [manifest_path()]
  end

  # Private

  defp compile_sources(sources, compiler_opts, verbose) do
    results =
      Enum.map(sources, fn source ->
        case Firebird.Compiler.compile(source, compiler_opts) do
          {:ok, result} ->
            if verbose do
              Mix.shell().info("  Compiled #{Path.relative_to_cwd(source)} → #{result.module}")
            end

            {:ok, Map.put(result, :source_file, source)}

          {:error, reason} ->
            {:error, {source, reason}}
        end
      end)

    diagnostics =
      Enum.flat_map(results, fn
        {:ok, result} ->
          [
            %Mix.Task.Compiler.Diagnostic{
              compiler_name: "wasm",
              file: Map.get(result, :source_file, "unknown"),
              message: "Compiled #{result.module} to WASM",
              position: 0,
              severity: :information
            }
          ]

        {:error, {source, reason}} ->
          [
            %Mix.Task.Compiler.Diagnostic{
              compiler_name: "wasm",
              file: source,
              message: format_error(reason),
              position: error_position(reason),
              severity: :error
            }
          ]
      end)

    successes = Enum.filter(results, &match?({:ok, _}, &1)) |> Enum.map(fn {:ok, r} -> r end)
    {successes, diagnostics}
  end

  defp find_stale_sources(sources, manifest) do
    manifest_map = Map.new(manifest, fn entry -> {entry.source, entry[:content_hash]} end)

    Enum.filter(sources, fn source ->
      cached_hash = Map.get(manifest_map, source)
      current_hash = Firebird.Target.ContentHash.hash_file(source)
      current_hash != cached_hash
    end)
  end

  defp build_manifest_entries(results) do
    Enum.map(results, fn result ->
      source = Map.get(result, :source_file, "unknown")
      base = Atom.to_string(result.module) |> String.replace(".", "_")
      output_dir = Path.dirname(Map.get(result, :wat_path, "_build/wasm/#{base}.wat"))

      %{
        source: source,
        module: result.module,
        wat_path: Path.join(output_dir, "#{base}.wat"),
        wasm_path: if(result.wasm, do: Path.join(output_dir, "#{base}.wasm")),
        compiled_at: :calendar.datetime_to_gregorian_seconds(:calendar.local_time()),
        content_hash: Firebird.Target.ContentHash.hash_file(source)
      }
    end)
  end

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

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

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

  defp manifest_path do
    Path.join(Mix.Project.manifest_path(), @manifest_file)
  rescue
    _ -> Path.join("_build", @manifest_file)
  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
    Enum.map(errors, fn {func, errs} -> "#{func}: #{Enum.join(errs, ", ")}" end)
    |> Enum.join("; ")
  end

  defp format_error({:wat2wasm_error, output}), do: "wat2wasm: #{output}"
  defp format_error({:file_error, reason, path}), do: "Cannot read #{path}: #{reason}"
  defp format_error(other), do: inspect(other)

  defp error_position({:parse_error, {line, _col}, _msg, _token}), do: line
  defp error_position(_), do: 0
end