# Copyright 2025 The RLinf Authors.
#
# 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
#
#     https://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 ast
import json
import logging
import os
import pickle
import re
from io import BytesIO
from typing import Any, Callable, Optional, Union

import numpy as np
import pandas as pd
import torch
from omegaconf import DictConfig
from PIL import Image
from torch.utils.data import Dataset
from tqdm import tqdm
from transformers import AutoProcessor, AutoTokenizer

from rlinf.data.datasets.item import DatasetItem, SftDatasetItem
from rlinf.data.utils import batch_pad_to_fixed_len


class VLMDatasetRegistry:
    registry: dict[str, Callable[..., "VLMBaseDataset"]] = {}

    @classmethod
    def register(
        cls, name: str
    ) -> Callable[[Callable[..., "VLMBaseDataset"]], Callable[..., "VLMBaseDataset"]]:
        def decorator(klass: Callable[..., "VLMBaseDataset"]):
            cls.registry[name] = klass
            return klass

        return decorator

    @classmethod
    def create(
        cls,
        dataset_name: Optional[str],
        *,
        data_paths: Union[list[str], str],
        config: DictConfig,
        tokenizer: AutoTokenizer,
        **kwargs,
    ) -> "VLMBaseDataset":
        key = dataset_name.lower()
        dataset_class = cls.registry.get(key)
        return dataset_class(
            data_paths=data_paths, config=config, tokenizer=tokenizer, **kwargs
        )


