# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.

import inspect
from typing import Optional, Callable, List, Dict, Iterator, Tuple, Any
from itertools import islice
import math
from tqdm import tqdm
from abc import ABC

import torch
import datasets
from torch.utils.data import DataLoader, Dataset, DistributedSampler, TensorDataset

from cosmos_rl.dispatcher.data.packer.base import BaseDataPacker
from cosmos_rl.policy.config import Config
from cosmos_rl.dispatcher.data import (
    CosmosDataset,
    RLDataset,
    RLPayload,
    CosmosValidationDataset,
)
from cosmos_rl.dispatcher.data import IdxAndRLPayload
from cosmos_rl.dispatcher.command import PolicyToRolloutUnicastCommand
from cosmos_rl.utils.checkpoint import CheckpointMananger
from cosmos_rl.utils.logging import logger
from cosmos_rl.utils.util import split_train_n_val_dataset


class DataFetcherBase(ABC):
    """
    DataFetcherBase is the base class for all data fetchers.
    """

    def __init__(
        self,
        config: Config,
        data_packer: BaseDataPacker,
        val_data_packer: BaseDataPacker,
        dataset: Optional[Callable[[Config], Dataset]] = None,
        val_dataset: Optional[Callable[[Config], Dataset]] = None,
        is_rl: bool = True,
    ):
        self.config = config
        self.data_packer = data_packer
        self.val_data_packer = val_data_packer
        self.dataset = dataset
        self.val_dataset = val_dataset
        self.is_rl = is_rl

    def load_dataset(self):
        if self.dataset is not None and isinstance(self.dataset, Callable):
            self.dataset = self.dataset(self.config)
        if self.val_dataset is not None and isinstance(self.val_dataset, Callable):
            self.val_dataset = self.val_dataset(self.config)

    def query_reference_answer(
        self, prompt_idx: int, dataset_type: str = "train"
    ) -> Any:
        """
        Query the reference answer from the dataset based on the prompt index.
        Args:
            prompt_idx (int): The index of the prompt in the dataset.
            dataset_type (str): The type of the dataset, either "train" or "val".
        Returns:
            Any: The reference answer corresponding to the prompt index.
        """
        if self.dataset is None:
            raise ValueError("Dataset is not loaded")
        if self.config.validation.enable and self.val_dataset is None:
            raise ValueError("Validation dataset is not loaded")

        if dataset_type == "train":
            return self.dataset.train_set.get_reference_answer(prompt_idx)
        elif dataset_type == "val":
            return self.val_dataset.val_set.get_reference_answer(prompt_idx)
        else:
            raise ValueError(f"Unknown dataset type: {dataset_type}")


