"""
hegm.hf_integration -- Transparent Hugging Face Trainer integration.

When transformers.trainer is imported, patches the Trainer class so that:
- HeGMProgressCallback replaces the default ProgressCallback (resume-aware progress bar)
- gradient_accumulation_steps is set for correct batch skip on resume

No changes to user training scripts are required.
Disable with HEGM_DISABLE_HF_INTEGRATION=1.
"""

import functools
import os
from typing import Any

from hegm._config import log
from hegm._hook import install_hook


def _get_hegm_progress_callback():
    """Build HeGMProgressCallback (lazy to avoid importing transformers at module load)."""
    try:
        from tqdm.auto import tqdm
        from transformers import ProgressCallback
    except ImportError:
        return None

    class HeGMProgressCallback(ProgressCallback):
        """ProgressCallback that initializes bar at HeGM-restored step."""

        def on_train_begin(self, args: Any, state: Any, control: Any, **kwargs: Any) -> Any:
            from hegm import global_step

            step = global_step()
            if step > 0:
                state.global_step = step
                log("Restored Trainer state.global_step = %d" % step)
            if state.is_world_process_zero:
                self.training_bar = tqdm(
                    total=state.max_steps,
                    initial=state.global_step,
                    dynamic_ncols=True,
                )
            self.current_step = state.global_step
            return control

    return HeGMProgressCallback


def _apply_trainer_patches(trainer_module: Any) -> None:
    """Patch Trainer.__init__ to inject HeGM callback and set gradient accumulation."""
    Trainer = getattr(trainer_module, "Trainer", None)
    if Trainer is None:
        return

    _original_init = Trainer.__init__

    @functools.wraps(_original_init)
    def _patched_init(self: Any, *args: Any, **kwargs: Any) -> None:
        _original_init(self, *args, **kwargs)
        try:
            from hegm import set_gradient_accumulation_steps
            from hegm.hf_integration import _get_hegm_progress_callback

            args_obj = getattr(self, "args", None)
            if args_obj is not None and hasattr(args_obj, "gradient_accumulation_steps"):
                set_gradient_accumulation_steps(args_obj.gradient_accumulation_steps)

            hegm_cb = _get_hegm_progress_callback()
            if hegm_cb is not None and hasattr(self, "callback_handler"):
                callbacks = getattr(self.callback_handler, "callbacks", [])
                new_callbacks = [
                    cb for cb in callbacks
                    if type(cb).__name__ != "ProgressCallback"
                ]
                new_callbacks.append(hegm_cb())
                self.callback_handler.callbacks = new_callbacks
        except Exception as exc:
            log("HeGM HF integration skipped: %s" % exc, level="DEBUG")

    Trainer.__init__ = _patched_init  # type: ignore[method-assign]
    log("HeGM: Trainer patched for transparent checkpoint resume", level="DEBUG")


def _maybe_install_hf_hook() -> None:
    """Install import hook for transformers.trainer (no-op if transformers not installed)."""
    try:
        import importlib.util
        spec = importlib.util.find_spec("transformers")
        if spec is not None:
            install_hook("transformers.trainer", _apply_trainer_patches)
    except ImportError:
        pass


# Install hook when this module is loaded (transformers may not be imported yet).
if os.environ.get("HEGM_DISABLE_HF_INTEGRATION", "").strip().lower() not in ("1", "true", "yes"):
    _maybe_install_hf_hook()
