from types import SimpleNamespace

import torch

import trl.experimental.online_dpo.online_dpo_trainer as module_under_test
from trl.experimental.online_dpo.online_dpo_trainer import OnlineDPOTrainer


class DummyProcessor:
    def __call__(self, text, return_tensors, padding, padding_side, add_special_tokens):
        return {"input_ids": torch.tensor([[101, 102, 103]])}


module_under_test.gather_object = lambda x: x
module_under_test.broadcast_object_list = lambda x, from_process=0: x

trainer = OnlineDPOTrainer.__new__(OnlineDPOTrainer)
trainer.accelerator = SimpleNamespace(is_main_process=True, process_index=0)
trainer.state = SimpleNamespace(global_step=0)
trainer._move_model_to_vllm = lambda: None
trainer._last_loaded_step = -1
trainer.num_generations = 2
trainer.repetition_penalty = 1.0
trainer.temperature = 1.0
trainer.top_p = 1.0
trainer.top_k = None
trainer.min_p = None
trainer.generation_config = SimpleNamespace(max_tokens=16)
trainer.args = SimpleNamespace(generation_kwargs={})
trainer.processing_class = DummyProcessor()
trainer.vllm_client = SimpleNamespace(
    generate=lambda **kwargs: {
        "completion_ids": [
            [11, 12, 13],
            [21, 22],
        ]
    }
)

completion_ids, prompt_ids = trainer._generate_vllm_server(["plain prompt", "plain prompt"])

print(completion_ids)
print(prompt_ids)

[[11], [12], [13], [21]]
[[101, 102, 103], [101, 102, 103]]

[[11, 12, 13], [21, 22]]
[[101, 102, 103], [101, 102, 103]]