class ControllerDataFetcher(DataFetcherBase):
    """
    ControllerDataFetcher is responsible for fetching data from the dataset for policy and rollout.
    """

    def __init__(
        self,
        config: Config,
        dataset: Optional[Callable[[Config], Dataset]] = None,
        val_dataset: Optional[Callable[[Config], Dataset]] = None,
        sampler: Optional[Callable] = None,
        batch_sampler: Optional[Callable] = None,
        val_sampler: Optional[Callable] = None,
        val_batch_sampler: Optional[Callable] = None,
        is_rl: bool = True,
    ):
        # ControllerDataFetcher doesn't need data packer.
        super().__init__(
            config,
            None,
            None,
            dataset,
            val_dataset,
            is_rl,
        )

        self.ckpt_extra_info = {}
        self.epoch = 1
        self.remain_samples_num = -1
        self.sampler = sampler
        self.batch_sampler = batch_sampler
        self.val_sampler = val_sampler
        self.val_batch_sampler = val_batch_sampler

        # Buffers for undispatched fetched data when data_dispatch_as_rank_in_mesh is enabled.
        self.fetched_data_buffer: List = []
        self.fetched_data_buffer_for_validation: List = []
        # Dict to track the number of data fetched for each policy at current step when data_dispatch_as_rank_in_mesh is enabled.
        self.data_fetched_for_each_policy_at_step = {}

        if self.config.train.train_policy.type == "sft":
            assert (
                self.config.train.train_policy.dataloader_batch_size
            ), "[DataFetcher] dataloader_batch_size must be set for SFT policy"
            # Set n_generation to 1 for SFT policy to avoid duplicated data counting when calculating the related value.
            self.config.rollout.n_generation = 1

        # Controller should always load the dataset and dataloader.
        self.load_dataset()

    def set_rollout_global_mesh_size(self, global_mesh_size: int):
        self.rollout_global_mesh_size = global_mesh_size

    def set_policy_global_mesh_size(self, global_mesh_size: int):
        self.policy_global_mesh_size = global_mesh_size

    def load_dataset(self):
        """
        Load the dataset and dataloader for epochs.
        """
        super().load_dataset()

        remain_samples_num = 0
        if self.is_rl:
            self.rollout_batch_size = (
                self.config.train.train_policy.dataloader_batch_size
                or self.config.rollout.batch_size
            )
            if self.dataset is not None:
                assert isinstance(self.dataset, Dataset)
                self.dataset = CosmosDataset(config=self.config, train_set=self.dataset)
                logger.info(
                    "[Controller] Using provided dataset for training, dataset specification in the toml config will be ignored"
                )
            else:
                self.dataset = CosmosDataset(config=self.config)

            if (
                self.config.validation.enable
                and self.val_dataset is None
                and not self.config.validation.dataset.name
            ):
                # If validation is enabled but no val_dataset or validation dataset name is provided, split from training dataset.
                train_dataset, val_dataset = split_train_n_val_dataset(
                    self.dataset.train_set.dataset, self.config
                )
                self.dataset.train_set.dataset = train_dataset
                self.val_dataset = val_dataset

            if self.config.train.local_dataset:
                train_index_set = RLDataset(
                    TensorDataset(torch.arange(len(self.dataset.train_set))),
                    self.config,
                )
                assert len(train_index_set) == len(self.dataset.train_set)
                self.dataset.train_set = train_index_set

            remain_samples_num = (
                (
                    len(self.dataset.train_set)
                    * self.config.rollout.n_generation
                    * self.config.train.epoch
                )
                if self.dataset is not None
                else 0
            )  # Total number of samples of policy training will consume.

            if self.sampler is not None:
                logger.info("[DataFetcher] Using provided sampler for training")
                if isinstance(self.sampler, Callable):
                    train_sampler = self.sampler(
                        self.dataset.train_set,
                        num_replicas=1,
                        rank=0,
                        shuffle=self.config.train.train_policy.dataloader_shuffle,
                        drop_last=False,
                    )
                else:
                    train_sampler = self.sampler
            else:
                train_sampler = DistributedSampler(
                    self.dataset.train_set,
                    num_replicas=1,
                    rank=0,
                    shuffle=self.config.train.train_policy.dataloader_shuffle,
                    drop_last=False,
                    seed=self.config.train.train_policy.dataloader_seed,
                )
            self.train_sampler = train_sampler
            if self.batch_sampler is not None and isinstance(
                self.batch_sampler, Callable
            ):
                sig = inspect.signature(self.batch_sampler)
                kwargs = {
                    "dataset": self.dataset.train_set.dataset,
                    "num_replicas": 1,
                    "rank": 0,
                    "num_workers": self.config.train.train_policy.dataloader_num_workers,
                    "config": self.config,
                    "sampler": self.train_sampler,
                    "batch_size": self.rollout_batch_size,
                    "drop_last": False,
                }
                # Filter kwargs to only those the function accepts
                filtered = {k: v for k, v in kwargs.items() if k in sig.parameters}
                self.batch_sampler = self.batch_sampler(**filtered)
            if self.config.train.resume:
                try:
                    # If resuming, disable the weight sync check flag for rollout to compare the received weight with the reference weight.
                    PolicyToRolloutUnicastCommand._do_weight_sync_check_flag = False
                    self.ckpt_manager = CheckpointMananger(self.config)
                    self.ckpt_extra_info = (
                        self.ckpt_manager.load_extra_info_from_checkpoint()
                    )
                    remain_samples_num = self.ckpt_extra_info.get(
                        "remain_samples_num", remain_samples_num
                    )
                    self.epoch = (
                        self.config.train.epoch
                        - (
                            math.ceil(
                                remain_samples_num
                                / (
                                    len(self.dataset.train_set)
                                    * self.config.rollout.n_generation
                                )
                            )
                        )
                        + 1
                    )
                    logger.info(
                        f"[DataFetcher] Resuming from checkpoint, current epoch: {self.epoch}, remaining samples: {remain_samples_num}"
                    )

                    train_dataloader_bias = max(
                        0,
                        len(self.dataset.train_set)
                        - (
                            (
                                math.ceil(
                                    remain_samples_num
                                    / self.config.rollout.n_generation
                                )
                            )
                            % len(self.dataset.train_set)
                        ),
                    )
                    logger.info(
                        f"[DataFetcher] Loaded extra info from checkpoint: {self.ckpt_extra_info}"
                    )
                    from cosmos_rl.policy.trainer.sampler import SkippingSampler

                    if hasattr(self.train_sampler, "set_epoch"):
                        # Here the epoch from 1 to total epoch count, not start from 0
                        self.train_sampler.set_epoch(self.epoch)

                    self.train_sampler = SkippingSampler(
                        base_sampler=self.train_sampler,
                        skip_samples=train_dataloader_bias
                        // (
                            len(list(islice(iter(self.train_sampler), 1))[0])
                            if isinstance(
                                list(islice(iter(self.train_sampler), 1))[0], list
                            )
                            else 1
                        ),
                    )
                    if self.batch_sampler is not None:
                        self.batch_sampler = SkippingSampler(
                            base_sampler=self.batch_sampler,
                            skip_samples=train_dataloader_bias
                            // (
                                len(list(islice(iter(self.batch_sampler), 1))[0])
                                if isinstance(
                                    list(islice(iter(self.batch_sampler), 1))[0], list
                                )
                                else 1
                            ),
                        )
                except Exception as e:
                    import traceback

                    traceback.print_exc()
                    logger.error(
                        f"[DataFetcher] Failed to load checkpoint extra info: {e}. Please check the checkpoint path and config."
                    )
            if hasattr(self.dataset.train_set.dataset, "data_loader"):
                # Use custom data loader if provided by dataset
                self.train_dataloader = self.dataset.train_set.dataset.data_loader
            elif self.batch_sampler is not None:
                logger.info(
                    "[DataFetcher] Using custom batch Sampler that yields list of indices for training dataset."
                )
                self.train_dataloader = DataLoader(
                    self.dataset.train_set,
                    num_workers=self.config.train.train_policy.dataloader_num_workers,
                    prefetch_factor=self.config.train.train_policy.dataloader_prefetch_factor,
                    collate_fn=RLPayload.collate_fn,
                    batch_sampler=self.batch_sampler,
                )
            else:
                self.train_dataloader = DataLoader(
                    self.dataset.train_set,
                    batch_size=self.rollout_batch_size,
                    shuffle=False,
                    num_workers=self.config.train.train_policy.dataloader_num_workers,
                    prefetch_factor=self.config.train.train_policy.dataloader_prefetch_factor,
                    collate_fn=RLPayload.collate_fn,
                    sampler=self.train_sampler,
                )
            self.train_dataloader_iter = iter(self.train_dataloader)

            if self.config.validation.enable:
                self.val_batch_size = (
                    self.config.train.train_policy.dataloader_batch_size
                    or self.config.validation.batch_size
                    or self.rollout_batch_size
                )
                assert (
                    self.val_batch_size > 0
                ), "[DataFetcher] val_batch_size should be greater than 0."
                if self.val_dataset is not None:
                    assert isinstance(self.val_dataset, Dataset) or isinstance(
                        self.val_dataset, datasets.arrow_dataset.Dataset
                    )
                    self.val_dataset = CosmosValidationDataset(
                        config=self.config,
                        val_set=self.val_dataset,
                    )
                    logger.info(
                        "[DataFetcher] Using provided validation dataset for validation, dataset specification in the toml config will be ignored"
                    )
                else:
                    self.val_dataset = CosmosValidationDataset(config=self.config)
                if self.config.train.local_dataset:
                    val_index_set = RLDataset(
                        TensorDataset(torch.arange(len(self.val_dataset.val_set))),
                        self.config,
                    )
                    assert len(val_index_set) == len(self.val_dataset.val_set)
                    self.val_dataset.val_set = val_index_set
                if self.val_sampler is not None:
                    logger.info("[DataFetcher] Using provided sampler for validation")
                    if isinstance(self.val_sampler, Callable):
                        self.val_sampler = self.val_sampler(
                            self.val_dataset.val_set,
                            num_replicas=1,
                            rank=0,
                            shuffle=False,
                            drop_last=False,
                        )

                if hasattr(self.val_dataset.val_set.dataset, "data_loader"):
                    # Use custom data loader if provided by dataset
                    self.val_dataloader = self.val_dataset.val_set.dataset.data_loader
                elif self.val_batch_sampler is not None:
                    logger.info(
                        "[DataFetcher] Using custom batch Sampler that yields list of indices for validation dataset."
                    )
                    if isinstance(self.val_batch_sampler, Callable):
                        sig = inspect.signature(self.val_batch_sampler)
                        kwargs = {
                            "dataset": self.val_dataset.val_set.dataset,
                            "num_replicas": 1,
                            "rank": 0,
                            "num_workers": self.config.train.train_policy.dataloader_num_workers,
                            "config": self.config,
                            "sampler": self.val_sampler
                            if self.val_sampler is not None
                            else DistributedSampler(
                                self.val_dataset.val_set,
                                num_replicas=1,
                                rank=0,
                                shuffle=False,
                                drop_last=False,
                            ),
                            "batch_size": self.val_batch_size,
                            "drop_last": False,
                        }
                        # Filter kwargs to only those the function accepts
                        filtered = {
                            k: v for k, v in kwargs.items() if k in sig.parameters
                        }
                        self.val_batch_sampler = self.val_batch_sampler(**filtered)
                        self.val_dataloader = DataLoader(
                            self.val_dataset.val_set,
                            num_workers=self.config.train.train_policy.dataloader_num_workers,
                            prefetch_factor=self.config.train.train_policy.dataloader_prefetch_factor,
                            collate_fn=RLPayload.collate_fn,
                            batch_sampler=self.val_batch_sampler,
                        )
                else:
                    self.val_dataloader = DataLoader(
                        self.val_dataset.val_set,
                        batch_size=self.val_batch_size,
                        shuffle=False,
                        num_workers=self.config.train.train_policy.dataloader_num_workers,
                        prefetch_factor=self.config.train.train_policy.dataloader_prefetch_factor,
                        collate_fn=RLPayload.collate_fn,
                        sampler=self.val_sampler,
                    )
            else:
                self.val_dataset = None
                self.val_dataloader = None
        else:
            self.val_dataset = None
            self.val_dataloader = None

        # validation
        self.val_datasize: Optional[int] = (
            0 if self.val_dataset is None else len(self.val_dataset.val_set)
        )
        self.val_iters: Dict[int, Iterator] = {}
        self.activated_val_iter: Optional[Iterator] = None
        self.activated_val_tqdm: Optional[tqdm] = None

        self.remain_samples_num = remain_samples_num

    def get_batched_prompt(
        self,
        n: int,
        validation_step: Optional[int] = None,
        rank_in_mesh: Optional[int] = None,
        weight_version: Optional[int] = None,
    ) -> Tuple[List[RLPayload], bool]:
        if weight_version is not None:
            self.data_fetched_for_each_policy_at_step.setdefault(weight_version, {})
        prompt_batch_per_replica = math.ceil(
            self.config.train.train_batch_per_replica / self.config.rollout.n_generation
        )
        add_answer = (
            self.config.rollout.multi_turn_config.enable
            or not self.config.train.local_dataset
        )
        # query n prompts from the dataset [idx, payload]
        payloads_list: List[RLPayload] = []
        is_end = False

        is_validation = validation_step is not None
        weight_version = None if is_validation else weight_version

        if is_validation:
            iterator = self.validation_get_dataloader(validation_step)
            batch_size = self.val_batch_size
            fetched_data_buffer = self.fetched_data_buffer_for_validation
        else:
            iterator = self.train_dataloader_iter
            batch_size = self.rollout_batch_size
            fetched_data_buffer = self.fetched_data_buffer

        def _next_payload(
            iterator, add_answer: bool
        ) -> tuple[List[int], List[RLPayload]]:
            idxs, payloads = next(iterator)
            assert len(idxs) <= batch_size
            assert len(payloads) <= batch_size
            assert len(idxs) == len(payloads)
            updated_payloads: List[RLPayload] = []
            for idx, payload in zip(idxs, payloads):
                if add_answer:
                    if is_validation:
                        payload.reference_answer = (
                            self.val_dataset.val_set.get_reference_answer(idx)
                        )
                    else:
                        payload.reference_answer = (
                            self.dataset.train_set.get_reference_answer(idx)
                        )
                updated_payloads.append(payload)
            return idxs, updated_payloads

        if self.config.train.train_policy.data_dispatch_as_rank_in_mesh:
            """
            First use the fetched_data_buffer to fill the payloads_list.
            Then fetch new data from the iterator until we have n payloads or the iterator is exhausted.
            """
            assert (
                rank_in_mesh is not None
            ), "rank_in_mesh should not be None when data_dispatch_as_rank_in_mesh is enabled"
            while n - len(payloads_list) > 0:
                found = False
                for index, data in enumerate(fetched_data_buffer):
                    if data[0] % self.rollout_global_mesh_size == rank_in_mesh and (
                        weight_version is None
                        or self.data_fetched_for_each_policy_at_step[
                            weight_version
                        ].get(data[0] % self.policy_global_mesh_size, 0)
                        < prompt_batch_per_replica
                    ):
                        payloads_list.append(data[1])
                        if weight_version is not None:
                            self.data_fetched_for_each_policy_at_step[weight_version][
                                data[0] % self.policy_global_mesh_size
                            ] = (
                                self.data_fetched_for_each_policy_at_step[
                                    weight_version
                                ].get(data[0] % self.policy_global_mesh_size, 0)
                                + 1
                            )
                        found = True
                        break
                if found:
                    del fetched_data_buffer[index]
                else:
                    break

        while n - len(payloads_list) > 0:
            for _ in range(math.ceil(n / batch_size)):
                payload: RLPayload | None = None
                try:
                    idxs, payloads = _next_payload(iterator, add_answer)
                except StopIteration:
                    if not is_validation:
                        self.epoch += 1
                        if hasattr(self.train_sampler, "set_epoch"):
                            # Here the epoch from 1 to total epoch count, not start from 0
                            self.train_sampler.set_epoch(self.epoch)
                        if self.epoch <= self.config.train.epoch:
                            logger.info(f"[Controller] Epoch {self.epoch} start.")
                            iterator = iter(self.train_dataloader)
                            self.train_dataloader_iter = iterator

                            idxs, payloads = _next_payload(iterator, add_answer)
                        else:
                            if self.epoch == self.config.train.epoch + 1:
                                # We only log this all finished information once.
                                logger.info(
                                    "[Controller] All epochs finished fetching rollout prompts, wait for rollouts generation and training to complete."
                                )
                            is_end = True
                            break
                    else:
                        is_end = True
                        break
                assert len(idxs) == len(payloads)
                for idx, payload in zip(idxs, payloads):
                    idx = idx.item() if isinstance(idx, torch.Tensor) else idx
                    if self.config.train.local_dataset:
                        # If local dataset is enabled, we set prompt to None. And rollout worker will query
                        # the prompt from local dataset.
                        payload.prompt = None
                        payload.conversation = None
                        if not self.config.rollout.multi_turn_config.enable:
                            # For non-multi-turn rollout, we set reference answer to None.
                            payload.reference_answer = None
                    if self.config.train.train_policy.data_dispatch_as_rank_in_mesh:
                        assert (
                            rank_in_mesh is not None
                        ), "rank_in_mesh should not be None when data_dispatch_as_rank_in_mesh is enabled"
                        if (
                            idx % self.rollout_global_mesh_size == rank_in_mesh
                            and (
                                weight_version is None
                                or self.data_fetched_for_each_policy_at_step[
                                    weight_version
                                ].get(idx % self.policy_global_mesh_size, 0)
                                < prompt_batch_per_replica
                            )
                            and len(payloads_list) < n
                        ):
                            payloads_list.append(payload)
                            if weight_version is not None:
                                self.data_fetched_for_each_policy_at_step[
                                    weight_version
                                ][idx % self.policy_global_mesh_size] = (
                                    self.data_fetched_for_each_policy_at_step[
                                        weight_version
                                    ].get(idx % self.policy_global_mesh_size, 0)
                                    + 1
                                )
                        else:
                            # For data_dispatch_as_rank_in_mesh, we store the fetched data into the buffer if not suitable for current rank_in_mesh.
                            fetched_data_buffer.append((idx, payload))
                    else:
                        payloads_list.append(payload)
            if (
                is_end
                or not self.config.train.train_policy.data_dispatch_as_rank_in_mesh
            ):
                break
        # For data_dispatch_as_rank_in_mesh, we only allow is_end to be True when there is no more data suitable for current rank_in_mesh.
        is_end = is_end and (
            len(fetched_data_buffer) == 0
            or not self.config.train.train_policy.data_dispatch_as_rank_in_mesh
        )
        if is_validation:
            self.fetched_data_buffer_for_validation = fetched_data_buffer
        else:
            self.fetched_data_buffer = fetched_data_buffer
        return payloads_list, is_end

    def validation_activate_dataloader(self, validation_step: int):
        if validation_step not in self.val_iters:
            logger.info(
                f"[DataFetcher] Activating validation dataloader for step {validation_step}, with length {(self.val_datasize or len(self.val_dataloader))}"
            )
            self.val_iters[validation_step] = iter(self.val_dataloader)
            self.activated_val_iter = self.val_iters[validation_step]
            self.activated_val_tqdm = tqdm(
                desc="validation",
                total=(self.val_datasize or len(self.val_dataloader)),
            )

    def validation_get_dataloader(
        self, validation_step: Optional[int] = None
    ) -> Iterator:
        if validation_step is None:
            return self.activated_val_iter
        else:
            return self.val_iters[validation_step]

    def clear_validation_status(self):
        self.activated_val_iter = None
        if self.activated_val_tqdm is not None:
            self.activated_val_tqdm.clear()
        self.activated_val_tqdm = None