@VLMDatasetRegistry.register("base")
class VLMBaseDataset(Dataset):
    def __init__(
        self,
        data_paths: Union[list[str], str],
        config: DictConfig,
        tokenizer: AutoTokenizer,
    ) -> None:
        super().__init__()
        self.cfg = config
        raw_paths = [data_paths] if isinstance(data_paths, str) else list(data_paths)
        # Expand directories into file lists recursively (json/jsonl/parquet)
        self.data_paths = self._expand_data_paths(raw_paths)
        self.tokenizer = tokenizer
        # Delay processor creation; only needed when use_chat_template is True
        self._processor = None

        self.system_prompt = config.data.get("system_prompt", None)
        self.use_chat_template = bool(config.data.use_chat_template)
        self.image_keys = list(config.data.image_keys or [])
        self.prompt_key = config.data.prompt_key
        self.choice_key = config.data.get("choice_key", None)
        self.answer_key = config.data.get("answer_key", None)
        self.solution_key = config.data.get("solution_key", None)
        self.max_prompt_length = int(config.data.max_prompt_length)
        self.eos_id = int(self.tokenizer.eos_token_id)

        # Loading mode
        self.lazy_loading = bool(getattr(config.data, "lazy_loading", False))

        self._records = []
        self._indices = []  # (path, fmt, row_index_or_offset)

        if self.lazy_loading:
            self._build_lazy_indices()
        else:
            self._eager_load_all()

    def __len__(self) -> int:
        return len(self._indices) if self.lazy_loading else len(self._records)

    def __getitem__(self, idx: int) -> DatasetItem:
        if self.lazy_loading:
            path, fmt, key = self._indices[idx]
            raw = self._load_single_lazy(path, fmt, key)
            return self._process_raw_record(raw, idx)
        else:
            raw = self._records[idx]
            return self._process_raw_record(raw, idx)

    # Ensure dataset is picklable for multi-process DataLoader by removing
    # unpicklable cache objects like pyarrow.ParquetFile from state.
    def __getstate__(self):
        state = self.__dict__.copy()
        # Drop heavy/unpicklable caches; they will be rebuilt on-demand in workers
        for k in ("_parquet_cache", "_parquet_df_cache", "_parquet_rg_bounds"):
            if k in state:
                state[k] = {}
        return state

    def __setstate__(self, state):
        # Restore state and ensure cache dicts exist
        self.__dict__.update(state)
        self._parquet_cache = getattr(self, "_parquet_cache", {})
        self._parquet_df_cache = getattr(self, "_parquet_df_cache", {})
        self._parquet_rg_bounds = getattr(self, "_parquet_rg_bounds", {})

    def get_image_list(self, dataitem: dict[str, Any]) -> list[Union[bytes, str, None]]:
        images: list[Union[bytes, str, None]] = []
        for k in self.image_keys:
            v = dataitem.get(k, None)
            if v is None:
                continue
            if isinstance(v, Image.Image):
                images.append(v)
            elif isinstance(v, dict) and "bytes" in v:
                images.append(v["bytes"])
            elif isinstance(v, np.ndarray) and "bytes" in v[0]:
                for p in v:
                    images.append(p["bytes"])
            else:
                images.append(v)  # path or url
        if not images:
            images = [None]
        return images

    def build_prompt_text(self, data_item: dict[str, Any]) -> str:
        # Default: prompt + optional choices rendered inline
        q = data_item.get(self.prompt_key, "")
        choices = data_item.get(self.choice_key, []) if self.choice_key else []
        if not isinstance(choices, list):
            choices = [choices]
        if choices:
            return f"{q}{choices}\n"
        return str(q)

    @staticmethod
    def process_inputs(
        processor: AutoProcessor,
        system_prompt: Optional[str],
        use_chat_template: bool,
        prompt_texts: list[str],
        images,
    ) -> tuple[str, dict[str, Any]]:
        """Shared helper to build chat-template text and processor inputs for VLM.

        This mirrors the logic in VLMBaseDataset.encode_prompt so it can be reused
        by both datasets and downstream modules (e.g., reward models).
        """
        prompt_text = prompt_texts[0]

        messages: list[dict[str, Any]] = []
        if system_prompt is not None:
            messages.append(
                {
                    "role": "system",
                    "content": [{"type": "text", "text": system_prompt}],
                }
            )

        content: list[dict[str, Any]] = []

        # Parse prompt_text for image placeholders and interleave text segments with images
        parts = re.split(r"<image>", prompt_text)
        if len(parts) == 1:
            # No placeholder found, fall back to original logic
            for _ in range(max(0, len(images))):
                content.append({"type": "image"})
            content.append({"type": "text", "text": prompt_text})
        else:
            for i, part in enumerate(parts):
                if part:
                    content.append({"type": "text", "text": part})
                if i < len(parts) - 1:
                    content.append({"type": "image"})

        messages.append({"role": "user", "content": content})

        if use_chat_template:
            rendered = processor.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        else:
            rendered = prompt_text

        images_inputs = []
        for image in images:
            image_obj = None
            if isinstance(image, Image.Image):
                image_obj = image.convert("RGB")
            if isinstance(image, (bytes, bytearray)):
                image_obj = Image.open(BytesIO(image)).convert("RGB")
            if image_obj is None:
                raise ValueError(
                    f"Unsupported image type: {type(image).__name__}. "
                    f"Expected PIL.Image.Image, bytes, or bytearray."
                )

            images_inputs.append(image_obj)

        inputs = processor(
            text=[rendered], images=images_inputs, padding=True, return_tensors="pt"
        )
        return rendered, inputs

    def encode_prompt(
        self, prompt_text: str, images
    ) -> tuple[torch.Tensor, int, Optional[str]]:
        """
        Return (token_ids[L], length, prompt_text_used). If using chat template, encode with processor.
        Subclasses may override to support alternative prompting.
        """
        if self.use_chat_template:
            if self._processor is None:
                self._processor = AutoProcessor.from_pretrained(
                    self.cfg.actor.model.model_path
                )

            rendered, inputs = self.process_inputs(
                processor=self._processor,
                system_prompt=self.system_prompt,
                use_chat_template=self.use_chat_template,
                prompt_texts=[prompt_text],
                images=images,
            )

            inputs.pop("attention_mask", None)

            if self.cfg.rollout.rollout_backend == "sglang":
                ids = inputs.pop("input_ids")
            elif self.cfg.rollout.rollout_backend == "vllm":
                inputs.pop("input_ids", None)
                ids = self._processor(
                    text=[rendered], images=None, padding=True, return_tensors="pt"
                )["input_ids"]
            else:
                raise ValueError(
                    f"Unsupported rollout backend {self.cfg.rollout.rollout_backend}"
                )
            if isinstance(ids, torch.Tensor):
                if ids.dim() == 2 and ids.size(0) == 1:
                    ids = ids.squeeze(0)
                ids = ids.to(dtype=torch.long)
            else:
                ids = torch.tensor(ids, dtype=torch.long)

            multi_modal_inputs = {}
            for k, v in inputs.items():
                multi_modal_inputs[k] = v
            return ids, int(ids.numel()), rendered, multi_modal_inputs
        else:
            # fallback: tokenizer only
            ids_list = self.tokenizer.encode(prompt_text)
            ids = torch.as_tensor(ids_list, dtype=torch.long)
            return ids, int(ids.numel()), prompt_text, {}

    def postprocess_dataset_item(
        self, item: DatasetItem, raw: dict[str, Any]
    ) -> DatasetItem:
        return item

    def _expand_data_paths(self, inputs: list[str]) -> list[str]:
        exts = {".jsonl", ".json", ".parquet"}
        files: list[str] = []
        for p in inputs:
            if os.path.isdir(p):
                for root, _, fnames in os.walk(p):
                    for fn in fnames:
                        ext = os.path.splitext(fn)[1].lower()
                        if ext in exts:
                            files.append(os.path.join(root, fn))
            else:
                files.append(p)
        files = sorted(set(files))
        return files

    def _eager_load_all(self) -> None:
        merged: list[dict[str, Any]] = []
        for path in tqdm(self.data_paths, desc="Loading dataset files", unit="file"):
            fmt = os.path.splitext(path)[1].lower()
            if fmt == ".jsonl":
                with open(path, "r", encoding="utf-8") as f:
                    merged.extend(json.loads(l) for l in f)
            elif fmt == ".json":
                with open(path, "r", encoding="utf-8") as f:
                    content = json.load(f)
                    if isinstance(content, list):
                        merged.extend(content)
                    else:
                        merged.append(content)
            elif fmt == ".parquet":
                try:
                    merged.extend(pd.read_parquet(path).to_dict(orient="records"))
                except Exception as e:
                    raise RuntimeError(f"Failed to load parquet eagerly: {path}: {e}")
            else:
                logging.warning(f"Unsupported format {fmt} for path {path}, skipping.")
        self._records = merged
        # Build indices for consistency
        self._indices = [("", "eager", i) for i in range(len(self._records))]

    def _build_parquet_rg_bounds(self, pf, path: str) -> list[tuple[int, int, int]]:
        """Build and cache row-group boundaries for a parquet file."""
        self._parquet_rg_bounds = getattr(self, "_parquet_rg_bounds", {})
        bounds = []
        offset = 0
        for g in range(pf.metadata.num_row_groups):
            rg_rows = pf.metadata.row_group(g).num_rows
            bounds.append((offset, offset + rg_rows, g))
            offset += rg_rows
        self._parquet_rg_bounds[path] = bounds
        return bounds

    def _build_lazy_indices(self) -> None:
        self._indices.clear()
        for path in tqdm(self.data_paths, desc="Loading dataset files", unit="file"):
            fmt = os.path.splitext(path)[1].lower()
            if fmt == ".jsonl":
                # index by byte offsets for each line
                offsets: list[int] = []
                with open(path, "rb") as fb:
                    pos = 0
                    for line in fb:
                        offsets.append(pos)
                        pos += len(line)
                self._indices.extend((path, "jsonl", off) for off in offsets)
            elif fmt == ".json":
                try:
                    with open(path, "r", encoding="utf-8") as f:
                        content = json.load(f)
                    if not isinstance(content, list):
                        content = [content]
                    # store the content to avoid re-reading
                    # keep perfile cache
                    self._json_cache = getattr(self, "_json_cache", {})
                    self._json_cache[path] = content
                    self._indices.extend((path, "json", i) for i in range(len(content)))
                except Exception as e:
                    raise RuntimeError(f"Failed to index json lazily: {path}: {e}")
            elif fmt == ".parquet":
                try:
                    import pyarrow.parquet as pq  # type: ignore

                    pf = pq.ParquetFile(path)
                    num_rows = pf.metadata.num_rows
                    # file handle cache
                    self._parquet_cache = getattr(self, "_parquet_cache", {})
                    self._parquet_cache[path] = pf
                    # Cache row group boundaries for correct lazy indexing
                    self._build_parquet_rg_bounds(pf, path)
                    self._indices.extend((path, "parquet", i) for i in range(num_rows))
                except Exception:
                    df = pd.read_parquet(path)
                    self._parquet_df_cache = getattr(self, "_parquet_df_cache", {})
                    self._parquet_df_cache[path] = df
                    self._indices.extend(
                        (path, "parquet_pd", i) for i in range(len(df))
                    )
            else:
                logging.warning(f"Unsupported format {fmt} for path {path}, skipping.")

    def _load_single_lazy(self, path: str, fmt: str, key: Any) -> dict[str, Any]:
        if fmt == "eager":
            return self._records[int(key)]
        if fmt == "jsonl":
            with open(path, "rb") as fb:
                fb.seek(int(key))
                line = fb.readline()
            return json.loads(line.decode("utf-8").strip())
        if fmt == "json":
            return self._json_cache[path][int(key)]  # type: ignore[attr-defined]
        if fmt == "parquet":
            # Try to use pyarrow lazily; rebuild cache if missing
            self._parquet_cache = getattr(self, "_parquet_cache", {})
            pf = self._parquet_cache.get(path)
            if pf is None:
                try:
                    import pyarrow.parquet as pq  # type: ignore

                    pf = pq.ParquetFile(path)
                    self._parquet_cache[path] = pf
                except Exception:
                    # Fall back to pandas-based cache
                    self._parquet_df_cache = getattr(self, "_parquet_df_cache", {})
                    df = self._parquet_df_cache.get(path)
                    if df is None:
                        df = pd.read_parquet(path)
                        self._parquet_df_cache[path] = df
                    return df.iloc[int(key)].to_dict()
            # Find the row group that contains the requested row
            self._parquet_rg_bounds = getattr(self, "_parquet_rg_bounds", {})
            bounds = self._parquet_rg_bounds.get(path)
            if bounds is None:
                # Rebuild bounds if cache was cleared (e.g. after unpickling)
                bounds = self._build_parquet_rg_bounds(pf, path)
            row = int(key)
            target_rg = 0
            local_idx = row
            for start, end, rg_idx in bounds:
                if start <= row < end:
                    target_rg = rg_idx
                    local_idx = row - start
                    break
            table = pf.read_row_group(target_rg, columns=None)
            df = table.to_pandas()
            return df.iloc[local_idx].to_dict()
        if fmt == "parquet_pd":
            self._parquet_df_cache = getattr(self, "_parquet_df_cache", {})
            df = self._parquet_df_cache.get(path)
            if df is None:
                df = pd.read_parquet(path)
                self._parquet_df_cache[path] = df
            return df.iloc[int(key)].to_dict()
        raise RuntimeError(f"Unknown lazy fmt {fmt}")

    def _process_raw_record(self, raw: dict[str, Any], idx: int) -> DatasetItem:
        images = self.get_image_list(raw)
        prompt_text = self.build_prompt_text(raw)
        prompt_ids, plen, rendered_text, multi_modal_inputs = self.encode_prompt(
            prompt_text, images
        )

        if plen > self.max_prompt_length:
            prompt_ids = prompt_ids[: self.max_prompt_length]
            plen = self.max_prompt_length
        prompt_ids = batch_pad_to_fixed_len(
            [prompt_ids], self.max_prompt_length, self.eos_id, left_pad=True
        )[0]

        answer_val = raw.get(self.answer_key, None) if self.answer_key else None
        solution_val = raw.get(self.solution_key, None) if self.solution_key else None
        item = DatasetItem(
            prompt=prompt_ids,
            length=plen,
            answer=str(answer_val) if answer_val is not None else None,
            idx=idx,
            image_data=images,
            prompt_text=rendered_text or prompt_text,
            solution=solution_val,
            meta=None,
            multi_modal_inputs=multi_modal_inputs,
        )
        return self.postprocess_dataset_item(item, raw)


