dux

defmodule Dux.Remote.Coordinator do
  @moduledoc """
  Orchestrates distributed query execution across a cluster of Dux workers.

  The coordinator:
  1. Discovers available workers via `:pg`
  2. Partitions the data source across workers
  3. Fans out the pipeline to each worker via `Worker.execute/2`
  4. Merges partial results on the coordinator node

  ## Usage

      # Execute a pipeline across all available workers
      result = Dux.Remote.Coordinator.execute(pipeline)

      # Execute with specific workers
      result = Dux.Remote.Coordinator.execute(pipeline, workers: [w1, w2, w3])

  The result is a `%Dux{}` struct with the merged data.
  """

  alias Dux.Remote.{Merger, Partitioner, PipelineSplitter, Shuffle, Worker}
  import Dux.SQL.Helpers, only: [qi: 1]

  # Broadcast threshold: 256MB serialized Arrow IPC
  @broadcast_threshold 256 * 1024 * 1024

  @doc """
  Execute a `%Dux{}` pipeline across distributed workers.

  Partitions the source, fans out to workers, collects Arrow IPC results,
  and merges on the coordinator node.

  ## Options

    * `:workers` — list of worker PIDs (default: all workers from `:pg`)
    * `:timeout` — per-worker timeout in ms (default: `:infinity`)
    * `:strategy` — partitioning strategy (default: `:round_robin`)

  Returns a `%Dux{}` struct with the merged results.
  """
  def execute(%Dux{} = pipeline, opts \\ []) do
    workers = Keyword.get_lazy(opts, :workers, &Worker.list/0)
    timeout = Keyword.get(opts, :timeout, :infinity)
    strategy = Keyword.get(opts, :strategy, :round_robin)
    bcast_threshold = Keyword.get(opts, :broadcast_threshold, @broadcast_threshold)

    if workers == [] do
      raise ArgumentError, "no workers available for distributed execution"
    end

    # Split pipeline: worker ops push down, coordinator ops apply post-merge
    %{worker_ops: worker_ops, coordinator_ops: coord_ops, agg_rewrites: rewrites} =
      PipelineSplitter.split(pipeline.ops)

    # Preprocess joins: broadcast/shuffle right sides that aren't worker-safe
    case preprocess_joins(worker_ops, workers, timeout, bcast_threshold) do
      {:ok, processed_ops, broadcast_tables} ->
        # All joins handled inline (broadcast or push-down)
        worker_pipeline = %{pipeline | ops: processed_ops}

        try do
          result = execute_fan_out(worker_pipeline, workers, strategy, timeout)
          result = apply_avg_rewrites(result, rewrites)
          apply_coordinator_ops(result, coord_ops)
        after
          cleanup_broadcast_tables(workers, broadcast_tables)
        end

      {:shuffle, ops_before, {right_computed, how, on_cols, suffix}, ops_after, broadcast_tables} ->
        # Pipeline needs a shuffle stage: execute pre-join → shuffle → post-join
        try do
          execute_with_shuffle(
            pipeline, ops_before, right_computed, how, on_cols, suffix,
            ops_after, coord_ops, rewrites, workers, strategy, timeout
          )
        after
          cleanup_broadcast_tables(workers, broadcast_tables)
        end
    end
  end

  @doc """
  Execute a pipeline across workers, returning raw Arrow IPC binaries.

  Useful when you want to handle merging yourself or stream results.
  """
  def fan_out_raw(%Dux{} = pipeline, opts \\ []) do
    workers = Keyword.get_lazy(opts, :workers, &Worker.list/0)
    timeout = Keyword.get(opts, :timeout, :infinity)

    assignments = Partitioner.assign(pipeline, workers)
    fan_out(assignments, timeout)
  end

  # ---------------------------------------------------------------------------
  # Internal
  # ---------------------------------------------------------------------------

  # Standard distributed execution: partition → fan out → merge
  defp execute_fan_out(worker_pipeline, workers, strategy, timeout) do
    assignments = Partitioner.assign(worker_pipeline, workers, strategy: strategy)
    results = fan_out(assignments, timeout)
    {successes, failures} = partition_results(results)

    if successes == [] do
      reasons = Enum.map(failures, fn {:error, reason} -> reason end)
      raise ArgumentError, "all workers failed: #{inspect(reasons)}"
    end

    Merger.merge_to_dux(successes, worker_pipeline)
  end

  # Multi-stage execution for shuffle joins:
  # 1. Execute ops before the join as a distributed query
  # 2. Shuffle join the result with the right side
  # 3. Apply remaining ops + coordinator ops
  defp execute_with_shuffle(
         pipeline, ops_before, right_computed, how, on_cols, _suffix,
         ops_after, coord_ops, rewrites, workers, strategy, timeout
       ) do
    # Stage 1: execute pre-join ops distributed
    left_result =
      if ops_before == [] do
        # No ops before join — just compute the source
        Dux.compute(%{pipeline | ops: []})
      else
        pre_pipeline = %{pipeline | ops: ops_before}
        execute_fan_out(pre_pipeline, workers, strategy, timeout)
      end

    # Stage 2: shuffle join
    # Convert on_cols [{left, right}, ...] to the format Shuffle expects
    shuffle_on =
      case on_cols do
        [{col, col}] -> String.to_atom(col)
        pairs -> Enum.map(pairs, fn {l, r} -> {String.to_atom(l), String.to_atom(r)} end)
      end

    shuffle_result =
      Shuffle.execute(left_result, right_computed,
        on: shuffle_on,
        how: how,
        workers: workers,
        timeout: timeout
      )

    # Stage 3: apply remaining ops + rewrites + coordinator ops
    # Any ops after the join in the worker list need to run on the shuffle result
    result =
      if ops_after == [] do
        shuffle_result
      else
        post_pipeline = %{shuffle_result | ops: ops_after}
        Dux.compute(post_pipeline)
      end

    result = apply_avg_rewrites(result, rewrites)
    apply_coordinator_ops(result, coord_ops)
  end

  defp fan_out(assignments, timeout) do
    # Execute in parallel via Task.async_stream
    assignments
    |> Task.async_stream(
      fn {worker, partition_pipeline} ->
        Worker.execute(worker, partition_pipeline, timeout)
      end,
      timeout: timeout,
      max_concurrency: length(assignments),
      ordered: true
    )
    |> Enum.map(fn
      {:ok, {:ok, ipc}} -> {:ok, ipc}
      {:ok, {:error, reason}} -> {:error, reason}
      {:exit, reason} -> {:error, {:worker_crash, reason}}
    end)
  end

  defp partition_results(results) do
    Enum.split_with(results, fn
      {:ok, _} -> true
      _ -> false
    end)
    |> then(fn {ok, err} ->
      {Enum.map(ok, fn {:ok, ipc} -> ipc end), err}
    end)
  end

  # ---------------------------------------------------------------------------
  # Join preprocessing — broadcast/shuffle non-worker-safe right sides
  # ---------------------------------------------------------------------------

  # Scan worker ops for joins where the right side holds a source that
  # workers can't resolve (e.g. a local {:table, ref}).  For each such join:
  #
  # - Small right side → broadcast to workers, replace join with broadcast ref
  # - Large right side → signal a pipeline split for shuffle join
  #
  # Returns:
  #   {:ok, processed_ops, broadcast_names} — all joins handled inline
  #   {:shuffle, ops_before, join_info, ops_after, broadcast_names} — need a shuffle stage
  defp preprocess_joins(ops, workers, timeout, threshold) do
    do_preprocess(ops, [], [], workers, timeout, threshold)
  end

  # Base case: all ops processed, no shuffle needed
  defp do_preprocess([], processed, broadcast_names, _workers, _timeout, _threshold) do
    {:ok, Enum.reverse(processed), broadcast_names}
  end

  # Join op: check if right side is worker-safe
  defp do_preprocess(
         [{:join, %Dux{} = right, how, on_cols, suffix} = op | rest],
         processed,
         broadcast_names,
         workers,
         timeout,
         threshold
       ) do
    if worker_safe_source?(right.source) and Enum.all?(right.ops, &worker_safe_op?/1) do
      # Worker-safe — pass through
      do_preprocess(rest, [op | processed], broadcast_names, workers, timeout, threshold)
    else
      # Non-worker-safe — compute right side and decide broadcast vs shuffle
      right_computed = Dux.compute(right)
      {:table, right_ref} = right_computed.source
      right_ipc = Dux.Native.table_to_ipc(right_ref)

      if byte_size(right_ipc) <= threshold do
        # Small → broadcast
        broadcast_name = "__bcast_#{:erlang.unique_integer([:positive])}"
        broadcast_to_workers(workers, broadcast_name, right_ipc, timeout)
        broadcast_right = Dux.from_query("SELECT * FROM #{qi(broadcast_name)}")
        new_op = {:join, broadcast_right, how, on_cols, suffix}

        do_preprocess(
          rest,
          [new_op | processed],
          [broadcast_name | broadcast_names],
          workers,
          timeout,
          threshold
        )
      else
        # Large → shuffle. Split the pipeline here.
        {:shuffle, Enum.reverse(processed),
         {right_computed, how, on_cols, suffix}, rest, broadcast_names}
      end
    end
  end

  # Non-join op: pass through
  defp do_preprocess([op | rest], processed, broadcast_names, workers, timeout, threshold) do
    do_preprocess(rest, [op | processed], broadcast_names, workers, timeout, threshold)
  end

  defp broadcast_to_workers(workers, name, ipc_binary, _timeout) do
    tasks =
      Enum.map(workers, fn worker ->
        Task.async(fn -> Worker.register_table(worker, name, ipc_binary) end)
      end)

    Task.await_many(tasks, 30_000)
  end

  defp cleanup_broadcast_tables(_workers, []), do: :ok

  defp cleanup_broadcast_tables(workers, broadcast_names) do
    Enum.each(broadcast_names, fn name ->
      Enum.each(workers, fn worker ->
        try do
          Worker.drop_table(worker, name)
        catch
          _, _ -> :ok
        end
      end)
    end)
  end

  # A source is worker-safe if every worker can independently resolve it.
  # File paths, SQL queries, and inline lists are worker-safe.
  # Table refs ({:table, ResourceArc}) are node-local — only the coordinator has them.
  defp worker_safe_source?({:parquet, _}), do: true
  defp worker_safe_source?({:csv, _, _}), do: true
  defp worker_safe_source?({:ndjson, _, _}), do: true
  defp worker_safe_source?({:sql, _}), do: true
  defp worker_safe_source?({:query, _}), do: true
  defp worker_safe_source?({:list, _}), do: true
  defp worker_safe_source?({:table, _}), do: false
  defp worker_safe_source?(_), do: false

  # Check if an op itself is worker-safe (no nested non-worker-safe sources).
  # Joins can nest other Dux structs in the right side.
  defp worker_safe_op?({:join, %Dux{} = right, _, _, _}) do
    worker_safe_source?(right.source) and Enum.all?(right.ops, &worker_safe_op?/1)
  end

  defp worker_safe_op?({:concat_rows, others}) do
    Enum.all?(others, fn %Dux{} = other ->
      worker_safe_source?(other.source) and Enum.all?(other.ops, &worker_safe_op?/1)
    end)
  end

  defp worker_safe_op?(_), do: true

  # ---------------------------------------------------------------------------

  # Apply aggregate rewrites (AVG, STDDEV, VARIANCE, COUNT DISTINCT)
  defp apply_avg_rewrites(dux, rewrites) when map_size(rewrites) == 0, do: dux

  defp apply_avg_rewrites(dux, rewrites) do
    computed_exprs =
      Enum.map(rewrites, fn
        {name, {:avg, sum_col, count_col}} ->
          {name, "#{qi(sum_col)} / #{qi(count_col)}"}

        {name, {:stddev, :stddev_samp, n, sum_x, sum_x2}} ->
          {name, stddev_formula(n, sum_x, sum_x2, :samp, true)}

        {name, {:stddev, :stddev_pop, n, sum_x, sum_x2}} ->
          {name, stddev_formula(n, sum_x, sum_x2, :pop, true)}

        {name, {:stddev, :var_samp, n, sum_x, sum_x2}} ->
          {name, stddev_formula(n, sum_x, sum_x2, :samp, false)}

        {name, {:stddev, :var_pop, n, sum_x, sum_x2}} ->
          {name, stddev_formula(n, sum_x, sum_x2, :pop, false)}

        {name, {:count_distinct, cd_col, _inner}} ->
          # ARRAY_AGG(DISTINCT) across workers → flatten and count distinct
          {name, "list_distinct(flatten(list(#{qi(cd_col)}))).__len"}
      end)

    intermediate_cols =
      Enum.flat_map(rewrites, fn
        {_, {:avg, sum_col, count_col}} ->
          [String.to_atom(sum_col), String.to_atom(count_col)]

        {_, {:stddev, _, n, sum_x, sum_x2}} ->
          [String.to_atom(n), String.to_atom(sum_x), String.to_atom(sum_x2)]

        {_, {:count_distinct, cd_col, _}} ->
          [String.to_atom(cd_col)]
      end)

    dux
    |> Dux.mutate_with(computed_exprs)
    |> Dux.discard(intermediate_cols)
  end

  defp stddev_formula(n_col, sum_col, sum2_col, pop_or_samp, sqrt?) do
    n = qi(n_col)
    sx = qi(sum_col)
    sx2 = qi(sum2_col)

    divisor = if pop_or_samp == :samp, do: "(#{n} - 1)", else: n

    variance = "GREATEST(0, (#{sx2} - #{sx} * #{sx} / #{n}::DOUBLE) / #{divisor}::DOUBLE)"

    if sqrt? do
      "CASE WHEN #{n} <= 1 AND '#{pop_or_samp}' = 'samp' THEN NULL ELSE SQRT(#{variance}) END"
    else
      "CASE WHEN #{n} <= 1 AND '#{pop_or_samp}' = 'samp' THEN NULL ELSE #{variance} END"
    end
  end

  # Apply coordinator-only ops to the merged result and compute locally.
  # The result must be fully materialized — if we return a %Dux{} with pending ops
  # and workers still set, compute() would re-distribute those ops.
  defp apply_coordinator_ops(dux, []), do: dux

  defp apply_coordinator_ops(dux, coord_ops) do
    pipeline = Enum.reduce(coord_ops, dux, &apply_single_op(&2, &1))

    # Compute locally to materialize the result — strip workers to prevent
    # re-distribution of coordinator-only ops
    local = %{pipeline | workers: nil}
    Dux.compute(local)
  end

  defp apply_single_op(dux, {:slice, offset, length}), do: Dux.slice(dux, offset, length)
  defp apply_single_op(dux, {:head, n}), do: Dux.head(dux, n)
  defp apply_single_op(dux, {:sort_by, spec}), do: %{dux | ops: dux.ops ++ [{:sort_by, spec}]}
  defp apply_single_op(dux, {:distinct, cols}), do: %{dux | ops: dux.ops ++ [{:distinct, cols}]}

  defp apply_single_op(dux, {:pivot_wider, names_col, values_col, agg}) do
    Dux.pivot_wider(dux, names_col, values_col, agg: agg)
  end

  defp apply_single_op(dux, {:pivot_longer, cols, names_to, values_to}) do
    Dux.pivot_longer(dux, cols, names_to: names_to, values_to: values_to)
  end

  defp apply_single_op(dux, op) do
    # Unknown op — append directly
    %{dux | ops: dux.ops ++ [op]}
  end

end