#!/usr/bin/env python

# Copyright 2025 HuggingFace Inc. team. All rights reserved.
#
# 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.

"""
Wall-X Utility Functions.

Contains data processing utilities, text formatting functions, and helper classes
for the Wall-X cross-embodiment robotic control model.
"""

import random
import re
from collections import OrderedDict
from typing import TYPE_CHECKING, Any

import torch

from lerobot.utils.import_utils import _transformers_available, _wallx_deps_available

if TYPE_CHECKING or _transformers_available:
    from transformers import BatchFeature
else:
    BatchFeature = None

if TYPE_CHECKING or _wallx_deps_available:
    from qwen_vl_utils.vision_process import smart_resize
else:
    smart_resize = None

from torchvision.transforms import InterpolationMode
from torchvision.transforms.v2 import functional as tv_functional

from lerobot.utils.constants import OBS_IMAGES

from .constant import (
    CAMERA_NAME_MAPPING,
    IMAGE_FACTOR,
    MAX_PIXELS,
    MIN_PIXELS,
    RESOLUTION,
)


def preprocesser_call(
    processor,
    images: list | Any | None = None,
    text: str | list[str] | None = None,
    videos: list | Any | None = None,
    device: torch.device | str | None = None,
    padding: bool | str = False,
    truncation: bool | None = None,
    max_length: int | None = None,
    return_tensors: str = "pt",
    target_spans: list[list[tuple[int, int]]] | None = None,
) -> BatchFeature:
    """Unified preprocessing function for Wall-X model handling text, image and video inputs.

    Processes inputs into format suitable for multimodal transformer models, including:
    - Text tokenization and special token handling
    - Image/video processing through image processor
    - Attention mask and label generation
    - Padding and truncation handling

    Args:
        processor: Multimodal processor containing tokenizer and image processor
        images: Input images (PIL, numpy arrays, or torch tensors)
        text: Text or list of texts to tokenize
        videos: Input videos (numpy arrays or torch tensors)
        device: Device on which image/video preprocessing should run
        padding: Whether to pad sequences to same length
        truncation: Whether to truncate sequences longer than max_length
        max_length: Maximum length for truncation/padding
        return_tensors: Format for returned tensors ('pt', 'np', etc.)
        target_spans: Character spans for text-supervised tokens in each prompt.
            ``None`` preserves the original action-only assistant-label path;
            an empty list for a sample explicitly disables text supervision.

    Returns:
        BatchFeature containing processed inputs with keys:
        - input_ids: Tokenized text
        - attention_mask: Attention mask for text
        - pixel_values: Processed image pixels
        - pixel_values_videos: Processed video frames
        - image_grid_thw: Image grid dimensions for LLM
        - video_grid_thw: Video grid dimensions for LLM
        - labels: Training labels with masking
    """
    # Process image inputs
    if images is not None and len(images) > 0:
        image_inputs = processor.image_processor(
            images=images,
            return_tensors=return_tensors,
            device=device,
        )
        image_grid_thw = image_inputs["image_grid_thw"]
    else:
        image_inputs = {}
        image_grid_thw = None

    # Process video inputs
    if videos is not None:
        videos_inputs = processor.image_processor(
            videos=videos,
            return_tensors=return_tensors,
            device=device,
        )
        video_grid_thw = videos_inputs["video_grid_thw"]
    else:
        videos_inputs = {}
        video_grid_thw = None

    # Ensure text input is in list format
    if text is None:
        raise ValueError("WALL-X preprocessing needs at least one prompt.")
    if not isinstance(text, list):
        text = [text]

    if target_spans is not None:
        if len(target_spans) != len(text):
            raise ValueError("WALL-X needs one target-span list for each prompt.")
        target_spans = [list(spans) for spans in target_spans]

    def replace_placeholder_tokens(
        prompt: str,
        placeholder: str,
        token_count: int,
        spans: list[tuple[int, int]],
    ) -> tuple[str, list[tuple[int, int]]]:
        position = prompt.find(placeholder)
        if position < 0:
            return prompt, spans
        replacement = "<|placeholder|>" * token_count
        placeholder_end = position + len(placeholder)
        for start, end in spans:
            if start < placeholder_end and end > position:
                raise ValueError("WALL-X text-supervision spans cannot contain image or video placeholders.")
        # The tokenizer receives the final prompt after each temporary
        # ``<|placeholder|>`` is restored to ``placeholder``. Shift spans
        # by that final length change, not by the temporary expansion.
        delta = len(placeholder) * (token_count - 1)
        spans = [
            (
                start + (delta if start >= placeholder_end else 0),
                end + (delta if end >= placeholder_end else 0),
            )
            for start, end in spans
        ]
        return prompt.replace(placeholder, replacement, 1), spans

    # Process image placeholder tokens in text.
    if image_grid_thw is not None:
        merge_length = processor.image_processor.merge_size**2
        index = 0
        for i in range(len(text)):
            while "<|image_pad|>" in text[i]:
                # Add bounds checking to avoid index overflow
                if index >= len(image_grid_thw):
                    print(
                        f"Warning: Number of image placeholders ({index + 1}) "
                        f"exceeds actual images ({len(image_grid_thw)}), "
                        f"skipping remaining placeholder processing"
                    )
                    break
                # Replace image placeholder with actual token count
                token_count = (image_grid_thw[index].prod() // merge_length).item()
                updated_text, updated_spans = replace_placeholder_tokens(
                    text[i],
                    "<|image_pad|>",
                    token_count,
                    target_spans[i] if target_spans is not None else [],
                )
                text[i] = updated_text
                if target_spans is not None:
                    target_spans[i] = updated_spans
                index += 1
            text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>")

    # Process video placeholder tokens in text
    if video_grid_thw is not None:
        merge_length = processor.image_processor.merge_size**2
        index = 0
        for i in range(len(text)):
            while "<|video_pad|>" in text[i]:
                # Replace video placeholder with actual token count
                token_count = (video_grid_thw[index].prod() // merge_length).item()
                updated_text, updated_spans = replace_placeholder_tokens(
                    text[i],
                    "<|video_pad|>",
                    token_count,
                    target_spans[i] if target_spans is not None else [],
                )
                text[i] = updated_text
                if target_spans is not None:
                    target_spans[i] = updated_spans
                index += 1
            text[i] = text[i].replace("<|placeholder|>", "<|video_pad|>")

    # Tokenize complete input text
    text_inputs = processor.tokenizer(
        text,
        return_tensors=return_tensors,
        padding=padding,
        truncation=truncation,
        max_length=max_length,
        return_offsets_mapping=target_spans is not None,
    )

    # Get pad token ID for label generation
    pad_token_id = processor.tokenizer.pad_token_id
    if pad_token_id is None:
        pad_token_id = processor.tokenizer.eos_token_id

    labels = torch.full_like(text_inputs.input_ids, -100)
    if target_spans is not None:
        offsets = text_inputs.pop("offset_mapping")
        for row, spans in enumerate(target_spans):
            for start, end in spans:
                overlap = (
                    (offsets[row, :, 1] > start)
                    & (offsets[row, :, 0] < end)
                    & text_inputs.attention_mask[row].bool()
                )
                labels[row, overlap] = text_inputs.input_ids[row, overlap]
    else:
        # Preserve the original action-only labeling path.
        assistant_marker = "<|im_start|>assistant\n"
        im_end_token_id = processor.tokenizer.convert_tokens_to_ids("<|im_end|>")
        assistant_tokens = processor.tokenizer("<|im_start|>assistant\n", add_special_tokens=False).input_ids

        for i in range(len(text)):
            assistant_regions = []
            parts = text[i].split(assistant_marker)
            num_left_pads = 0
            for token_id in text_inputs.input_ids[i]:
                if token_id == pad_token_id:
                    num_left_pads += 1
                else:
                    break
            current_pos = num_left_pads

            for j, part in enumerate(parts):
                part_tokens = processor.tokenizer(part, add_special_tokens=False).input_ids
                if j == 0:
                    current_pos += len(part_tokens)
                    continue
                for k in range(current_pos + 1, len(text_inputs.input_ids[i])):
                    if text_inputs.input_ids[i][k] == im_end_token_id:
                        assistant_regions.append((current_pos + len(assistant_tokens), k + 2))
                        break
                current_pos += len(part_tokens) + 3

            for start, end in assistant_regions:
                labels[i][start:end] = text_inputs.input_ids[i][start:end]

    # Mask special action tokens in labels
    action_token_id = processor.tokenizer.encode("<|action|>")[0]
    propri_token_id = processor.tokenizer.encode("<|propri|>")[0]
    labels[labels == action_token_id] = -100
    labels[labels == propri_token_id] = -100
    labels[labels == pad_token_id] = -100

    # Set labels to None if all are invalid to skip cross entropy loss
    if (labels != -100).any().item():
        text_inputs["labels"] = labels
    else:
        text_inputs["labels"] = None

    return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs})


