from __future__ import annotations

from collections.abc import Mapping

import torch
from transformers import AutoConfig, AutoModel, AutoModelForImageTextToText, PretrainedConfig, PreTrainedModel

try:
    from .action_head import MiniCPMV_VLA_ActionHead
    from .configuration_minicpm_vla import MiniCPMVLAConfig
except ImportError:  # pragma: no cover - fallback for non-package loaders
    from action_head import MiniCPMV_VLA_ActionHead
    from configuration_minicpm_vla import MiniCPMVLAConfig


def _torch_dtype(name: str) -> torch.dtype:
    try:
        return getattr(torch, name)
    except AttributeError as exc:
        raise ValueError(f"Unsupported torch dtype: {name}") from exc


def _build_vlm_config(config_dict: dict) -> PretrainedConfig:
    values = dict(config_dict)
    model_type = values.pop("model_type")
    return AutoConfig.for_model(model_type, **values)


class MiniCPMV_VLA(PreTrainedModel):
    """Complete VLA model loadable with ``AutoModel.from_pretrained``."""

    config_class = MiniCPMVLAConfig
    base_model_prefix = ""
    _tied_weights_keys = {
        "vlm.lm_head.weight": "vlm.model.language_model.embed_tokens.weight",
    }
    _no_split_modules = ["BasicTransformerBlock"]

    def __init__(self, config: MiniCPMVLAConfig):
        super().__init__(config)
        if config.vlm_config is None:
            raise ValueError("config.vlm_config is required")
        vlm_config = _build_vlm_config(config.vlm_config)
        self.vlm_dtype = _torch_dtype(config.vlm_dtype)
        self.action_head_dtype = _torch_dtype(config.action_head_dtype)
        self.vlm = AutoModelForImageTextToText.from_config(vlm_config).to(self.vlm_dtype)
        self.vlm.config.hidden_size = self.vlm.config.text_config.hidden_size
        self.action_head = MiniCPMV_VLA_ActionHead(
            action_dim=config.action_dim,
            state_dim=config.state_dim,
            action_horizon=config.action_horizon,
            num_inference_timesteps=config.num_inference_timesteps,
            max_num_embodiments=config.max_num_embodiments,
        ).to(self.action_head_dtype)
        self.post_init()

    def get_input_embeddings(self):
        return self.vlm.get_input_embeddings()

    def set_input_embeddings(self, value):
        return self.vlm.set_input_embeddings(value)

    def _vlm_forward(self, data: Mapping[str, object]):
        inputs = {key: value for key, value in data.items() if value is not None}
        if "language_attention_mask" in inputs and "attention_mask" not in inputs:
            inputs["attention_mask"] = inputs.pop("language_attention_mask")
        inputs.setdefault("output_hidden_states", True)
        inputs.setdefault("return_dict", True)
        inputs.setdefault("use_cache", False)
        inputs.setdefault("logits_to_keep", 1)
        with torch.autocast("cuda", dtype=self.vlm_dtype, enabled=torch.cuda.is_available()):
            try:
                return self.vlm(**inputs)
            except TypeError as exc:
                if "logits_to_keep" not in str(exc):
                    raise
                inputs.pop("logits_to_keep")
                return self.vlm(**inputs)

    @torch.no_grad()
    def predict_action(
        self,
        state: torch.Tensor,
        embodiment_id: torch.Tensor,
        **vlm_inputs,
    ) -> torch.Tensor:
        vl_embs = self._vlm_forward(vlm_inputs).hidden_states[-1].to(self.action_head_dtype)
        return self.action_head.predict_action(vl_embs, state, embodiment_id)


AutoConfig.register("minicpm_vla", MiniCPMVLAConfig)
AutoModel.register(MiniCPMVLAConfig, MiniCPMV_VLA)