# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Forward only the config arguments the installed TRL still accepts.

The `Unsloth<X>Config.__init__` generated by `unsloth/models/rl.py` ends in a
`**kwargs` catch-all splatted into `super().__init__()`, so an argument TRL has
since retired crashes the base dataclass (`max_prompt_length` was removed from
`GRPOConfig` in 0.28.0, and fourteen `DPOConfig` fields in 0.29.0), breaking
pinned scripts that were correct when written. Stdlib only, so the generated
cache file can import it behind a try/except without pulling in torch.

Three outcomes per unrecognised argument: TRL documented a rename and the new
name exists, so the value is moved across; TRL removed the field, so it is
dropped and reported by name with TRL's advice; anything else is dropped with
a warning.

That last case means `GRPOConfig(learnign_rate = 3e-4)` now warns instead of
raising `TypeError`. The trade is accepted: a typo is still surfaced loudly,
whereas the crash it replaces hits correct code on a TRL upgrade and cannot be
fixed without editing library code.
"""

import dataclasses
import inspect

__all__ = [
    "filter_config_init_kwargs",
    "TRL_CONFIG_RENAMES",
    "TRL_REMOVED_FIELD_ADVICE",
]


# Straight renames from TRL's own deprecation warnings. Only entries where the
# old value can be handed to the new name unchanged belong here, so anything
# needing a different value or setting goes in the advice table below.
TRL_CONFIG_RENAMES = {
    # Deprecated in TRL 0.26.0 (DPO) / 0.27.0 (GRPO), removed in 0.28.0.
    "use_liger_loss": "use_liger_kernel",
    # Deprecated in TRL 0.27.0, removed in 0.28.0.
    "vllm_guided_decoding_regex": "vllm_structured_outputs_regex",
    # Deprecated in TRL 0.26.0, removed in 0.27.0.
    "wandb_log_unique_prompts": "log_unique_prompts",
}


# What TRL tells users to do instead, keyed on the retired argument name. Only
# consulted for arguments the installed config rejects, so an entry is harmless
# for a config that still accepts the field (`max_completion_length` is retired
# on `DPOConfig` but current on `GRPOConfig`).
TRL_REMOVED_FIELD_ADVICE = {
    # GRPOConfig, removed in TRL 0.28.0.
    "max_prompt_length": ("filter overlong prompts out of your dataset before training instead"),
    # DPOConfig, removed in TRL 0.29.0.
    "max_completion_length": "use `max_length` to cap total sample length instead",
    "base_model_attribute_name": "the base model is now retrieved via `get_decoder`",
    "force_use_ref_model": "no longer needed, pass `ref_model` and it is used automatically",
    "generate_during_eval": "use a `TrainerCallback` instead",
    "label_pad_token_id": "this value is no longer configurable",
    "model_adapter_name": "only the default adapter is supported now",
    "ref_adapter_name": "the trainer handles the reference adapter itself now",
    "ref_model_init_kwargs": ("build the reference model yourself and pass it as `ref_model`"),
    "reference_free": "use `CPOTrainer` for a reference-free objective",
    "rpo_alpha": ("add 'sft' to `loss_type` and set its weight in `loss_weights` instead"),
    "tools": "pass tools through the dataset instead",
    "use_logits_to_keep": "the DPO trainer no longer uses this setting",
    "padding_value": "this value is no longer configurable",
}


# Memoised on the class object, which is module level and outlives the process.
_ACCEPTED_CACHE = {}


def _accepted_parameters(config_class):
    """Names `config_class.__init__` accepts, and whether it takes `**kwargs`.

    A config with its own `**kwargs` may forward anywhere, so it is reported as
    accepting everything and the caller forwards unfiltered.
    """
    try:
        cached = _ACCEPTED_CACHE.get(config_class)
    except TypeError:  # unhashable, so uncacheable; fall through and recompute
        cached = None
    if cached is not None:
        return cached

    names = set()
    takes_var_keyword = False
    try:
        parameters = inspect.signature(config_class.__init__).parameters
    except (TypeError, ValueError):
        parameters = None

    if parameters is not None:
        for name, parameter in parameters.items():
            if name == "self":
                continue
            if parameter.kind is inspect.Parameter.VAR_KEYWORD:
                takes_var_keyword = True
            elif parameter.kind is not inspect.Parameter.VAR_POSITIONAL:
                names.add(name)
    elif dataclasses.is_dataclass(config_class):
        # Unreadable `__init__`: the field list is the next best description.
        names = {f.name for f in dataclasses.fields(config_class) if f.init}
    else:
        # Nothing to go on. Forward everything rather than guess.
        takes_var_keyword = True

    result = (frozenset(names), takes_var_keyword)
    try:
        _ACCEPTED_CACHE[config_class] = result
    except TypeError:  # unhashable class, vanishingly rare but not fatal
        pass
    return result


_MISSING = object()


def _field_default(config_class, name):
    """The declared default for `name`, or `_MISSING` if it cannot be read."""
    try:
        if dataclasses.is_dataclass(config_class):
            for field in dataclasses.fields(config_class):
                if field.name != name:
                    continue
                if field.default is not dataclasses.MISSING:
                    return field.default
                if field.default_factory is not dataclasses.MISSING:
                    return field.default_factory()
                return _MISSING
        parameter = inspect.signature(config_class.__init__).parameters.get(name)
        if parameter is not None and parameter.default is not inspect.Parameter.empty:
            return parameter.default
    except Exception:
        pass
    return _MISSING


def _is_untouched(config_class, name, value):
    """True if `value` is indistinguishable from `name`'s declared default."""
    default = _field_default(config_class, name)
    if default is _MISSING:
        return False
    try:
        return bool(default == value)
    except Exception:
        return default is value


