# Copyright 2026 The Orbax 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.

"""Benchmarks for orbax.checkpoint.PyTreeCheckpointHandler."""

from __future__ import annotations

import dataclasses
import pprint
import time
from typing import Any

from absl import logging
import jax
from jax.experimental import multihost_utils
import numpy as np
from orbax.checkpoint import v1 as ocp
from orbax.checkpoint._src.testing.benchmarks.core import core as benchmarks_core
from orbax.checkpoint._src.testing.benchmarks.core import metric as metric_lib
from orbax.checkpoint.experimental.caching import client


SERVICE_URL = "http://service-dns/"
LUSTRE_PATH_PREFIX = "/lustre/"
GCS_PATH_PREFIX = "gs://"


def _metrics_to_measure(options: LustreBenchmarkOptions) -> list[str]:
  """Returns the list of metrics to measure."""
  del options
  metrics = ["time", "rss", "io"]
  return metrics


# ==============================================================================
# 1. Define the Options Dataclass for this specific benchmark
# ==============================================================================
@dataclasses.dataclass(frozen=True)
class LustreBenchmarkOptions(benchmarks_core.BenchmarkOptions):
  """Configuration options for benchmarks targeting PyTreeCheckpointHandler.

  Each attribute can be a single value or a list of values to create
  a parameter sweep.

  TODO(cpgaffney): Add meaningful benchmark options.

  Attributes:
    use_ocdbt: Whether to use OCDBT for checkpointing. Currently ignored.
  """

  use_ocdbt: bool = True

  def is_valid(self):
    return True


def _get_xid() -> int:
  """Returns the XID for this run."""
  xid = multihost_utils.broadcast_one_to_all(
      np.asarray(int(time.time()))
  ).item()
  logging.info("XID: %s", xid)
  return xid


# ==============================================================================
# 2. Implement the Benchmark Generator
# ==============================================================================
@benchmarks_core.benchmark_options(LustreBenchmarkOptions)
class LustreBenchmark(benchmarks_core.BenchmarksGenerator):
  """Benchmark for caching with Lustre and GCS, using the experimental service."""

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self._client = client.StorageServiceClient()
    self._xid = _get_xid()

  def _clear_pytree(self, pytree: Any) -> Any:
    """Clears the pytree to free up memory."""
    return jax.tree.map(
        lambda x: x.delete() if isinstance(x, jax.Array) else None, pytree
    )

  def test_fn(
      self, context: benchmarks_core.TestContext
  ) -> benchmarks_core.TestResult:
    """The core test logic for a single save/restore cycle.

    This function is called for each combination of options generated by the
    framework. It uses the `context.options` to configure the handler
    dynamically for each run.

    Args:
      context: The test context containing the pytree, path, and options.

    Returns:
      The test result containing the metrics.
    """
    logging.info(
        "JAX info: %s processes, %s devices, %s process index",
        jax.process_count(),
        jax.device_count(),
        jax.process_index(),
    )
    metrics = metric_lib.Metrics()
    assert isinstance(context.options, LustreBenchmarkOptions)

    logging.info("Benchmark options: %s", pprint.pformat(context.options))

    metrics_to_measure = _metrics_to_measure(context.options)

    step = context.repeat_index or 0

    with metrics.measure("resolve_cache", metrics_to_measure):
      resolved_path: str = self._client.resolve(self._xid, step)
      assert resolved_path.startswith(LUSTRE_PATH_PREFIX), resolved_path
    with metrics.measure("save_cache", metrics_to_measure):
      ocp.save_pytree(resolved_path, context.pytree)
    with metrics.measure("finalize_cache", metrics_to_measure):
      self._client.finalize(self._xid, step)

    with metrics.measure("save", metrics_to_measure):
      ocp.save_pytree(context.path / str(step), context.pytree)

    abstract_pytree = jax.tree.map(
        ocp.arrays.to_shape_dtype_struct, context.pytree
    )
    context.pytree = self._clear_pytree(context.pytree)

    # Only exercise restore on the previous step.
    if step > 0:
      resolved_path: str = self._client.resolve(self._xid, step - 1)
      # Previous step should have been moved to GCS by now.
      assert resolved_path.startswith(GCS_PATH_PREFIX), resolved_path
      with metrics.measure("prefetch_cache", metrics_to_measure):
        self._client.prefetch(self._xid, step - 1)
      with metrics.measure("wait_prefetch_cache", metrics_to_measure):
        self._client.await_transfer(self._xid, step - 1)
      with metrics.measure("restore_cache", metrics_to_measure):
        restored_pytree = ocp.load_pytree(resolved_path, abstract_pytree)
      restored_pytree = self._clear_pytree(restored_pytree)
      del restored_pytree
      with metrics.measure("restore", metrics_to_measure):
        restored_pytree = ocp.load_pytree(
            context.path / str(step), abstract_pytree
        )
      restored_pytree = self._clear_pytree(restored_pytree)
      del restored_pytree

    return benchmarks_core.TestResult(metrics=metrics)