def _wall_x_resize_dimensions(height: int, width: int) -> tuple[int, int, int, int]:
    """Return the intermediate and final Wall-X resize dimensions as ``(H, W, H, W)``."""
    if RESOLUTION == -1:
        intermediate_height, intermediate_width = height, width
    elif width > height:
        intermediate_width = RESOLUTION
        intermediate_height = int(RESOLUTION * height / width)
    else:
        intermediate_height = RESOLUTION
        intermediate_width = int(RESOLUTION * width / height)

    resized_height, resized_width = smart_resize(
        intermediate_height,
        intermediate_width,
        factor=IMAGE_FACTOR,
        min_pixels=MIN_PIXELS,
        max_pixels=MAX_PIXELS,
    )
    return intermediate_height, intermediate_width, resized_height, resized_width


def _resize_wall_x_image_batch(images: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int, int, int]]:
    """Quantize and resize a BCHW camera batch without leaving its current device."""
    if images.ndim != 4:
        raise ValueError(f"Wall-X images must be BCHW tensors, got shape {tuple(images.shape)}")

    original_height, original_width = images.shape[-2:]
    intermediate_height, intermediate_width, resized_height, resized_width = _wall_x_resize_dimensions(
        original_height, original_width
    )

    if images.is_floating_point():
        images = (images * 255).to(torch.uint8)
    elif images.dtype != torch.uint8:
        raise TypeError(f"Wall-X images must be floating point or uint8, got {images.dtype}")

    if images.shape[-2:] != (intermediate_height, intermediate_width):
        images = tv_functional.resize(
            images,
            [intermediate_height, intermediate_width],
            interpolation=InterpolationMode.BICUBIC,
            antialias=True,
        )
    if images.shape[-2:] != (resized_height, resized_width):
        images = tv_functional.resize(
            images,
            [resized_height, resized_width],
            interpolation=InterpolationMode.BICUBIC,
            antialias=True,
        )

    return images, (original_height, original_width, resized_height, resized_width)


