/* Copyright 2025 The JAX Authors

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <Python.h>

#include <algorithm>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "nanobind/nanobind.h"
#include "nanobind/stl/vector.h"  // IWYU pragma: keep
#include "jaxlib/nb_class_ptr.h"
#include "jaxlib/py_array.h"
#include "jaxlib/py_client.h"
#include "jaxlib/py_user_context.h"
#include "jaxlib/to_ifrt_sharding.h"
#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
#include "xla/pjrt/status_casters.h"
#include "xla/python/ifrt/array.h"
#include "xla/python/ifrt/array_spec.h"
#include "xla/python/ifrt/device.h"
#include "xla/python/ifrt/device_list.h"
#include "xla/python/ifrt/dtype.h"
#include "xla/python/ifrt/remap_plan.h"
#include "xla/python/ifrt/shape.h"
#include "xla/python/ifrt/sharding.h"
#include "xla/python/ifrt/user_context.h"
#include "xla/python/nb_absl_flat_hash_map.h"  // IWYU pragma: keep
#include "xla/python/nb_absl_flat_hash_set.h"  // IWYU pragma: keep
#include "xla/python/nb_absl_span.h"  // IWYU pragma: keep
#include "xla/python/pjrt_ifrt/pjrt_client.h"
#include "xla/python/types.h"
#include "xla/tsl/concurrency/ref_count.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/util.h"

namespace jax {

namespace nb = ::nanobind;

namespace {
using xla::ifrt::Array;
using xla::ifrt::ArrayCopySemantics;
using xla::ifrt::ArrayRef;
using xla::ifrt::ArraySpec;
using xla::ifrt::DType;
using xla::ifrt::DeviceListRef;
using xla::ifrt::RemapPlan;
using xla::ifrt::Shape;
using xla::ifrt::ShardingRef;

Shape GetShardShape(const PyArray& py_array) {
  DCHECK(PyGILState_Check());
  const nb::object& py_aval = py_array.aval();
  nb::object py_shape = py_aval.attr("shape");
  nb::object py_shard_shape = py_array.sharding().attr("shard_shape")(py_shape);
  std::vector<int64_t> dims = nb::cast<std::vector<int64_t>>(py_shard_shape);
  return Shape(Shape::Dimensions(dims.begin(), dims.end()));
}

std::vector<Shape> GetShardShapes(
    absl::Span<const PyArray> py_arrays) {
  std::vector<Shape> shard_shapes;
  shard_shapes.reserve(py_arrays.size());
  for (const PyArray& py_array : py_arrays) {
    shard_shapes.push_back(GetShardShape(py_array));
  }
  return shard_shapes;
}

// If `backend` is nullptr, sets it to `array.py_client()`; otherwise checks
// that `backend` equals `array.py_client()`.
absl::Status PyClientFromPyArray(const PyArray& array,
                                 nb_class_ptr<PyClient>& backend) {
  if (array.py_client().get() == nullptr) {
    return absl::InternalError("Unexpected array with py_client as nullptr.");
  }
  if (backend.get() == nullptr) {
    backend = array.py_client();
  } else if (backend.get() != array.py_client().get()) {
    std::string old_description =
        absl::StrFormat("%p/%s/%s/%s/%s", backend.get(),
                        backend->platform_name(), backend->platform_version(),
                        backend->runtime_type(), backend->raw_platform_name());
    std::string new_description =
        absl::StrFormat("%p/%s/%s/%s/%s", array.py_client().get(),
                        array.py_client()->platform_name(),
                        array.py_client()->platform_version(),
                        array.py_client()->runtime_type(),
                        array.py_client()->raw_platform_name());
    return absl::InvalidArgumentError(absl::StrCat(
        "py_client mismatch: ", old_description, " vs ", new_description));
  }
  return absl::OkStatus();
}


// Runs `xla::ifrt::Client::ReshardArrays`.
absl::StatusOr<nb::list> ExperimentalReshardArrays(nb::sequence py_arrays,
                                                   nb::sequence out_shardings,
                                                   bool donate_input) {
  const int num_arrays = nb::len(py_arrays);

  if (nb::len(out_shardings) != num_arrays) {
    return absl::InvalidArgumentError(
        absl::StrCat("Number of out_shardings must match number of arrays: ",
                     nb::len(out_shardings), num_arrays));
  }

  if (num_arrays == 0) {
    return nb::list();
  }

  PyUserContextScope user_context_scope;
  nb_class_ptr<PyClient> backend;
  std::vector<ArrayRef> ifrt_arrays;
  std::vector<ArraySpec> ifrt_specs;
  ifrt_arrays.reserve(num_arrays);
  ifrt_specs.reserve(num_arrays);

  for (int i = 0; i < num_arrays; ++i) {
    PyArray array = nb::cast<PyArray>(py_arrays[i]);
    if (array.ifrt_array_ref() == nullptr) {
      return absl::InvalidArgumentError(
          absl::StrCat("Input array ", i, " has been donated or deleted."));
    }
    ABSL_RETURN_IF_ERROR(PyClientFromPyArray(array, backend));
    ifrt_arrays.push_back(array.ifrt_array_ref());

    ABSL_ASSIGN_OR_RETURN(DType ifrt_dtype,
                          xla::DtypeToIfRtDType(array.dtype()));
    Shape ifrt_shape(array.shape());
    ABSL_ASSIGN_OR_RETURN(ShardingRef ifrt_sharding,
                          GetIfrtHloSharding(out_shardings[i], ifrt_shape));
    ifrt_specs.push_back(ArraySpec{
        /*dtype=*/std::move(ifrt_dtype),
        /*shape=*/std::move(ifrt_shape),
        /*sharding=*/std::move(ifrt_sharding),
    });
  }

  const ArrayCopySemantics copy_semantics =
      donate_input ? ArrayCopySemantics::kDonateInput
                   : ArrayCopySemantics::kAlwaysCopy;

  std::vector<ArrayRef> outputs;
  {
    nb::gil_scoped_release gil_release;
    ABSL_ASSIGN_OR_RETURN(
        outputs, backend->ifrt_client()->ReshardArrays(
                     absl::MakeSpan(ifrt_arrays), ifrt_specs, copy_semantics));
  }

  nb::list result;
  for (int i = 0; i < num_arrays; ++i) {
    PyArray new_py_array = PyArray::MakeFromIfrtArrayAndSharding(
        backend, std::move(outputs[i]), out_shardings[i],
        /*weak_type=*/false,
        /*committed=*/true,
        /*skip_checks=*/true);
    result.append(std::move(new_py_array));
  }

  return result;
}