@VLMDatasetRegistry.register("robo2vlm")
class Robo2VLMDataset(VLMBaseDataset):
    def __init__(
        self,
        data_paths: Union[list[str], str],
        config: DictConfig,
        tokenizer: AutoTokenizer,
    ) -> None:
        super().__init__(data_paths, config, tokenizer)
        self.system_prompt = (
            "You are a helpful robotic vision assistant specialized in "
            "answering questions about robotic manipulation tasks. "
            "Use <think></think> tags to show your reasoning process, "
            "then provide your final answer in <answer></answer> tags."
        )

    def get_image_list(self, dataitem: dict[str, Any]) -> list[Union[bytes, str, None]]:
        images: list[Any] = []
        if "images" in dataitem:
            v = dataitem.get("images")
            if isinstance(v, list):
                images = list(v)
            elif v is not None:
                images = [v]
            else:
                images = [None]
        elif "image" in dataitem:
            v = dataitem.get("image")
            if v is not None:
                images = [v]
            else:
                images = [None]
        else:
            return super().get_image_list(dataitem)

        normed: list[Union[bytes, str, None]] = []
        for v in images:
            if v is None:
                continue
            if isinstance(v, Image.Image):
                normed.append(v)
            elif isinstance(v, dict) and "bytes" in v:
                normed.append(v["bytes"])  # raw bytes
            else:
                normed.append(v)  # path/uri/string
        if not normed:
            normed = [None]
        return normed

    def build_prompt_text(self, data_item: dict[str, Any]) -> str:
        # Use 'question' and 'choices' if present; else fallback to base using configured prompt/choice keys
        question = data_item.get("question", None)
        choices = data_item.get("choices", None)
        if question is None:
            return super().build_prompt_text(data_item)
        # normalize choices
        if isinstance(choices, str):
            try:
                choices = ast.literal_eval(choices)
            except Exception:
                choices = [choices]
        if not isinstance(choices, list):
            choices = [choices] if choices is not None else []

        text = f"{question}\n"
        if choices:
            text += "Choices:\n"
            for i, c in enumerate(choices):
                text += f"{chr(65 + i)}. {c}\n"
        return text

    def postprocess_dataset_item(
        self, item: DatasetItem, raw: dict[str, Any]
    ) -> DatasetItem:
        answer_dict = {
            "choices": raw.get("choices", None),
            "correct_answer": raw.get("correct_answer", None),
        }
        item.answer = answer_dict

        return item