def prepare_wall_x_image_inputs(
    batch: dict[str, Any], image_keys: list[str]
) -> tuple[list[list[torch.Tensor]], dict[str, tuple[int, int, int, int]]]:
    """Resize each camera batch and restore sample-major/camera-minor ordering."""
    resized_by_key: dict[str, torch.Tensor] = {}
    dimensions_by_key: dict[str, tuple[int, int, int, int]] = {}
    for key in image_keys:
        resized_by_key[key], dimensions_by_key[key] = _resize_wall_x_image_batch(batch[key])

    batch_size = batch[image_keys[0]].shape[0]
    image_inputs = [[resized_by_key[key][index] for key in image_keys] for index in range(batch_size)]
    return image_inputs, dimensions_by_key


def process_grounding_points(
    text: str,
    orig_height: int,
    orig_width: int,
    resized_height: int,
    resized_width: int,
    model_type: str,
) -> str:
    """Process grounding point coordinates in text based on image resizing.

    Adjusts coordinate values in <point> tags to match resized image dimensions
    for different model types (qwen2, qwen2_5).

    Args:
        text: Input text containing <point> tags with coordinates
        orig_height: Original image height
        orig_width: Original image width
        resized_height: Resized image height
        resized_width: Resized image width
        model_type: Model type for coordinate processing ('qwen2' or 'qwen2_5')

    Returns:
        Text with adjusted coordinate values
    """
    # Regex pattern to match <point> tags and their contents
    point_pattern = re.compile(r"<point>(.*?)</point>")

    def process_match(match):
        """Process a single point match and adjust coordinates."""
        coords_str = match.group(1)
        try:
            # Extract coordinates from string
            coords = list(map(int, re.findall(r"\d+", coords_str)))

            # Calculate resize scale factors
            scale_w = resized_width / orig_width
            scale_h = resized_height / orig_height

            if len(coords) == 2:
                x, y = coords
                if model_type == "qwen2_5":
                    # Qwen2.5 uses pixel coordinates
                    new_x = max(0, min(round(x * scale_w), resized_width - 1))
                    new_y = max(0, min(round(y * scale_h), resized_height - 1))
                elif model_type == "qwen2":
                    # Qwen2 normalizes to [0, 1000) range
                    new_x = max(0, min(999.999, (x / orig_width) * 1000))
                    new_y = max(0, min(999.999, (y / orig_height) * 1000))
                else:
                    raise ValueError(f"Unsupported model type: {model_type}")
                coords = [new_x, new_y]

            elif len(coords) == 4:
                x1, y1, x2, y2 = coords
                if model_type == "qwen2_5":
                    new_x1 = max(0, min(round(x1 * scale_w), resized_width - 1))
                    new_y1 = max(0, min(round(y1 * scale_h), resized_height - 1))
                    new_x2 = max(0, min(round(x2 * scale_w), resized_width - 1))
                    new_y2 = max(0, min(round(y2 * scale_h), resized_height - 1))
                elif model_type == "qwen2":
                    new_x1 = max(0, min(999.999, (x1 / orig_width) * 1000))
                    new_y1 = max(0, min(999.999, (y1 / orig_height) * 1000))
                    new_x2 = max(0, min(999.999, (x2 / orig_width) * 1000))
                    new_y2 = max(0, min(999.999, (y2 / orig_height) * 1000))
                else:
                    raise ValueError(f"Unsupported model type: {model_type}")
                coords = [new_x1, new_y1, new_x2, new_y2]

            # Return processed point tag
            return f"<point>[{', '.join(map(str, coords))}]</point>"

        except (ValueError, TypeError):
            # Return original content if processing fails
            return match.group(0)

    # Replace all matching point tags
    processed_text = point_pattern.sub(process_match, text)
    return processed_text