absl::StatusOr<std::vector<std::vector<nb::object>>>
ExperimentalSplitByMeshAxis(
    nb::object py_arrays_py, absl::Span<const int> sharded_dim_idxs,
    absl::Span<const int> mesh_axis_sizes, int mesh_axis_idx,
    absl::Span<const int> mesh_axis_sections,
    absl::Span<const std::vector<nb::object>> submesh_shardings, bool donate) {
  // Using `nb_class_ptr<PyClient>` requires GIL.
  DCHECK(PyGILState_Check());

  auto py_arrays = nb::cast<std::vector<PyArray>>(py_arrays_py);
  if (py_arrays.empty()) {
    return std::vector<std::vector<nb::object>>();
  }

  if (sharded_dim_idxs.size() != py_arrays.size()) {
    return absl::InvalidArgumentError(
        absl::StrCat("Number of sharded_dim_idxs must match number of arrays: ",
                     sharded_dim_idxs.size(), " vs ", py_arrays.size()));
  }
  if (submesh_shardings.size() != py_arrays.size()) {
    return absl::InvalidArgumentError(absl::StrCat(
        "Number of submesh_shardings must match number of arrays: ",
        submesh_shardings.size(), " vs ", py_arrays.size()));
  }

  int num_submeshes = submesh_shardings[0].size();
  if (mesh_axis_sections.size() != num_submeshes) {
    return absl::InvalidArgumentError(absl::StrCat(
        "Number of mesh_axis_sections must match number of submeshes: ",
        mesh_axis_sections.size(), " vs ", num_submeshes));
  }

  PyUserContextScope user_context_scope;
  // All input arrays are expected to use the same mesh.
  ABSL_ASSIGN_OR_RETURN(DeviceListRef device_list,
                        GetIfrtDeviceList(py_arrays[0].sharding()));
  int num_devices = device_list->size();
  // The last entry in `mexh_axis_sections` contains the mesh axis size.
  int mesh_axis_size = mesh_axis_sections.back();
  if (num_devices % mesh_axis_size != 0) {
    return absl::InvalidArgumentError(absl::StrCat(
        "Number of devices must be divisible by the mesh axis size: ",
        num_devices, " vs ", mesh_axis_size));
  }

  std::vector<ArraySpec> input_specs;
  std::vector<ArraySpec> output_specs;
  absl::flat_hash_map<int, std::vector<RemapPlan::InputDeviceRange>>
      input_devices_for_output_map;

  nb_class_ptr<PyClient> backend;
  std::vector<ArrayRef> input_ifrt_arrays;
  input_ifrt_arrays.reserve(py_arrays.size());
  for (int array_idx = 0; array_idx < py_arrays.size(); ++array_idx) {
    ABSL_RETURN_IF_ERROR(PyClientFromPyArray(py_arrays[array_idx], backend));
    xla::ifrt::ArrayRef array = py_arrays[array_idx].ifrt_array_ref();
    if (array == nullptr) {
      return xla::InvalidArgument("Input array #%d has been donated or deleted",
                                  array_idx);
    }

    input_specs.push_back(ArraySpec{/*dtype=*/array->dtype(),
                                    /*shape=*/array->shape(),
                                    /*sharding=*/array->shared_ptr_sharding()});

    for (int submesh_idx = 0; submesh_idx < num_submeshes; ++submesh_idx) {
      int submesh_axis_size = mesh_axis_sections[submesh_idx];
      if (submesh_idx > 0) {
        submesh_axis_size -= mesh_axis_sections[submesh_idx - 1];
      }
      Shape subshape = [&]() -> Shape {
        if (sharded_dim_idxs[array_idx] >= 0) {
          std::vector<int64_t> dims(array->shape().dims().begin(),
                                    array->shape().dims().end());
          dims[sharded_dim_idxs[array_idx]] =
              dims[sharded_dim_idxs[array_idx]] / mesh_axis_size *
              submesh_axis_size;
          return Shape(dims);
        } else {
          // The arrays is replicated, so its shape does not change.
          return array->shape();
        }
      }();
      ABSL_ASSIGN_OR_RETURN(
          ShardingRef ifrt_submesh_sharding,
          GetIfrtHloSharding(submesh_shardings[array_idx][submesh_idx],
                             subshape));
      int out_idx = output_specs.size();
      input_devices_for_output_map[out_idx].push_back(
          RemapPlan::InputDeviceRange{
              .in_array = array_idx,
              .input_devices = ifrt_submesh_sharding->devices(),
          });
      output_specs.push_back(
          ArraySpec{/*dtype=*/array->dtype(),
                    /*shape=*/std::move(subshape),
                    /*sharding=*/std::move(ifrt_submesh_sharding)});
    }

    input_ifrt_arrays.push_back(array);
  }

  RemapPlan remap_plan(std::move(input_specs), std::move(output_specs),
                       std::move(input_devices_for_output_map));
  DCHECK_OK(remap_plan.Validate());

  std::vector<ArrayRef> result_ifrt_arrays;
  {
    nb::gil_scoped_release gil_release;
    ABSL_ASSIGN_OR_RETURN(result_ifrt_arrays,
                          backend->ifrt_client()->RemapArrays(
                              remap_plan, absl::MakeSpan(input_ifrt_arrays),
                              donate ? ArrayCopySemantics::kDonateInput
                                     : ArrayCopySemantics::kReuseInput));
  }

  DCHECK_EQ(result_ifrt_arrays.size(), py_arrays.size() * num_submeshes);

  // Wrap IFRT arrays as JAX arrays.
  std::vector<std::vector<nb::object>> py_results;
  int offset_in_results = 0;
  for (int array_idx = 0; array_idx < py_arrays.size(); ++array_idx) {
    auto& py_submesh_results = py_results.emplace_back();
    for (int submesh_idx = 0; submesh_idx < num_submeshes; ++submesh_idx) {
      PyArray new_py_array = PyArray::MakeFromIfrtArrayAndSharding(
          backend,
          std::move(result_ifrt_arrays[offset_in_results + submesh_idx]),
          submesh_shardings[array_idx][submesh_idx],
          py_arrays[array_idx].weak_type(),
          /*committed=*/true,
          /*skip_checks=*/true);
      py_submesh_results.push_back(new_py_array);
    }
    offset_in_results += num_submeshes;
  }

  return py_results;
}