@VLMDatasetRegistry.register("robo2vlmsft")
class Robo2VLMSFTDataset(Robo2VLMDataset):
    # The reason for overriding the current function is that the SFT dataset requires masking the attention for the answer separately.
    def __init__(
        self,
        data_paths: Union[list[str], str],
        config: DictConfig,
        tokenizer: AutoTokenizer,
        eval_dataset: bool = False,
    ) -> None:
        super().__init__(data_paths, config, tokenizer)
        self.apply_chat_template = config.data.apply_chat_template
        self.eval_dataset = eval_dataset

    def _process_raw_record(self, raw: dict[str, Any], idx: int) -> DatasetItem:
        images = self.get_image_list(raw)
        (
            prompt_ids,
            plen,
            prompt_text,
            answer,
            attention_mask,
            label_mask,
            multi_modal_inputs,
        ) = self.encode_prompt(raw)

        item = SftDatasetItem(
            prompt=prompt_ids,
            length=plen,
            idx=idx,
            image_data=images,
            answer=answer,
            prompt_text=prompt_text,
            attention_mask=attention_mask,
            label_mask=label_mask,
            meta=None,
            multi_modal_inputs=multi_modal_inputs,
        )
        return item

    def vlm_vision_convert(
        self,
        prompt: dict[str, Any],
        images_inputs: list[Any],
    ):
        if self._processor is None:
            self._processor = AutoProcessor.from_pretrained(
                self.cfg.actor.model.model_path
            )

        prompt_text = self._processor.apply_chat_template(
            prompt, tokenize=False, add_generation_prompt=True
        )

        assert isinstance(prompt_text[0], str), (
            f"{type(prompt_text)} {prompt_text} is not str"
        )

        prompt_inputs = self._processor(
            text=prompt_text, images=images_inputs, return_tensors="pt", padding=True
        )

        return prompt_inputs

    def encode_prompt(
        self, raw: dict[str, Any]
    ) -> tuple[torch.Tensor, int, dict[str, Any], torch.Tensor, Optional[str]]:
        """
        Return (
            input_ids: torch.Tensor,      # [L] or [1, L] token ids
            int(input_ids.numel()),       # length of token ids
            prompt: dict[str, Any],       # raw prompt/messages
            answer: Optional[str],        # raw answer text (may be None)
            attention_mask: torch.Tensor, # [L] or [1, L] attention mask
            label_mask: torch.Tensor,     # [L] or [1, L] SFT label mask (prompt part masked)
            multi_modal_inputs: dict[str, Any], # processor outputs (e.g. pixel_values, image_grid_thw)
        ).

        If using chat template, encode with processor.
        Subclasses may override to support alternative prompting.

        Expected format:
        {
            "messages": [
                {"role": "user", "content": [
                    {"type": "image1", "image1": <PIL.Image1>}
                    {"type": "image2", "image2": <PIL.Image2>}
                    {"type": "text", "text": "User message"},
                ]},
                {"role": "assistant", "content": [
                    {"type": "text", "text": "Assistant response"}
                ]}
            ]
        }

        """
        question = raw.get(self.prompt_key, None)
        assert question is not None, "question is required"
        choices = raw.get(self.choice_key, None)
        assert choices is not None, "choices is required"
        correct_answer = raw.get(self.answer_key, None)
        assert correct_answer is not None, "correct_answer is required"
        images = self.get_image_list(raw)

        str_images = []
        for k in self.image_keys:
            v = raw.get(k)
            if isinstance(v, str):
                str_images.append(v)
            elif isinstance(v, list):
                str_images.extend([x for x in v if isinstance(x, str)])

        # convert the images to PIL Image objects
        images_inputs = []
        for image in images:
            image_obj = None
            if isinstance(image, Image.Image):
                image_obj = image.convert("RGB")
            if isinstance(image, (bytes, bytearray)):
                image_obj = Image.open(BytesIO(image)).convert("RGB")
            images_inputs.append(image_obj)

        # Handle both letter format (A, B, C, D) and index format (0, 1, 2, 3)
        if isinstance(correct_answer, str):
            # Map the answer letter (A, B, C, D) to index (0, 1, 2, 3)
            letter_to_idx = {"A": 0, "B": 1, "C": 2, "D": 3, "E": 4}
            correct_answer_idx = letter_to_idx.get(correct_answer, 0)
            correct_answer_letter = correct_answer
        else:
            # It's already an index
            correct_answer_idx = correct_answer
            correct_answer_letter = chr(
                65 + correct_answer_idx
            )  # Convert 0->A, 1->B, etc.

        # Parse choices if they're in string format
        if isinstance(choices, str):
            try:
                choices = ast.literal_eval(choices)
            except (ValueError, SyntaxError) as e:
                logging.warning(
                    "Failed to parse choices as list, fallback to single choice. "
                    "choices=%r, error=%s",
                    choices,
                    e,
                )
                choices = [str(choices)]

        # Create a formatted question with choices
        formatted_question = f"{question}\nChoices:\n"
        for i, choice in enumerate(choices):
            formatted_question += f"{chr(65 + i)}. {choice}\n"

        # Create the answer using the correct choice (matching finetune example format)
        answer = f"{correct_answer_letter}. {choices[correct_answer_idx]}"

        image_context = [{"type": "image", "image": img} for img in images_inputs]
        image_context.append({"type": "text", "text": formatted_question})

        prompt = [
            (
                {"role": "user", "content": image_context},
                {"role": "assistant", "content": [{"type": "text", "text": answer}]},
            )
        ]

        without_assistant_prompt = [({"role": "user", "content": image_context},)]

        # get the vlm model prompt inputs
        # prompt_inputs = input_ids, attention_mask, pixel_values, image_grid_thw
        prompt_inputs = self.vlm_vision_convert(prompt, images_inputs)
        # get the vlm model prompt inputs without assistant for label
        label_prompt_inputs = self.vlm_vision_convert(
            without_assistant_prompt, images_inputs
        )

        if self.eval_dataset:
            # for eval dataset, we need to use the label prompt inputs
            input_ids = label_prompt_inputs.pop("input_ids")
            attention_mask = label_prompt_inputs.pop("attention_mask")
            label_mask = attention_mask
        else:
            # for train dataset, we need to use the prompt prompt inputs
            input_ids = prompt_inputs.pop("input_ids")
            attention_mask = prompt_inputs.pop("attention_mask")
            label_mask = label_prompt_inputs.pop("attention_mask")

        if isinstance(input_ids, torch.Tensor):
            if input_ids.dim() == 2 and input_ids.size(0) == 1:
                input_ids = input_ids.squeeze(0)
            input_ids = input_ids.to(dtype=torch.long)
        else:
            input_ids = torch.tensor(input_ids, dtype=torch.long)

        multi_modal_inputs = {}
        if self.eval_dataset:
            for k, v in label_prompt_inputs.items():
                multi_modal_inputs[k] = v
        else:
            for k, v in prompt_inputs.items():
                multi_modal_inputs[k] = v

        return (
            input_ids,
            int(input_ids.numel()),
            prompt,
            answer,
            attention_mask,
            label_mask,
            multi_modal_inputs,
        )