def get_frame_instruction(
    instruction_info: dict[str, Any],
    frame_idx: int | None = None,
    truncate_keys: list[str] | None = None,
) -> tuple[dict[str, Any], int | None]:
    """Extract frame-specific instruction from instruction dictionary.

    Args:
        instruction_info: Dictionary containing instruction components
        frame_idx: Current frame index
        truncate_keys: Keys that trigger truncation when found

    Returns:
        Tuple of (frame_instruction_dict, split_end_frame)
    """
    if truncate_keys is None:
        truncate_keys = [
            "subtask_generation",
            "distribute",
            "subtask_generation_zh",
            "distribute_zh",
        ]

    instruction_for_frame = {}
    split_end = None

    for key, value in instruction_info.items():
        if isinstance(value, dict):
            if frame_idx is None:
                raise ValueError("frame_idx is required to select frame-range instructions.")
            # Handle frame-range specific instructions
            for frame_range, frame_instruction in value.items():
                start_frame, end_frame = map(int, frame_range.split(" "))
                if start_frame <= frame_idx < end_frame or (start_frame == frame_idx):
                    instruction_for_frame[key] = frame_instruction
                    if truncate_keys is not None and split_end is None and key in truncate_keys:
                        split_end = end_frame + 1
                    break
        else:
            instruction_for_frame[key] = value

    return instruction_for_frame, split_end


