# Copyright 2026 The 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.

"""Qwen processor construction and latent placeholder-token configuration."""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from lerobot.utils.import_utils import _transformers_available, require_package

if TYPE_CHECKING or _transformers_available:
    from transformers import AutoProcessor


@dataclass(frozen=True)
class LatentWorldProcessorSpec:
    """Portable inputs required to reconstruct the LaWAM VLM processor."""

    model_id: str
    placeholder_token: str


def build_latent_world_processor_spec(*, policy_cfg: Any, vlm_model_id: str) -> LatentWorldProcessorSpec:
    """Build a processor specification from the native policy configuration."""
    return LatentWorldProcessorSpec(
        model_id=str(vlm_model_id),
        placeholder_token=str(policy_cfg.latent_action_placeholder_token),
    )


def configure_latent_world_processor(
    processor: Any,
    *,
    placeholder_token: str,
) -> tuple[Any, Any, int]:
    """Install the latent placeholder token and return its tokenizer ID."""
    if processor.chat_template is None and getattr(
        getattr(processor, "tokenizer", None), "chat_template", None
    ):
        processor.chat_template = processor.tokenizer.chat_template

    tokenizer = processor.tokenizer
    tokenizer.add_special_tokens({"additional_special_tokens": [str(placeholder_token)]})
    placeholder_token_id = int(tokenizer.convert_tokens_to_ids(str(placeholder_token)))
    if placeholder_token_id < 0:
        raise ValueError(f"Invalid placeholder token id for `{placeholder_token}`.")

    return processor, tokenizer, placeholder_token_id


def load_latent_world_processor(
    spec: LatentWorldProcessorSpec,
) -> tuple[Any, Any, int]:
    """Load and configure the Qwen processor described by a portable spec."""
    require_package("transformers", extra="lawam")
    processor = AutoProcessor.from_pretrained(spec.model_id)
    return configure_latent_world_processor(
        processor,
        placeholder_token=spec.placeholder_token,
    )
