"""
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. Patches ``builtins.iter`` so frameworks using ``iter(dataloader)``
       (e.g. Hugging Face Trainer, SFTTrainer) also skip batches on resume.
    3. Patches Hugging Face Trainer when transformers is imported (progress
       bar resume, gradient_accumulation_steps). Disable with
       HEGM_DISABLE_HF_INTEGRATION=1.
    4. Optionally patches ``builtins.range`` so epoch loops resume transparently
       (see HEGM_DISABLE_EPOCH_RANGE_PATCH).
    5. Registers the SIGUSR1 signal handler.
    6. Installs PEP-451 import hooks for every registered backend
       (PyTorch today, TensorFlow in the future).
    7. Exposes the public API: :func:`global_step`, :func:`resumable`.

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

    # Unchanged user code:
    for epoch in range(num_epochs):
        for step, batch in enumerate(dataloader):
            ...   # epoch and step resume from checkpoint automatically

    # Hugging Face Trainer/SFTTrainer: progress bar and batch skip auto-resume
    trainer = SFTTrainer(model=..., args=...)
    trainer.train()
"""

import builtins as _builtins
import inspect
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.
# PyTorch backend is optional; HeGM works without it (e.g. before user installs torch).
try:
    import hegm.backends.pytorch  # noqa: F401
except Exception as exc:
    log("PyTorch backend not loaded: %s" % exc, level="DEBUG")
# Transparent HF Trainer integration (patches Trainer when transformers is imported).
try:
    import hegm.hf_integration  # noqa: F401
except Exception as exc:
    log("HF integration not loaded: %s" % exc, level="DEBUG")

# ---------------------------------------------------------------------------
# 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]
            try:
                from hegm.backends import pytorch as _pt  # type: ignore[attr-defined]
            except ImportError:
                return original(iterable, start)
            skip = _pt._dataloader_skip
            if skip >= 0:
                # Detect DataLoader by type without importing torch (avoids circular
                # import when enumerate is used during stdlib/torch bootstrap).
                cls = type(iterable)
                is_dataloader = (
                    cls.__module__ == "torch.utils.data.dataloader"
                    and cls.__name__ == "DataLoader"
                )
                if is_dataloader:
                    # Auto-infer steps_per_epoch for epoch-mode (option 3)
                    if _pt._steps_per_epoch is None and hasattr(iterable, "__len__"):
                        try:
                            _pt._steps_per_epoch = len(iterable)
                        except (TypeError, NotImplementedError):
                            pass
                    if skip > 0:
                        _pt._dataloader_skip = 0
                        log("Resuming dataloader -- skipping %d batches." % skip)
                        return original(
                            itertools.islice(iterable, skip, None),
                            start=start + skip,
                        )

            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


# ---------------------------------------------------------------------------
# Transparent iter() patch (Hugging Face Trainer / SFTTrainer compatibility)
# ---------------------------------------------------------------------------
# The Trainer uses iter(train_dataloader) and next(), not enumerate().
# Without this patch, batches would restart from 0 on resume.

_original_iter = _builtins.iter


def _iter_proxy(obj, *args, **kwargs):
    """Proxy for iter() that skips DataLoader batches after checkpoint restore."""
    if args or kwargs:
        return _original_iter(obj, *args, **kwargs)
    # Detect DataLoader by type without importing torch (avoids circular import
    # when iter is used during stdlib/torch bootstrap, e.g. inspect.getsourcelines).
    cls = type(obj)
    is_dataloader = (
        (cls.__module__ == "torch.utils.data.dataloader" and cls.__name__ == "DataLoader")
        or ("accelerate" in cls.__module__ and "DataLoader" in cls.__name__)
    )
    if is_dataloader:
        try:
            from hegm.backends import pytorch as _pt
            if _pt._dataloader_skip > 0:
                skip = _pt._dataloader_skip
                _pt._dataloader_skip = 0
                log("Resuming dataloader (iter) -- skipping %d batches." % skip)
                return itertools.islice(obj, skip, None)
        except ImportError:
            pass
    return _original_iter(obj)


def _should_patch_iter() -> bool:
    """Return True if we should patch ``builtins.iter`` for DataLoader skip."""
    disable_env = os.environ.get("HEGM_DISABLE_ITER_PATCH", "").strip().lower()
    if disable_env in {"1", "true", "yes"}:
        return False
    return True


