#!/usr/bin/env python
"""
Reproduce the GRPO tool-mask shape mismatch *inside* `GRPOTrainer.train()`.

This script intentionally forces a tool-call round where:
- `completion_ids` is effectively shortened by re-tokenization logic, but
- `tool_mask` is only extended (not truncated),
which leads to a mismatch at:
`trl/trainer/grpo_trainer.py`, `_compute_loss`, line with
`completion_mask * inputs["tool_mask"]`.
"""

from __future__ import annotations

import types
from typing import Any

import torch
import torch.distributed.fsdp as fsdp


# Compatibility shims for this environment.
if not hasattr(fsdp, "FSDPModule"):
    fsdp.FSDPModule = fsdp.FullyShardedDataParallel
if not hasattr(torch.backends.mps, "is_macos_or_newer"):
    torch.backends.mps.is_macos_or_newer = lambda major, minor: False

from datasets import load_dataset

import trl.trainer.grpo_trainer as grpo_mod
from trl import GRPOConfig, GRPOTrainer


def multiply_tool(a: int, b: int) -> int:
    """Multiply two integers.

    Args:
        a: First integer.
        b: Second integer.

    Returns:
        Product of `a` and `b`.
    """
    return a * b


def constant_reward(completions: list[Any], **kwargs) -> list[float]:
    del kwargs
    return [0.0] * len(completions)


def main() -> None:
    dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train[:3]")

    args = GRPOConfig(
        output_dir="/tmp/trl_bug_repro_train",
        per_device_train_batch_size=3,
        num_generations=3,
        max_completion_length=128,
        max_steps=1,
        logging_steps=1,
        report_to="none",
        use_cpu=True,
    )

    trainer = GRPOTrainer(
        model="trl-internal-testing/tiny-Qwen3MoeForCausalLM",
        reward_funcs=constant_reward,
        args=args,
        train_dataset=dataset,
        tools=[multiply_tool],
    )

    # Keep references so we can restore global/module state.
    original_generate_single_turn = trainer._generate_single_turn
    original_parse_response = grpo_mod.parse_response
    original_tool_call_loop = trainer._tool_call_loop
    original_get_per_token_logps_and_entropies = trainer._get_per_token_logps_and_entropies

    state = {"turn": 0}

    def fake_generate_single_turn(self, prompts):
        del prompts
        state["turn"] += 1
        if state["turn"] == 1:
            # First generation batch (size 3).
            prompt_ids = [[100, 101], [110, 111], [120, 121]]
            completion_ids = [
                [900, 901, 902, 903],  # will be interpreted as a tool call
                [700, 701, 702],
                [710, 711, 712],
            ]
            return prompt_ids, completion_ids, None, {}
        if state["turn"] == 2:
            # Tool-round generation for the single tool-calling sample.
            # Intentionally short prompt+completion+tool sequence to induce
            # negative `tool_length` in `_tool_call_loop`.
            prompt_ids = [[100, 101, 201, 202]]
            completion_ids = [[600, 601]]
            return prompt_ids, completion_ids, None, {}
        raise RuntimeError(f"Unexpected fake generation turn: {state['turn']}")

    def fake_parse_response(processing_class, ids):
        del processing_class
        if ids and ids[0] == 900:
            return {
                "role": "assistant",
                "content": "",
                "tool_calls": [
                    {
                        "type": "function",
                        "function": {
                            "name": "multiply_tool",
                            "arguments": {"a": 3, "b": 4},
                        },
                    }
                ],
            }
        return {"role": "assistant", "content": "ok"}

    def wrapped_tool_call_loop(*args, **kwargs):
        out = original_tool_call_loop(*args, **kwargs)
        tool_mask, _completions, completion_ids, _logprobs, *_ = out
        print("DEBUG completion lengths:", [len(x) for x in completion_ids])
        print("DEBUG tool_mask lengths :", [len(x) for x in tool_mask])
        return out

    def fake_get_per_token_logps_and_entropies(
        self,
        model,
        input_ids,
        attention_mask,
        logits_to_keep,
        batch_size=None,
        compute_entropy=False,
        **kwargs,
    ):
        # Return synthetic tensors to keep this script focused on mask-shape behavior.
        del model, attention_mask, batch_size, kwargs
        batch = input_ids.shape[0]
        shape = (batch, logits_to_keep)
        logps = torch.zeros(shape, dtype=torch.float32, device=input_ids.device, requires_grad=True)
        entropies = torch.zeros(shape, dtype=torch.float32, device=input_ids.device) if compute_entropy else None
        return logps, entropies

    trainer._generate_single_turn = types.MethodType(fake_generate_single_turn, trainer)
    trainer._tool_call_loop = wrapped_tool_call_loop
    trainer._get_per_token_logps_and_entropies = types.MethodType(fake_get_per_token_logps_and_entropies, trainer)
    grpo_mod.parse_response = fake_parse_response

    try:
        trainer.train()
        print("\nTraining finished without tool_mask/completion_mask shape mismatch.")
    finally:
        trainer._generate_single_turn = original_generate_single_turn
        trainer._tool_call_loop = original_tool_call_loop
        trainer._get_per_token_logps_and_entropies = original_get_per_token_logps_and_entropies
        grpo_mod.parse_response = original_parse_response


if __name__ == "__main__":
    main()