"""
HeGM -- Heterogeneous GPU Migration Framework.

Importing this package:
    1. Patches ``builtins.enumerate`` to transparently skip DataLoader
       batches after a checkpoint restore (zero code changes required).
    2. Registers the SIGUSR1 signal handler.
    3. Installs PEP-451 import hooks for every registered backend
       (PyTorch today, TensorFlow in the future).
    4. Exposes the public API: :func:`global_step`, :func:`resumable`.

Typical training scripts need **no modifications** -- HeGM is injected via
``sitecustomize.py`` and everything works transparently::

    # Unchanged user code:
    for step, batch in enumerate(dataloader):
        ...   # step resumes from checkpoint automatically
"""

import builtins as _builtins
import itertools
import os
import signal
import sys
from typing import Any, Iterable, Iterator, Tuple

from hegm._config import CHECKPOINT_PATH, EXIT_CODE, log, warn
from hegm._hook import install_hook
from hegm.backends import get_active, get_backends

# Re-import all built-in backends so they auto-register.
import hegm.backends.pytorch  # noqa: F401

# ---------------------------------------------------------------------------
# Transparent enumerate() patch
# ---------------------------------------------------------------------------

_original_enumerate = _builtins.enumerate


def _make_enumerate_proxy():
    """Create a class-like proxy for ``enumerate`` that stays compatible with
    torch._dynamo while still injecting HeGM's DataLoader skip behavior.
    """
    original = _original_enumerate

    class _EnumerateMeta(type):
        def __call__(self, iterable, start=0):  # type: ignore[override]
            # Import the backend module -- always available because we import it above.
            from hegm.backends import pytorch as _pt  # type: ignore[attr-defined]

            skip = _pt._dataloader_skip
            if skip > 0:
                try:
                    from torch.utils.data import DataLoader

                    if isinstance(iterable, DataLoader):
                        _pt._dataloader_skip = 0
                        log("Resuming dataloader -- skipping %d batches." % skip)
                        return original(
                            itertools.islice(iterable, skip, None),
                            start=start + skip,
                        )
                except ImportError:
                    pass

            return original(iterable, start)

    # Create a dummy class object whose metaclass controls the call behavior.
    EnumerateProxy = _EnumerateMeta("enumerate", (), {})
    return EnumerateProxy


def _should_patch_enumerate() -> bool:
    """Return True if we should patch ``builtins.enumerate``.

    You can force-disable the patch by setting
    ``HEGM_DISABLE_ENUMERATE_PATCH=1`` in the environment.
    """
    disable_env = os.environ.get("HEGM_DISABLE_ENUMERATE_PATCH", "").strip().lower()
    if disable_env in {"1", "true", "yes"}:
        log(
            "HEGM_DISABLE_ENUMERATE_PATCH set; leaving enumerate() unmodified. "
            "Use hegm.resumable(dataloader) for checkpoint-aware iteration.",
            level="INFO",
        )
        return False

    return True


if _should_patch_enumerate():
    _builtins.enumerate = _make_enumerate_proxy()  # type: ignore[assignment]

# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def global_step(optimizer: Any = None) -> int:
    """Return the current global training step.

    Parameters
    ----------
    optimizer : optional
        If omitted, returns the step count of the **first** tracked
        optimizer (covers the single-optimizer common case).

    Returns
    -------
    int
        Number of ``optimizer.step()`` calls completed so far.
        Returns ``0`` on a fresh run (no checkpoint).
    """
    backend = get_active()
    if backend is None:
        return 0
    return backend.global_step(optimizer)


def resumable(
    dataloader: Iterable,
    optimizer: Any = None,
) -> Iterator[Tuple[int, Any]]:
    """Wrap a dataloader so iteration resumes from the last checkpoint.

    .. note::

        Since v0.2.0 HeGM patches ``enumerate()`` transparently, so this
        helper is no longer required for the common case.  It is kept for
        backward compatibility and for advanced use cases (e.g. multiple
        dataloaders or custom skip logic).

    Drop-in replacement for ``enumerate(dataloader)``::

        # Before
        for step, batch in enumerate(dataloader):
            ...

        # After (checkpoint-aware)
        for step, batch in hegm.resumable(dataloader):
            ...

    On a fresh run this is identical to ``enumerate(dataloader)``.
    After a checkpoint restore it skips already-processed batches
    and yields ``(step, batch)`` tuples starting from the saved step.

    Parameters
    ----------
    dataloader : Iterable
        Any iterable (typically a ``torch.utils.data.DataLoader``).
    optimizer : optional
        Passed to :func:`global_step` to look up the step count.
        Omit for the common single-optimizer case.
    """
    start = global_step(optimizer)
    if start > 0:
        log("Resuming dataloader from step %d (skipping %d batches)." % (start, start))
    # islice is O(n) for the skipped items but avoids GPU work.
    it = itertools.islice(dataloader, start, None)
    yield from _original_enumerate(it, start=start)


# ---------------------------------------------------------------------------
# SIGUSR1 signal handler
# ---------------------------------------------------------------------------


def _sigusr1_handler(signum: int, frame: Any) -> None:
    """Save checkpoint via the active backend, then ``sys.exit(EXIT_CODE)``."""
    backend = get_active()

    if backend is None:
        log("SIGUSR1 received but no ML framework imported -- nothing to save.")
        sys.exit(EXIT_CODE)

    log("SIGUSR1 received. Saving checkpoint...")
    try:
        backend.save_checkpoint(CHECKPOINT_PATH)
    except Exception as exc:
        warn(f"Error saving checkpoint: {exc}")

    log("Checkpoint saved. Exiting for migration.")

    try:
        backend.release_gpu()
    except Exception:
        pass

    sys.exit(EXIT_CODE)


if hasattr(signal, "SIGUSR1"):
    signal.signal(signal.SIGUSR1, _sigusr1_handler)


# ---------------------------------------------------------------------------
# Install import hooks for all registered backends
# ---------------------------------------------------------------------------

for _backend in get_backends():
    install_hook(_backend.target_module, _backend.apply_patches)
