# 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.

from typing import Optional

import numpy as np
import torch

try:
    from qwen_vl_utils import extract_vision_info, process_vision_info
except ImportError:
    print("qwen_vl_utils is not available. Reason1 model will not work.")

from transformers.models.auto.processing_auto import AutoProcessor

from cosmos_rl.utils.logging import logger

_LOCK_TIMEOUT_SECONDS = 60


def build_tokenizer(
    tokenizer_type: str,
    cache_dir: Optional[str] = None,
):
    return Processor(tokenizer_type, cache_dir)


class Processor:  # TODO (maxzhaoshuol): This should be inherit AutoProcessor, instead of using self.processor?
    # This is a wrapper around the AutoProcessor class to add some helper functions
    def __init__(self, name="Qwen/Qwen2.5-VL-3B-Instruct", cache_dir=None):
        self.name = name

        if name in [
            "Qwen/Qwen2.5-VL-7B-Instruct",
            "Qwen/Qwen2.5-VL-3B-Instruct",
            "Qwen/Qwen2-VL-2B-Instruct",
        ]:
            self.processor = AutoProcessor.from_pretrained(name)
            logger.info("Successfully loaded processor from local cache")
        else:
            raise ValueError(
                f"Error loading processor {name}, please check if the tokenizer is available"
            )

        if hasattr(self.processor, "image_token"):
            self.image_token_id = self.processor.tokenizer.convert_tokens_to_ids(
                self.processor.image_token
            )
        else:
            self.image_token_id = None
        if hasattr(self.processor, "video_token"):
            self.video_token_id = self.processor.tokenizer.convert_tokens_to_ids(
                self.processor.video_token
            )
        else:
            self.video_token_id = None
        self.eos_id = (
            self.processor.tokenizer.eos_token_id
        )  # TODO (maxzhaoshuol): should be self.eos_token_id??
        self.pad_id = self.processor.tokenizer.pad_token_id

    def apply_chat_template(
        self,
        messages,
        add_generation_prompt=False,
        return_tensors="pt",
        tokenize=True,
        add_vision_id=False,
    ):
        assert tokenize, "tokenize must be True"
        if self.name.startswith("Qwen/Qwen2"):
            # Use manual workaround for add_vision_id bug
            text = self.processor.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=add_generation_prompt,
                add_vision_id=add_vision_id,
            )
            image_inputs, video_inputs, _ = process_vision_info(
                messages, return_video_kwargs=True
            )

            # add fps ourselves, as videos have been presampled according to the specified token length
            vision_infos = extract_vision_info(messages)
            fps_list = []
            for vision_info in vision_infos:
                if "video" in vision_info:
                    fps_list.append(vision_info["fps"])

            # process inputs
            inputs = self.processor(
                text=[text],
                images=image_inputs,
                videos=video_inputs,
                padding=False,
                return_tensors=return_tensors,
                fps=fps_list,
                # padding="max_length",
                # max_length=5000,  # Set the fixed length (adjust based on model requirements)
                # truncation=True,  # Ensures text longer than `max_length` is truncated
            )

            # save raw text
            inputs["text"] = text

            # Convert batch features into single features
            # By default, the processor returns a batch of features, but we use processor in dataloader, so we need to convert it to single features
            inputs["input_ids"] = inputs["input_ids"][
                0
            ]  # N_dialogue, N_token -> N_token
            inputs["attention_mask"] = inputs["attention_mask"][
                0
            ]  # N_dialogue, N_token -> N_token
            # inputs["image_grid_thw"]: N_img, 3
            # inputs["video_grid_thw"]: N_video, 3
        else:
            raise ValueError(
                f"apply_chat_template is not implemented for tokenizer_type {self.name}"
            )

        return inputs

    def add_assistant_tokens_mask(self, tokens):
        """
        Add a mask to the assistant tokens.
        This is used to mask out tokens that are not generated by the assistant (e.g.,  system prompts, user prompts, chat templates), such that in the loss computation, only the tokens generated by the assistant are used.
        If there are multiple turns in the conversation, the mask will mask all the assistant tokens in each turn.

        Args:
            tokens (Union[List[int], torch.Tensor]): The tokens to add the mask to.
        Returns:
            Union[List[bool], torch.Tensor]: The mask. True for tokens generated by the assistant (i.e. should apply loss on), False for tokens not generated by the assistant.
        """
        if isinstance(tokens, torch.Tensor) and tokens.ndim == 2:
            mask = torch.stack(
                [
                    self.add_assistant_tokens_mask(tokens[i])
                    for i in range(tokens.shape[0])
                ]
            )
            assert mask.shape == tokens.shape
            return mask
        np_tokens = (
            tokens.cpu().numpy()
            if isinstance(tokens, torch.Tensor)
            else np.array(tokens)
        )
        assert np_tokens.ndim == 1

        if self.name.startswith("Qwen/Qwen2"):
            # Constants defining bos, eos and fixed offsets.
            BOS_TOKEN = "<|im_start|>"
            EOS_TOKEN = "<|im_end|>"
            ROLE = "assistant"
            # Offsets: skip the bos + "assistant\n" (always 3 tokens) and include the eos (+1) for supervision
            START_OFFSET = 3
            END_OFFSET = 1

            # Retrieve token IDs for the markers and the role.
            bos_token_id = self.processor.tokenizer.convert_tokens_to_ids(BOS_TOKEN)
            eos_token_id = self.processor.tokenizer.convert_tokens_to_ids(EOS_TOKEN)
            role_id = self.processor.tokenizer.convert_tokens_to_ids(ROLE)

            # Locate all positions where the start and end markers appear.
            start_indices = np.where(np_tokens == bos_token_id)[0]
            end_indices = np.where(np_tokens == eos_token_id)[0]

            # Initialize the mask with False values.
            masks = np.zeros_like(np_tokens, dtype=bool)
            assert len(start_indices) == len(end_indices)
            # For each pair of bos/eos, check if the role is 'assistant'
            # and apply the mask accordingly.
            for start, end in zip(start_indices, end_indices):
                if np_tokens[start + 1] == role_id:
                    # Mask tokens from after the assistant header (start+3) to include the end marker (end+1)
                    masks[start + START_OFFSET : end + END_OFFSET] = True
        else:
            raise ValueError(
                f"add_assistant_tokens_mask is not implemented for tokenizer_type {self.name}"
            )

        assert masks.shape == np_tokens.shape
        if isinstance(tokens, torch.Tensor):
            return torch.from_numpy(masks)
        else:
            return masks.tolist()

    def encode(self, *args, **kwargs):
        return self.processor.encode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        return self.processor.decode(*args, **kwargs)


if __name__ == "__main__":
    """
    PYTHONPATH=. python3 projects/cosmos/reason1/tokenizer/processor.py
    """
    processor = Processor("Qwen/Qwen2.5-VL-3B-Instruct")
    print("done")