def get_task_instruction(
    frame_instruction_info: dict[str, Any], priority_order: OrderedDict | None = None
) -> str:
    """Construct task instruction from available instruction fields using priority sampling.

    Args:
        frame_instruction_info: Dictionary containing instruction fields
        priority_order: OrderedDict specifying sampling probability for each field

    Returns:
        Combined instruction string with priority components
    """
    # Default priority settings
    default_priority_order = OrderedDict(
        {
            "subtask_generation": 0.25,
            "subtask_generation_zh": 0.25,
            "distribute": 0.25,
            "distribute_zh": 0.25,
        }
    )

    priority_order = OrderedDict(priority_order) if priority_order is not None else default_priority_order

    got_instruction = False
    task_instruction = ""

    # Sample instruction components based on priority probabilities
    for key, prob in priority_order.items():
        if key in frame_instruction_info and frame_instruction_info[key] != "":
            if got_instruction and random.random() >= prob:
                continue

            task_instruction += f"\n{frame_instruction_info[key]}"
            got_instruction = True
            break

    # Fall back to base instruction if no priority components found
    if not got_instruction:
        task_instruction = frame_instruction_info.get("instruction", "")

    return task_instruction


def get_wallx_normal_text(
    instruction_info: dict[str, Any],
    action_chunk_size: int,
    frame_idx: int,
    priority_order: OrderedDict | None = None,
    img_keys: list[str] | None = None,
    generate_subtask_ratio: float = 0.0,
) -> tuple[str, bool]:
    """Construct complete multimodal prompt text for Wall-X model.

    Formats input using special tokens including:
    - System message
    - User observations (with image placeholders)
    - Task instructions
    - Proprioception prompts
    - Assistant responses (with action tokens)

    Args:
        instruction_info: Dictionary containing instruction components
        action_chunk_size: Number of action tokens to generate
        frame_idx: Current frame index
        priority_order: Priority order for instruction sampling
        img_keys: List of image keys
        generate_subtask_ratio: Probability of generating subtask instead of actions

    Returns:
        Tuple of (formatted_prompt_text, is_subtask_generation)
    """
    # Special tokens for formatting
    role_start_symbol = "<|im_start|>"
    role_end_symbol = "<|im_end|>"
    vision_start_symbol = "<|vision_start|>"
    vision_end_symbol = "<|vision_end|>"
    image_pad_symbol = "<|image_pad|>"
    propri_symbol = "<|propri|>"
    action_symbol = "<|action|>"
    action_fast_symbol = "<|action_fast|>"

    # System prologue
    prologue = f"{role_start_symbol}system\nYou are a helpful assistant.{role_end_symbol}\n"

    # User request with observation
    user_request = f"{role_start_symbol}user\nObservation:"
    if img_keys:
        img_keys = img_key_mapping(img_keys)
        for key in img_keys:
            user_request += f" {key}: {vision_start_symbol}{image_pad_symbol}{vision_end_symbol}"
    user_request += "\nInstruction:"

    # Get frame-specific instruction
    frame_instruction_info, _ = get_frame_instruction(instruction_info, frame_idx=frame_idx)

    generate_subtask = False
    priority_keys = ["subtask_generation", "distribute"]

    # Decide whether to generate subtask or actions
    if (
        bool(set(frame_instruction_info.keys()) & set(priority_keys))
        and random.random() < generate_subtask_ratio
    ):
        # Generate subtask (equivalent to VQA task)
        instruction = frame_instruction_info.get("instruction", "")
        text_prompt = "\nPredict the next action in language.\n"
        user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n"

        # Find output instruction from priority keys
        for key in priority_keys:
            if key in frame_instruction_info:
                output_instruction = frame_instruction_info[key]
                break

        assistant_output = f"{role_start_symbol}assistant\n{output_instruction}\n{role_end_symbol}"
        generate_subtask = True
    else:
        # Generate actions
        instruction = get_task_instruction(frame_instruction_info, priority_order=priority_order)
        text_prompt = f"\nPredict the next action in robot action.\nProprioception: {propri_symbol}\n"
        user_message = f"{user_request} {instruction}{text_prompt}{role_end_symbol}\n"
        assistant_output = f"{role_start_symbol}assistant\n{action_fast_symbol}{role_end_symbol}\n{action_symbol * action_chunk_size}"

    complete_text = prologue + user_message + assistant_output
    return complete_text, generate_subtask