def _default_notifier(message):
    print(message)


def filter_config_init_kwargs(
    config_class,
    kwargs,
    notify = None,
):
    """Return `kwargs` reduced to what `config_class.__init__` will accept.

    Renames are applied where TRL documented one; everything else the config
    rejects is dropped. Each decision is reported through `notify` (`print` by
    default, which shows up reliably in a notebook).
    """
    if not kwargs:
        return kwargs

    accepted, takes_var_keyword = _accepted_parameters(config_class)
    if takes_var_keyword:
        return kwargs

    if notify is None:
        notify = _default_notifier
    config_name = getattr(config_class, "__name__", str(config_class))

    # Two passes so the outcome does not depend on dict ordering: take what the
    # config accepts, then rename or drop the leftovers.
    forwarded = {key: value for key, value in kwargs.items() if key in accepted}

    for key, value in kwargs.items():
        if key in accepted:
            continue

        renamed = TRL_CONFIG_RENAMES.get(key)
        if renamed is not None and renamed in accepted:
            # The new name is a mirrored parameter, so it is already in the dict
            # carrying either the caller's value or the class default. Only
            # overwrite the default: an explicitly set new name wins.
            existing = kwargs.get(renamed, _MISSING)
            if existing is _MISSING or _is_untouched(config_class, renamed, existing):
                forwarded[renamed] = value
                notify(
                    f"Unsloth: TRL renamed `{key}` to `{renamed}`. Forwarding your "
                    f"value to `{renamed}` - update your code when convenient."
                )
            else:
                notify(
                    f"Unsloth: `{key}` was renamed to `{renamed}` by TRL and this "
                    f"{config_name} accepts only the new name. You set both, so "
                    f"`{key}` is ignored and your `{renamed}` is kept."
                )
            continue

        advice = TRL_REMOVED_FIELD_ADVICE.get(key)
        if advice:
            notify(
                f"Unsloth: `{key}` is not supported by the installed TRL's "
                f"{config_name} and will be IGNORED - {advice}."
            )
        else:
            notify(
                f"Unsloth: `{key}` is not a valid {config_name} argument for the "
                f"installed TRL and will be IGNORED. Check the spelling, or your "
                f"TRL version if this argument used to work."
            )

    return forwarded