def _should_patch_range() -> bool:
    """Return True if we should patch ``builtins.range`` for transparent epoch resume.

    Disable with ``HEGM_DISABLE_EPOCH_RANGE_PATCH=1``.
    Uses AST to find training loops; only patches range() at those sites.
    Set HEGM_EPOCH_RANGE_USE_AST=0 for legacy first-matching-call behavior.
    """
    disable_env = os.environ.get("HEGM_DISABLE_EPOCH_RANGE_PATCH", "").strip().lower()
    if disable_env in {"1", "true", "yes"}:
        log(
            "HEGM_DISABLE_EPOCH_RANGE_PATCH set; leaving range() unmodified.",
            level="INFO",
        )
        return False
    return True


_original_range = _builtins.range

# Fallback to "first matching call" when AST cannot determine (e.g. exec, -c).
_USE_AST_FOR_RANGE_PATCH = os.environ.get("HEGM_EPOCH_RANGE_USE_AST", "1").strip().lower() in ("1", "true", "yes")


def _range_proxy(*args, **kwargs):
    """Proxy for range() that transparently resumes epoch loops.

    Uses AST to find training loops (for X in range(...): ... for Y, Z in enumerate(...)).
    Only patches range() calls at detected training-loop sites.
    Set HEGM_EPOCH_RANGE_USE_AST=0 for legacy first-matching-call behavior (e.g. exec/-c).
    """
    if len(args) != 1 or len(kwargs) != 0:
        return _original_range(*args, **kwargs)
    stop = args[0]
    try:
        stop_int = int(stop)
    except (TypeError, ValueError):
        return _original_range(*args, **kwargs)

    try:
        from hegm.backends import pytorch as _pt
    except ImportError:
        return _original_range(*args, **kwargs)

    use_ast = _USE_AST_FOR_RANGE_PATCH
    should_patch = False

    if use_ast:
        try:
            frame = inspect.currentframe()
            if frame is not None and frame.f_back is not None:
                caller = frame.f_back
                filename = caller.f_code.co_filename
                lineno = caller.f_lineno
                from hegm._ast_scan import is_training_loop_site

                should_patch = is_training_loop_site(filename, lineno)
        except Exception:
            pass
    else:
        # Legacy: consider all range(stop) for patch (first matching wins in _maybe_patch_range)
        should_patch = True

    if not should_patch:
        return _original_range(*args, **kwargs)

    patched = _pt._maybe_patch_range(stop_int)
    if patched is not None:
        return patched
    return _original_range(*args, **kwargs)


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

if _should_patch_iter():
    _builtins.iter = _iter_proxy  # type: ignore[assignment]

if _should_patch_range():
    _builtins.range = _range_proxy  # type: ignore[assignment]

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


def set_gradient_accumulation_steps(n: int) -> None:
    """Set batches per optimizer step for Hugging Face Trainer.

    HeGM counts optimizer steps; the Trainer consumes
    gradient_accumulation_steps batches per step. Call before trainer.train()::

        hegm.set_gradient_accumulation_steps(training_args.gradient_accumulation_steps)
        trainer.train()
    """
    try:
        from hegm.backends import pytorch as _pt
        _pt.set_gradient_accumulation_steps(n)
    except ImportError:
        pass


def set_steps_per_epoch(n: int) -> None:
    """Enable epoch-mode checkpointing for transparent epoch-loop resume.

    When set, checkpoints store (epoch, step_in_epoch). On resume,
    ``start_epoch()`` returns the saved epoch and ``enumerate(DataLoader)``
    skips only batches within the current epoch.

    Call before the epoch loop::

        hegm.set_steps_per_epoch(len(dataloader))
        for epoch in range(hegm.start_epoch(), num_epochs):
            for step, batch in enumerate(dataloader):
                ...
    """
    try:
        from hegm.backends import pytorch as _pt
        _pt.set_steps_per_epoch(n)
    except ImportError:
        pass


def find_training_loop_sites(source_or_path):
    """Find range() call sites that are training loops (epoch + enumerate pattern).

    Parses Python source and returns a list of :class:`TrainingLoopSite` with
    lineno, col_offset, loop_var for each detected ``for X in range(...): ...
    for Y, Z in enumerate(...)`` pattern.

    Useful for debugging or tooling. Import the full API::

        from hegm._ast_scan import find_training_loop_sites, TrainingLoopSite
    """
    from hegm._ast_scan import find_training_loop_sites as _find

    return _find(source_or_path)


def start_epoch() -> int:
    """Return the epoch index to start (or resume) from.

    On a fresh run returns 0. After checkpoint restore returns the saved
    epoch so the epoch loop can resume transparently.
    """
    try:
        from hegm.backends import pytorch as _pt
        return _pt.start_epoch()
    except ImportError:
        return 0


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)