def _resolve_video_path(path: str, data_root: Optional[str]) -> str:
    """Resolve a video path, rewriting it with data_root if the original path is missing."""
    if not isinstance(path, str):
        return str(path)
    if os.path.isfile(path):
        return path
    if not data_root:
        return path
    # Extract the data/ suffix from absolute paths and rewrite with data_root
    # to handle host-to-container path differences.
    idx = path.find("data/")
    if idx >= 0:
        resolved = os.path.join(data_root, path[idx:])
        if os.path.isfile(resolved):
            return resolved
    # Relative path: join it directly with data_root.
    if not os.path.isabs(path):
        resolved = os.path.join(data_root, path)
        if os.path.isfile(resolved):
            return resolved
    return path


@VLMDatasetRegistry.register("qwentrend_progress_sft")
class QwenTrendProgressSFTDataset(VLMBaseDataset):
    """SFT dataset for QwenTrend progress: full_video + video_clip input, aligned with the pilot.

    Each record: full_video (mp4), video_clip (mp4), question, answer.
    Qwen3-VL uses videos input, matching the dataset format.
    If JSONL paths are host absolute paths, set data.data_root to resolve them in the container.
    """

    def __init__(
        self,
        data_paths: Union[list[str], str],
        config: DictConfig,
        tokenizer: AutoTokenizer,
        eval_dataset: bool = False,
    ) -> None:
        super().__init__(data_paths, config, tokenizer)
        self.eval_dataset = eval_dataset
        self._data_root = config.data.get("data_root") or os.environ.get(
            "RLINF_DATA_ROOT"
        )

    @classmethod
    def _build_video_user_content(
        cls, processor: AutoProcessor, prompt_text: str
    ) -> str:
        video_tok = getattr(processor, "video_token", "<|video_pad|>")
        return f"{video_tok}\n{video_tok}\n\n{prompt_text}"

    @classmethod
    def process_inputs(
        cls,
        processor: AutoProcessor,
        system_prompt: Optional[str],
        use_chat_template: bool,
        prompt_texts: list[str] | list[list[str]],
        videos: list[Any] | list[list[Any]],
        answer_text: Optional[str] | list[Optional[str]] = None,
    ) -> tuple[str | list[str], dict[str, Any], dict[str, Any]]:
        """
        Build Qwen3-VL processor inputs for QwenTrend progress SFT.
        """

        def _render_prompt_text(
            prompt_text: str, answer_text_i: Optional[str]
        ) -> tuple[str, str]:
            user_content = cls._build_video_user_content(processor, prompt_text)

            try:
                if answer_text_i is None:
                    rendered_prompt = processor.apply_chat_template(
                        [{"role": "user", "content": user_content}],
                        tokenize=False,
                        add_generation_prompt=True,
                    )
                    rendered_label = rendered_prompt
                else:
                    rendered_prompt = processor.apply_chat_template(
                        [
                            {"role": "user", "content": user_content},
                            {"role": "assistant", "content": answer_text_i},
                        ],
                        tokenize=False,
                        add_generation_prompt=True,
                    )
                    rendered_label = processor.apply_chat_template(
                        [{"role": "user", "content": user_content}],
                        tokenize=False,
                        add_generation_prompt=True,
                    )
            except Exception:
                rendered_prompt = (
                    f"<|im_start|>user\n{user_content}<|im_end|>\n"
                    f"<|im_start|>assistant\n"
                )
                rendered_label = rendered_prompt

            return rendered_prompt, rendered_label

        is_batch_input = bool(prompt_texts) and isinstance(prompt_texts[0], list)
        if is_batch_input:
            prompt_texts_batch: list[list[str]] = prompt_texts
            videos_batch: list[list[Any]] = videos
            batch_size = len(prompt_texts_batch)

            if isinstance(answer_text, list):
                assert len(answer_text) == batch_size, (
                    f"answer_text list size {len(answer_text)} does not match batch size {batch_size}"
                )
                answer_text_batch = answer_text
            else:
                answer_text_batch = [answer_text for _ in range(batch_size)]

            rendered_prompts: list[str] = []
            rendered_labels: list[str] = []
            videos_kwargs = {"video_metadata": []}

            for prompt_texts_i, videos_i, answer_text_i in zip(
                prompt_texts_batch, videos_batch, answer_text_batch
            ):
                rendered_prompt_i, rendered_label_i = _render_prompt_text(
                    prompt_texts_i[0], answer_text_i
                )
                rendered_prompts.append(rendered_prompt_i)
                rendered_labels.append(rendered_label_i)
                videos_kwargs["video_metadata"].append(
                    [
                        {"total_num_frames": len(video), "fps": 24.0}
                        for video in videos_i
                    ]
                )

            full_inputs = processor(
                text=rendered_prompts,
                videos=videos_batch,
                return_tensors="pt",
                padding=True,
                videos_kwargs=videos_kwargs,
            )

            if all(answer is None for answer in answer_text_batch):
                label_inputs = {"attention_mask": full_inputs["attention_mask"]}
            else:
                label_inputs = processor(
                    text=rendered_labels,
                    videos=videos_batch,
                    return_tensors="pt",
                    padding=True,
                    videos_kwargs=videos_kwargs,
                )

            return rendered_prompts, full_inputs, label_inputs

        prompt_text = prompt_texts[0]
        rendered_prompt, rendered_label = _render_prompt_text(prompt_text, answer_text)
        videos_kwargs = {
            "video_metadata": [
                {"total_num_frames": len(video), "fps": 24.0} for video in videos
            ],
        }

        full_inputs = processor(
            text=[rendered_prompt],
            videos=videos,
            return_tensors="pt",
            padding=True,
            videos_kwargs=videos_kwargs,
        )

        if answer_text is None:
            # Inference: avoid an extra processor call for label_text.
            # Downstream expects `label_inputs` to at least provide `attention_mask`.
            label_inputs = {"attention_mask": full_inputs["attention_mask"]}
        else:
            label_inputs = processor(
                text=[rendered_label],
                videos=videos,
                return_tensors="pt",
                padding=True,
                videos_kwargs=videos_kwargs,
            )

        return rendered_prompt, full_inputs, label_inputs

    def encode_prompt(
        self,
        prompt_text: str,
        videos: list[str],
        answer_text: str,
    ) -> tuple[torch.Tensor, int, torch.Tensor, torch.Tensor, dict[str, Any]]:
        """
        Encode prompt into input_ids + masks for SFT training.

        Returns:
          - input_ids: token ids for (user + assistant)
          - length: number of tokens in input_ids
          - attention_mask: mask for input_ids
          - label_mask: mask for prompt part only (used to keep answer tokens)
          - multi_modal_inputs: processor outputs excluding text-only tensors
        """
        if self._processor is None:
            self._processor = AutoProcessor.from_pretrained(
                self.cfg.actor.model.model_path
            )

        _, full_inputs, label_inputs = self.process_inputs(
            processor=self._processor,
            system_prompt=self.system_prompt,
            use_chat_template=self.use_chat_template,
            prompt_texts=[prompt_text],
            videos=videos,
            answer_text=answer_text,
        )

        input_ids = full_inputs.pop("input_ids")
        attention_mask = full_inputs.pop("attention_mask")
        label_mask = label_inputs.pop("attention_mask")

        if isinstance(input_ids, torch.Tensor):
            if input_ids.dim() == 2 and input_ids.size(0) == 1:
                input_ids = input_ids.squeeze(0)
            input_ids = input_ids.to(dtype=torch.long)
        else:
            input_ids = torch.tensor(input_ids, dtype=torch.long)

        plen = int(input_ids.numel())
        multi_modal_inputs = dict(full_inputs)
        return input_ids, plen, attention_mask, label_mask, multi_modal_inputs

    @classmethod
    def _parse_raw_record(
        cls,
        raw: dict[str, Any],
        idx: int,
        data_root: Optional[str],
    ) -> tuple[str, str, list[Any], list[Any]]:
        full_video = raw.get("full_video")
        video_clip = raw.get("video_clip")
        question = str(raw.get("question", ""))
        answer_text = str(raw.get("answer", ""))

        if full_video and video_clip:
            full_video = _resolve_video_path(str(full_video), data_root)
            video_clip = _resolve_video_path(str(video_clip), data_root)
            return (
                question,
                answer_text,
                [full_video, video_clip],
                [full_video, video_clip],
            )

        pkl_path = raw.get("pkl_path")
        if not pkl_path:
            raise ValueError(f"Sample {idx} missing full_video/video_clip or pkl_path")

        resolved_pkl_path = _resolve_video_path(str(pkl_path), data_root)
        with open(resolved_pkl_path, "rb") as f:
            payload = pickle.load(f)
        main_frames = payload.get("main_frames")
        extra_view_frames = payload.get("extra_view_frames")
        if main_frames is None or extra_view_frames is None:
            raise ValueError(f"Sample {idx} pkl missing dual-view frame arrays")
        return (
            question,
            answer_text,
            [main_frames, extra_view_frames],
            [resolved_pkl_path, resolved_pkl_path],
        )

    def _process_raw_record(self, raw: dict[str, Any], idx: int) -> "SftDatasetItem":
        prompt_text, answer_text, videos, image_data = self._parse_raw_record(
            raw, idx, self._data_root
        )
        input_ids, plen, attention_mask, label_mask, multi_modal_inputs = (
            self.encode_prompt(
                prompt_text=prompt_text,
                videos=videos,
                answer_text=answer_text,
            )
        )
        if plen > self.max_prompt_length:
            input_ids = input_ids[: self.max_prompt_length]
            attention_mask = attention_mask[..., : self.max_prompt_length]
            label_mask = label_mask[..., : self.max_prompt_length]
            plen = self.max_prompt_length
        return SftDatasetItem(
            prompt=input_ids,
            length=plen,
            idx=idx,
            image_data=image_data,
            answer=answer_text,
            prompt_text=prompt_text,
            attention_mask=attention_mask,
            label_mask=label_mask,
            meta=None,
            multi_modal_inputs=multi_modal_inputs,
        )


