"""
hegm.hf_trainer -- Hugging Face Trainer integration.

Provides a callback that restores TrainerState.global_step and the progress
bar after HeGM checkpoint restore (e.g. 126/300 instead of 0/300).
"""

from typing import TYPE_CHECKING, Any

from hegm._config import log

if TYPE_CHECKING:
    from transformers import TrainerCallback, TrainerState, TrainingArguments


def get_hegm_callback() -> "TrainerCallback":
    """Return a TrainerCallback that restores progress after HeGM checkpoint.

    Normally not needed: HeGM automatically patches Trainer on import
    (see hegm.hf_integration). Use this only if you disabled automatic
    integration (HEGM_DISABLE_HF_INTEGRATION=1) or need manual control.
    """
    try:
        from tqdm.auto import tqdm
        from transformers import ProgressCallback
    except ImportError:
        raise ImportError("transformers is required for get_hegm_callback()") from 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()
