from __future__ import annotations

from typing import Any, Dict, Optional

from transformers import AutoConfig, PretrainedConfig


# Map a short name to (clamp_max, replacement_token).
# When the MLM tokenizer's vocab is larger than the AR model's, tokens above the AR
# vocab range must be replaced with something safe (typically the AR EOS) before
# being fed to the AR model. The replacement is also used for AR's BOS at position 0.
AR_TOKEN_PROCESSORS: Dict[str, Dict[str, int]] = {
    "identity": {"clamp_above": -1, "replacement": 0},
    "olmo": {"clamp_above": 50279, "replacement": 50279},  # OLMo eos
}


class MariaConfig(PretrainedConfig):
    """Config for a MARIA model.

    A MARIA model wraps a frozen autoregressive (AR) backbone, a frozen masked
    language model (MLM) backbone, and a trained linear fusion head that maps
    [ar_hidden ; mlm_hidden] -> vocab logits.

    The full sub-configs of both backbones are stored on this config so that
    ``from_pretrained`` can rebuild the backbones from a single bundled
    checkpoint without any external lookups.
    """

    model_type = "maria"
    is_composition = True

    def __init__(
        self,
        ar_config: Optional[Dict[str, Any]] = None,
        mlm_config: Optional[Dict[str, Any]] = None,
        ar_token_processor: str = "olmo",
        ar_attn_implementation: str = "sdpa",
        mlm_attn_implementation: str = "sdpa",
        torch_dtype: Optional[str] = "bfloat16",
        **kwargs,
    ):
        # AR and MLM configs are serialised as plain dicts to keep this config JSON-clean.
        # They may be None during fresh construction; build_hf_checkpoint.py fills them in.
        self.ar_config = ar_config
        self.mlm_config = mlm_config

        if ar_token_processor not in AR_TOKEN_PROCESSORS:
            raise ValueError(
                f"Unknown ar_token_processor={ar_token_processor!r}. "
                f"Available: {sorted(AR_TOKEN_PROCESSORS)}"
            )
        self.ar_token_processor = ar_token_processor
        self.ar_attn_implementation = ar_attn_implementation
        self.mlm_attn_implementation = mlm_attn_implementation

        super().__init__(torch_dtype=torch_dtype, **kwargs)

    # ---- convenience accessors -------------------------------------------------
    @property
    def ar_hidden_size(self) -> int:
        return self.ar_config["hidden_size"]

    @property
    def mlm_hidden_size(self) -> int:
        return self.mlm_config["hidden_size"]

    @property
    def vocab_size(self) -> int:
        return self.mlm_config["vocab_size"]

    @property
    def max_position_embeddings(self) -> int:
        return min(
            self.ar_config["max_position_embeddings"],
            self.mlm_config["max_position_embeddings"],
        )

    # ---- (de)serialisation -----------------------------------------------------
    @staticmethod
    def _config_from_dict(d: Dict[str, Any]) -> PretrainedConfig:
        d = dict(d)
        model_type = d.pop("model_type")
        # Drop transient fields that can confuse re-construction across HF versions.
        for k in ("_name_or_path", "transformers_version", "_attn_implementation"):
            d.pop(k, None)
        return AutoConfig.for_model(model_type, **d)

    def get_ar_pretrained_config(self) -> PretrainedConfig:
        return self._config_from_dict(self.ar_config)

    def get_mlm_pretrained_config(self) -> PretrainedConfig:
        return self._config_from_dict(self.mlm_config)

    @classmethod
    def from_backbones(
        cls,
        ar_config: PretrainedConfig,
        mlm_config: PretrainedConfig,
        ar_token_processor: str = "olmo",
        **kwargs,
    ) -> "MariaConfig":
        return cls(
            ar_config=ar_config.to_dict(),
            mlm_config=mlm_config.to_dict(),
            ar_token_processor=ar_token_processor,
            **kwargs,
        )