@VLMDatasetRegistry.register("simple_qwentrend_sft")
class SimpleQwenTrendSFTDataset(QwenTrendProgressSFTDataset):
    """SFT dataset for a single-video, single-word QwenTrend format."""

    @classmethod
    def _build_video_user_content(
        cls, processor: AutoProcessor, prompt_text: str
    ) -> str:
        video_tok = getattr(processor, "video_token", "<|video_pad|>")
        return f"{video_tok}\n\n{prompt_text}"

    @classmethod
    def _parse_raw_record(
        cls,
        raw: dict[str, Any],
        idx: int,
        data_root: Optional[str],
    ) -> tuple[str, str, list[str], list[str]]:
        clip_path = raw.get("clip_path") or raw.get("video_clip")
        if not clip_path:
            raise ValueError(f"Sample {idx} missing clip_path or video_clip")

        clip_path = _resolve_video_path(str(clip_path), data_root)
        prompt_text = str(raw.get("prompt") or raw.get("question") or "").strip()
        if not prompt_text:
            raise ValueError(f"Sample {idx} missing prompt or question")

        supervision = raw.get("supervision")
        supervision_label = (
            supervision.get("label", "") if isinstance(supervision, dict) else ""
        )
        answer_text = (
            str(raw.get("answer") or raw.get("label") or supervision_label)
            .strip()
            .lower()
        )
        return prompt_text, answer_text, [clip_path], [clip_path]