class WorkerDataFetcher(DataFetcherBase):
    """
    WorkerDataFetcher is responsible for fetching data locally for policy and rollout, according to the index returned by the controller.
    WorkerDataFetcher is much more simpler than ControllerDataFetcher, because it only supports query data by index.
    """

    def __init__(
        self,
        config: Config,
        data_packer: BaseDataPacker,
        val_data_packer: BaseDataPacker,
        dataset: Optional[Callable[[Config], Dataset]] = None,
        val_dataset: Optional[Callable[[Config], Dataset]] = None,
        is_rl: bool = True,
    ):
        super().__init__(
            config,
            data_packer,
            val_data_packer,
            dataset,
            val_dataset,
            is_rl,
        )

        if self.config.train.local_dataset:
            self.load_dataset()

    def load_dataset(self):
        super().load_dataset()

        if self.dataset is not None:
            assert isinstance(self.dataset, Dataset)
            self.dataset = CosmosDataset(config=self.config, train_set=self.dataset)
            logger.info(
                "[DataFetcher] Using provided dataset for training, dataset specification in the toml config will be ignored"
            )
        else:
            self.dataset = CosmosDataset(config=self.config)

        if self.config.validation.enable:
            if self.val_dataset is not None:
                assert isinstance(self.val_dataset, Dataset)
                self.val_dataset = CosmosValidationDataset(
                    config=self.config,
                    val_set=self.val_dataset,
                )
                logger.info(
                    "[DataFetcher] Using provided validation dataset for validation, dataset specification in the toml config will be ignored"
                )
            elif not self.config.validation.dataset.name:
                # If validation is enabled but no val_dataset or validation dataset name is provided, split from training dataset.
                train_dataset, val_dataset = split_train_n_val_dataset(
                    self.dataset.train_set.dataset, self.config
                )
                self.dataset.train_set.dataset = train_dataset
                self.val_dataset = val_dataset
                self.val_dataset = CosmosValidationDataset(
                    config=self.config,
                    val_set=self.val_dataset,
                )
            else:
                self.val_dataset = CosmosValidationDataset(config=self.config)

    def get_payload_by_index(
        self, index: int, is_validation: bool = False, attr: str = "prompt"
    ) -> RLPayload:
        row: IdxAndRLPayload = None
        if is_validation:
            if self.val_dataset is None or not self.config.validation.enable:
                raise ValueError(
                    "[DataFetcher] Validation dataset is not loaded or validation is not enabled"
                )
            row = self.val_dataset.val_set[index]
        else:
            if self.dataset is None:
                raise ValueError("[DataFetcher] Local dataset is not loaded")
            row = self.dataset.train_set[index]
        return getattr(row[1], attr)