absl::StatusOr<std::vector<nb::object>> ExperimentalConcatenateByMeshAxis(
    nb::object py_arrays_py, absl::Span<const int> sharded_dim_idxs,
    absl::Span<const int> mesh_axis_sizes, int mesh_axis_idx,
    absl::Span<const int> mesh_axis_sections,
    absl::Span<const nb::object> out_shardings, bool donate) {
  DCHECK(PyGILState_Check());

  auto py_arrays_list =
      nb::cast<std::vector<std::vector<PyArray>>>(py_arrays_py);
  if (py_arrays_list.empty()) {
    return std::vector<nb::object>();
  }
  int num_output_arrays = py_arrays_list.size();
  int num_input_arrays_per_output = py_arrays_list[0].size();

  if (sharded_dim_idxs.size() != num_output_arrays) {
    return absl::InvalidArgumentError("sharded_dim_idxs size mismatch");
  }
  if (out_shardings.size() != num_output_arrays) {
    return absl::InvalidArgumentError("out_shardings size mismatch");
  }

  for (const auto& py_arrays : py_arrays_list) {
    if (py_arrays.size() != num_input_arrays_per_output) {
      return absl::InvalidArgumentError(absl::StrCat(
          "`ConcatenateByMeshAxis` expects all output arrays to have the same ",
          "number of input arrays.  Saw ", py_arrays.size(), " vs ",
          num_input_arrays_per_output, "."));
    }
    std::vector<Shape> shard_shapes = GetShardShapes(py_arrays);
    for (int idx = 1; idx < shard_shapes.size(); ++idx) {
      if (shard_shapes[idx] != shard_shapes[0]) {
        return absl::InvalidArgumentError(absl::StrCat(
            "`ConcatenateByMeshAxis` expects all input arrays for a given ",
            "output to have the same shard shape. Input array #",
            idx, " has shard shape ", shard_shapes[idx],
            ", which differs from shard shape of arrays with ",
            "smaller indices ",
            shard_shapes[0]));
      }
    }
  }

  std::vector<ArraySpec> input_specs;
  std::vector<ArraySpec> output_specs;
  absl::flat_hash_map<int, std::vector<RemapPlan::InputDeviceRange>>
      input_devices_for_output_map;

  std::vector<ArrayRef> input_ifrt_arrays;
  input_ifrt_arrays.reserve(num_output_arrays * num_input_arrays_per_output);

  nb_class_ptr<PyClient> backend;
  ABSL_RETURN_IF_ERROR(PyClientFromPyArray(py_arrays_list[0][0], backend));

  for (int array_idx = 0; array_idx < num_output_arrays; ++array_idx) {
    const auto& py_arrays = py_arrays_list[array_idx];
    if (py_arrays.size() != num_input_arrays_per_output) {
      return absl::InvalidArgumentError("Inconsistent number of input arrays");
    }

    int64_t concatenated_sharded_dim_size = 0;
    for (int input_idx = 0; input_idx < num_input_arrays_per_output;
         ++input_idx) {
      const PyArray& py_array = py_arrays[input_idx];
      xla::ifrt::ArrayRef array = py_array.ifrt_array_ref();
      if (array == nullptr) {
        return xla::InvalidArgument(
            "Input array #%d for output #%d has been donated or deleted",
            input_idx, array_idx);
      }

      if (sharded_dim_idxs[array_idx] >= 0) {
        concatenated_sharded_dim_size +=
            array->shape().dims()[sharded_dim_idxs[array_idx]];
      }

      int in_array_idx = input_specs.size();
      input_devices_for_output_map[array_idx].push_back(
          RemapPlan::InputDeviceRange{
              .in_array = in_array_idx,
              .input_devices = array->sharding().devices(),
          });
      input_specs.push_back(
          ArraySpec{.dtype = array->dtype(),
                    .shape = array->shape(),
                    .sharding = array->shared_ptr_sharding()});
      input_ifrt_arrays.push_back(array);
    }

    xla::ifrt::ArrayRef first_array = py_arrays[0].ifrt_array_ref();
    if (sharded_dim_idxs[array_idx] < 0) {
      ABSL_ASSIGN_OR_RETURN(
          ShardingRef ifrt_sharding,
          GetIfrtHloSharding(out_shardings[array_idx], first_array->shape()));
      output_specs.push_back(ArraySpec{.dtype = first_array->dtype(),
                                       .shape = first_array->shape(),
                                       .sharding = std::move(ifrt_sharding)});
    } else {
      std::vector<int64_t> concatenated_dims(
          first_array->shape().dims().begin(),
          first_array->shape().dims().end());
      concatenated_dims[sharded_dim_idxs[array_idx]] =
          concatenated_sharded_dim_size;
      Shape concatenated_shape = Shape(concatenated_dims);
      ABSL_ASSIGN_OR_RETURN(
          ShardingRef ifrt_sharding,
          GetIfrtHloSharding(out_shardings[array_idx], concatenated_shape));
      output_specs.push_back(ArraySpec{.dtype = first_array->dtype(),
                                       .shape = std::move(concatenated_shape),
                                       .sharding = std::move(ifrt_sharding)});
    }
  }

  PyUserContextScope user_context_scope;
  std::vector<ArrayRef> result_ifrt_arrays;
  {
    nb::gil_scoped_release gil_release;
    RemapPlan remap_plan(std::move(input_specs), std::move(output_specs),
                         std::move(input_devices_for_output_map));
    DCHECK_OK(remap_plan.Validate());

    ABSL_ASSIGN_OR_RETURN(result_ifrt_arrays,
                          backend->ifrt_client()->RemapArrays(
                              remap_plan, absl::MakeSpan(input_ifrt_arrays),
                              donate ? ArrayCopySemantics::kDonateInput
                                     : ArrayCopySemantics::kReuseInput));
    DCHECK_EQ(result_ifrt_arrays.size(), num_output_arrays);
  }

  std::vector<nb::object> py_results;
  py_results.reserve(num_output_arrays);
  for (int array_idx = 0; array_idx < num_output_arrays; ++array_idx) {
    const auto& py_array_template = py_arrays_list[array_idx][0];
    PyArray new_py_array = PyArray::MakeFromIfrtArrayAndSharding(
        backend, std::move(result_ifrt_arrays[array_idx]),
        out_shardings[array_idx], py_array_template.weak_type(),
        /*committed=*/true,
        /*skip_checks=*/true);
    py_results.push_back(std::move(new_py_array));
  }
  return py_results;
}