def img_key_mapping(img_keys: list[str]) -> list[str]:
    """Map image keys to camera names.

    Args:
        img_keys: List of image keys

    Returns:
        List of camera names
    """
    processed_img_keys = []
    for key in img_keys:
        key = key.replace(OBS_IMAGES + ".", "")
        if key in CAMERA_NAME_MAPPING:
            key = CAMERA_NAME_MAPPING[key]
        else:
            key = key.replace("_", " ") if "view" in key else key + " view"
        processed_img_keys.append(key)
    return processed_img_keys


def get_action_tokens(normalized_actions: torch.Tensor | list, action_tokenizer) -> list[list[str]]:
    """Convert normalized actions to action token strings.

    Args:
        normalized_actions: Normalized action arrays/tensors
        action_tokenizer: Tokenizer for converting actions to tokens

    Returns:
        List of action token string lists for each sample
    """
    if isinstance(normalized_actions, torch.Tensor):
        normalized_actions = normalized_actions.cpu().numpy()

    all_action_tokens = []
    for i in range(len(normalized_actions)):
        if isinstance(normalized_actions[i], torch.Tensor):
            normalized_actions[i] = normalized_actions[i].cpu().numpy()

        token_id = action_tokenizer(normalized_actions[i])
        action_tokens = [f"<|action_token_{j}|>" for j in token_id[0]]
        all_action_tokens.append(action_tokens)

    return all_action_tokens


def pad_action_token_strs(
    actions_token_lists: list[list[str]],
    pad_token: str = "<|endoftext|>",  # nosec B107
) -> list[str]:
    """Pad action token lists to same length and join as strings.

    Args:
        actions_token_lists: List of action token lists for each sample
        pad_token: Token used for padding

    Returns:
        List of padded action token strings
    """
    max_len = max(len(tokens) for tokens in actions_token_lists)
    padded_action_strs = []

    for tokens in actions_token_lists:
        padded_tokens = tokens + ["<|im_end|>\n"] + [pad_token] * (max_len - len(tokens))
        padded_action_strs.append("".join(padded_tokens))

    return padded_action_strs


def replace_action_token(
    text: list[str],
    norm_action: torch.Tensor | None,
    action_tokenizer,
    dof_masks: torch.Tensor | None = None,
) -> list[str]:
    """Replace action placeholders in text with actual action tokens.

    Args:
        text: List of text strings with action placeholders
        norm_action: Normalized action tensors
        action_tokenizer: Tokenizer for converting actions to tokens
        dof_masks: Masks for degrees of freedom

    Returns:
        List of text strings with action tokens replaced
    """
    if action_tokenizer is not None and norm_action is not None:
        if dof_masks is None:
            raise ValueError("dof_masks are required to tokenize actions.")
        # Extract actions based on chunk sizes and DOF masks
        norm_action = [action[:32, dof_masks[i, 0].bool()] for i, action in enumerate(norm_action)]

        # Convert to action tokens and pad
        actions_fast_tokens = get_action_tokens(norm_action, action_tokenizer)
        actions_fast_token_strs = pad_action_token_strs(actions_fast_tokens)

        # Replace action placeholders with actual tokens
        actions_fast_token_idx = 0
        for i in range(len(text)):
            if "<|action_fast|>" in text[i]:
                text[i] = text[i].replace(
                    "<|action_fast|><|im_end|>\n",
                    actions_fast_token_strs[actions_fast_token_idx],
                )
                actions_fast_token_idx += 1

        # Remove remaining action placeholders
        text = [t.replace("<|action|>", "") for t in text]
    else:
        # Remove action placeholders when no tokenizer available
        text = [t.replace("<|action_fast|><|im_end|>\n", "") for t in text]

    return text
