# Copyright 2025 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.

"""ScalarLeafHandler that implements the types.LeafHandler Protocol.

The primary purpose of this handler is to provide serialization and
deserialization for scalar values.
"""

import asyncio
import dataclasses
from typing import Awaitable, Sequence, Type

from absl import logging
import jax.numpy as jnp
import numpy as np
from orbax.checkpoint._src.futures import future
from orbax.checkpoint._src.serialization import type_handlers as type_handlers_v0
from orbax.checkpoint.experimental.v1._src.context import context as context_lib
from orbax.checkpoint.experimental.v1._src.serialization import types


Scalar = int | float | np.number
ScalarSerializationParam = types.SerializationParam[Scalar]
ScalarDeserializationParam = types.DeserializationParam["AbstractScalar"]


@dataclasses.dataclass
class AbstractScalar:
  """Abstract scalar representation of a scalar value.

  This an abstract scalar that can be used to represent
  the metadata belonging to a scalar value.

  dtype:
    Dtype of the scalar value. Note extra np.dtype and jnp.dtype are needed
    to support the v0 ScalarMetadata.
  """

  dtype: Type[Scalar] | np.dtype | jnp.dtype


def _create_v0_saving_paraminfo(
    param: ScalarSerializationParam,
    context: context_lib.Context,
    serialization_context: types.SerializationContext,
) -> type_handlers_v0.ParamInfo:
  """Creates a V0 ParamInfo from V1 params andn contexts for saving."""

  saving_options = context.array_options.saving

  return type_handlers_v0.ParamInfo(
      name=param.name,
      path=serialization_context.parent_dir.path / param.name,
      parent_dir=serialization_context.parent_dir.path,
      byte_limiter=serialization_context.byte_limiter,
      is_ocdbt_checkpoint=saving_options.use_ocdbt,
      use_zarr3=saving_options.use_zarr3,
      ocdbt_target_data_file_size=saving_options.ocdbt_target_data_file_size,
      ts_context=serialization_context.ts_context,
      value_typestr="scalar",
  )


def _create_v0_savearg(
    param: ScalarSerializationParam,
    context: context_lib.Context,
) -> type_handlers_v0.SaveArgs:
  """Creates a V0 SaveArgs from V1 params and context for saving."""

  fn = context.pytree_options.saving.create_array_storage_options_fn

  if fn:
    storage_options = fn(param.keypath, param.value)
    savearg = type_handlers_v0.SaveArgs(
        dtype=storage_options.dtype,
        chunk_byte_size=storage_options.chunk_byte_size,
        shard_axes=storage_options.shard_axes,
    )
  else:
    savearg = type_handlers_v0.SaveArgs()

  return savearg


def _create_v0_restore_paraminfo(
    param: types.DeserializationParam[None | AbstractScalar],
    context: context_lib.Context,
    deserialization_context: types.DeserializationContext,
) -> type_handlers_v0.ParamInfo:
  """Creates a V0 ParamInfo from V1 params and contexts for loading."""

  loading_options = context.array_options.Loading

  return type_handlers_v0.ParamInfo(
      name=param.name,
      path=deserialization_context.parent_dir / param.name,
      parent_dir=deserialization_context.parent_dir,
      skip_deserialize=False,
      byte_limiter=deserialization_context.byte_limiter,
      is_ocdbt_checkpoint=deserialization_context.ocdbt_checkpoint,
      ts_context=deserialization_context.ts_context,
      raise_array_data_missing_error=loading_options.raise_array_data_missing_error,
      use_zarr3=deserialization_context.zarr3_checkpoint,
  )


def _create_v0_restorearg(
    param: ScalarDeserializationParam,
) -> type_handlers_v0.RestoreArgs:
  """Creates a V0 RestoreArgs from V1 params."""

  if param.value is None:
    # use default restore args
    return type_handlers_v0.RestoreArgs()
  else:
    v = param.value
    if isinstance(v, AbstractScalar):
      restored_dtype = v.dtype
    elif isinstance(v, Scalar):
      restored_dtype = type(v)
    else:
      raise ValueError(f"Unsupported type: {type(v)}")

    logging.vlog(1, "name: %s, v.dtype: %s", param.name, restored_dtype)
    return type_handlers_v0.RestoreArgs(
        restore_type=restored_dtype,
        dtype=v.dtype,
    )


async def _async_futures(commit_futures: Sequence[future.Future]):
  await asyncio.gather(*[asyncio.to_thread(f.result) for f in commit_futures])


class ScalarLeafHandler(types.LeafHandler[Scalar, AbstractScalar]):
  """ScalarLeafHandler that implements the types.LeafHandler Protocol."""

  def __init__(
      self,
      *,
      context: context_lib.Context | None = None,
  ):
    self._context = context_lib.get_context(context)
    self._handler_impl = type_handlers_v0.ScalarHandler()

    logging.info("ScalarLeafHandler created.")

  async def serialize(
      self,
      params: Sequence[ScalarSerializationParam],
      serialization_context: types.SerializationContext,
  ) -> Awaitable[None]:
    """Serializes scalar values as a checkpointable to a storage location.

    Args:
      params: a sequence of ScalarSerializationParam per leaf.
      serialization_context: SerializationContext for the scalar leaf handler.

    Returns:
      Sequence of commit futures which can be awaited to complete the save
      operation.
    """
    values = [p.value for p in params]
    paraminfos = [
        _create_v0_saving_paraminfo(p, self._context, serialization_context)
        for p in params
    ]
    saveargs = [_create_v0_savearg(p, self._context) for p in params]

    await serialization_context.parent_dir.await_creation()

    commit_futures = await self._handler_impl.serialize(
        values, paraminfos, saveargs
    )
    assert commit_futures

    return _async_futures(commit_futures)

  async def deserialize(
      self,
      params: Sequence[types.DeserializationParam[AbstractScalar]],
      deserialization_context: types.DeserializationContext,
  ) -> Awaitable[Sequence[Scalar]]:
    """Returns sequence of Scalar values from a stored checkpointable location.

    Args:
      params: sequence of ScalarDeserializationParam per leaf.
      deserialization_context: ScalarDeserializationContext for the leaf
        handler.

    Returns:
      The deserialized sequence of scalar values as leaves.
    """

    # validate all parameters
    paraminfos = [
        _create_v0_restore_paraminfo(p, self._context, deserialization_context)
        for p in params
    ]
    restoreargs = [_create_v0_restorearg(p) for p in params]

    return asyncio.create_task(
        self._handler_impl.deserialize(paraminfos, restoreargs)
    )

  async def metadata(
      self,
      params: Sequence[types.DeserializationParam[None]],
      deserialization_context: types.DeserializationContext,
  ) -> Sequence[AbstractScalar]:
    """Returns a squence of AbstractScalar from a stored checkpointable location.

    Args:
      params: sequence of ScalarDeserializationParam per scalar value leaf.
      deserialization_context: DeserializationContext for the scalar leaf
        handler.

    Returns:
      Sequence of ScalarMetadata for each provided DeserializationParam.
    """
    paraminfos = [
        _create_v0_restore_paraminfo(p, self._context, deserialization_context)
        for p in params
    ]

    async def _convert_to_scalar_metadata() -> Sequence[AbstractScalar]:
      v0_metadatas = await self._handler_impl.metadata(paraminfos)

      ret = []
      for meta in v0_metadatas:
        scalar_metadata = AbstractScalar(
            dtype=meta.dtype,
        )
        ret.append(scalar_metadata)

        if logging.vlog_is_on(1):
          logging.vlog(1, "scalar_metadata: %r", scalar_metadata)

      return ret

    return await _convert_to_scalar_metadata()