jax::nb_class_ptr<jax::PyClient> CreateCpuClient(
    absl::flat_hash_set<int> addressable_devices,
    absl::flat_hash_map<int, int> device_id_to_process_index) {
  xla::CpuClientOptions options;
  options.asynchronous = true;
  options.cpu_device_count = addressable_devices.size();
  auto cpu_client =
      xla::ValueOrThrow(xla::GetXlaPjrtCpuClient(std::move(options)));

  xla::ifrt::PjRtClient::CreateOptions create_options;
  create_options.pjrt_client = std::move(cpu_client);
  xla::ifrt::PjRtClient::CreateOptions::GlobalDeviceMapping&
      global_device_mapping = create_options.global_device_mapping.emplace();
  for (int device_id : addressable_devices) {
    global_device_mapping.addressable_device_ids.insert(
        xla::ifrt::DeviceId(device_id));
  }
  for (const auto& [device_id, process_index] : device_id_to_process_index) {
    global_device_mapping
        .device_id_to_process_index[xla::ifrt::DeviceId(device_id)] =
        process_index;
  }
  create_options.sort_devices_by_process_index = false;
  return jax::PyClient::Make(
      xla::ValueOrThrow(xla::ifrt::PjRtClient::Create(create_options)));
}

}  // namespace

NB_MODULE(_pathways, m) {
  m.def("_transfer_to_shardings",
        xla::ValueOrThrowWrapper(ExperimentalReshardArrays), nb::arg("arrays"),
        nb::arg("out_shardings"), nb::arg("donate") = false);
  m.def("_split_by_mesh_axis",
        xla::ValueOrThrowWrapper(ExperimentalSplitByMeshAxis),
        nb::arg("arrays"), nb::arg("sharded_dim_idxs"),
        nb::arg("mesh_axis_sizes"), nb::arg("mesh_axis_idx"),
        nb::arg("mesh_axis_sections"), nb::arg("submesh_shardings"),
        nb::arg("donate"));
  m.def("_concatenate_by_mesh_axis",
        xla::ValueOrThrowWrapper(ExperimentalConcatenateByMeshAxis),
        nb::arg("arrays"), nb::arg("sharded_dim_idxs"),
        nb::arg("mesh_axis_sizes"), nb::arg("mesh_axis_idx"),
        nb::arg("mesh_axis_sections"), nb::arg("out_shardings"),
        nb::arg("donate"));
  m.def("_create_cpu_client", CreateCpuClient, nb::arg("addressable_devices"),
        nb::arg("device_id_to_process_index"),
        // clang-format off
        nb::sig(
          "def _create_cpu_client("
          "addressable_devices: collections.abc.Set[int], "
          "device_id_to_process_index: collections.abc.Mapping[int, int]"
          ") -> Client")
        // clang-format on
        );
}

}  // namespace jax