"""HuggingFace Transformers integration for Tyko."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

try:
    from transformers import TrainerCallback as _Base  # type: ignore[import-not-found]
except ImportError:
    _Base = object  # type: ignore[assignment, misc]

if TYPE_CHECKING:
    from transformers import TrainerControl, TrainerState, TrainingArguments  # type: ignore[import-not-found]

    from tyko._types import Run

# Metrics to skip — cumulative counters, not time-series values.
_SKIP_METRICS: frozenset[str] = frozenset(
    {
        "total_flos",
    }
)

# Curated subset of TrainingArguments fields useful as run params.
# args.to_dict() returns 100+ fields; most are internal config.
_TRAINING_ARG_PARAMS: frozenset[str] = frozenset(
    {
        "learning_rate",
        "num_train_epochs",
        "per_device_train_batch_size",
        "per_device_eval_batch_size",
        "gradient_accumulation_steps",
        "weight_decay",
        "adam_beta1",
        "adam_beta2",
        "adam_epsilon",
        "max_grad_norm",
        "lr_scheduler_type",
        "warmup_ratio",
        "warmup_steps",
        "seed",
        "fp16",
        "bf16",
        "optim",
        "max_steps",
    }
)


def _rewrite_logs(logs: dict[str, Any]) -> dict[str, float | int]:
    """Rewrite HF Trainer log keys to use slash separators.

    Converts underscore prefixes to slash grouping for the Tyko dashboard:
        eval_loss  -> eval/loss
        test_loss  -> test/loss
        loss       -> train/loss
    """
    rewritten: dict[str, float | int] = {}
    for key, value in logs.items():
        if key in _SKIP_METRICS:
            continue
        if not isinstance(value, (int, float)):
            continue
        if key.startswith("eval_"):
            rewritten[f"eval/{key[5:]}"] = value
        elif key.startswith("test_"):
            rewritten[f"test/{key[5:]}"] = value
        else:
            rewritten[f"train/{key}"] = value
    return rewritten


class TykoCallback(_Base):
    """A TrainerCallback that logs metrics and training arguments to Tyko.

    Usage::

        from tyko import TykoClient
        from tyko.integrations import TykoCallback

        client = TykoClient()
        with client.start_run(project="my-project") as run:
            trainer = Trainer(
                model=model,
                args=training_args,
                train_dataset=train_dataset,
                callbacks=[TykoCallback(run)],
            )
            trainer.train()

    Args:
        run: An active Tyko Run instance (from ``client.start_run()``).
        log_training_args: If True (default), log a curated subset of
            TrainingArguments as run params when training begins.
    """

    def __init__(self, run: Run, *, log_training_args: bool = True) -> None:
        self._run = run
        self._log_training_args = log_training_args

    def on_train_begin(
        self,
        args: TrainingArguments,
        state: TrainerState,
        control: TrainerControl,
        **kwargs: Any,
    ) -> None:
        if not state.is_world_process_zero:
            return
        if self._log_training_args:
            args_dict = args.to_dict()
            params = {k: v for k, v in args_dict.items() if k in _TRAINING_ARG_PARAMS}
            if params:
                self._run.params.update(params)

    def on_log(
        self,
        args: TrainingArguments,
        state: TrainerState,
        control: TrainerControl,
        logs: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        if not state.is_world_process_zero:
            return
        if logs is None:
            return
        metrics = _rewrite_logs(logs)
        if metrics:
            self._run.log(metrics)
