# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

from typing import Any

import dataclasses
from dataclasses import dataclass

from .configuration_nemotron_h import NemotronHConfig


@dataclass
class MoeConfig:
    moe_intermediate_size: int | None = None
    num_experts_per_tok: int | None = None
    block_type: str = "moe"


@dataclass
class AttentionConfig:
    block_type: str = "attention"


@dataclass
class MambaConfig:
    block_type: str = "mamba"


BLOCK_TYPE_TO_CONFIG_CLASS = {config_class.block_type: config_class
                              for config_class in (AttentionConfig, MoeConfig, MambaConfig)}


class NemotronHPuzzleConfig(NemotronHConfig):
    model_type = "nemotron_h_puzzle"

    def __init__(self, **kwargs):
        self.block_configs = []
        if "block_configs" in kwargs:
            self.block_configs, layers_block_type = build_block_configs(kwargs.pop("block_configs"))
            kwargs["layers_block_type"] = layers_block_type

        self.mtp_block_configs = []
        if "mtp_block_configs" in kwargs:
            self.mtp_block_configs, mtp_layers_block_type = build_block_configs(kwargs.pop("mtp_block_configs"))
            kwargs["num_nextn_predict_layers"] = len(self.mtp_block_configs) // 2
            if mtp_layers_block_type:
                kwargs["mtp_layers_block_type"] = mtp_layers_block_type
                assert mtp_layers_block_type == ["attention", "moe"], \
                    "The MTP module must either not exist or be ['attention', 'moe'] to be supported by vLLM."

        super().__init__(**kwargs)

        self._validate_block_field_consistency()
        self._delete_blockwise_members_from_global_config()

    def to_dict(self) -> dict[str, Any]:
        """
        Remove None values from block_configs and mtp_block_configs.
        """
        output = super().to_dict()
        output["block_configs"] = [{k: v for k, v in dataclasses.asdict(block_config).items() if v is not None}
                                   for block_config in self.block_configs]
        output["mtp_block_configs"] = [{k: v for k, v in dataclasses.asdict(block_config).items() if v is not None}
                                       for block_config in self.mtp_block_configs]
        return output

    def get_nemotron_h_config_for_layer(self, layer_idx: int) -> NemotronHConfig:
        """
        Builds a vanilla NemotronHConfig that matches the block configuration for a specific layer.
        MTP layers are considered to be in location self.num_hidden_layers + mtp_layer_idx.
        """
        config_dict = self.to_dict()
        
        if layer_idx < self.num_hidden_layers:
            block_configs = config_dict["block_configs"]
        else:
            assert layer_idx < self.num_hidden_layers + len(self.mtp_block_configs), \
                f"layer_idx={layer_idx} is out of bounds for " \
                f"{self.num_hidden_layers=} + {len(self.mtp_block_configs)=}"
            layer_idx -= self.num_hidden_layers
            block_configs = config_dict["mtp_block_configs"]
        
        block_config = block_configs[layer_idx]
        
        config_dict.update(block_config)
        del config_dict["block_configs"]
        del config_dict["mtp_block_configs"]
        
        nemotron_h_config = NemotronHConfig.from_dict(config_dict)
        nemotron_h_config._attn_implementation = self._attn_implementation

        return nemotron_h_config

    @property
    def mtp_n_routed_experts(self) -> int:
        n_routed_experts = self.n_routed_experts
        assert isinstance(n_routed_experts, int)
        return n_routed_experts

    def _validate_block_field_consistency(self):
        """Validate that no block field is set on only some of the blocks that define it.

        A block field is either global for every block (left unset / None) or explicitly set for
        every block that defines it - mixing the two (some blocks set it, others fall back to the
        global value) is not supported.
        """
        field_is_set: dict[str, set[bool]] = {}
        for block_config in self.block_configs + self.mtp_block_configs:
            for field in dataclasses.fields(block_config):
                if field.name == "block_type":
                    continue
                field_is_set.setdefault(field.name, set()).add(getattr(block_config, field.name) is not None)
        for name, statuses in field_is_set.items():
            if len(statuses) > 1:
                raise ValueError(
                    f"Block field '{name}' must be either set for all blocks that define it, or left "
                    f"unset (global) for all of them - mixing per-block and global values is not supported."
                )

    def _delete_blockwise_members_from_global_config(self):
        """Move blockwise fields off the global config.

        A field is blockwise iff it is set (non-None) in the blocks that define it; otherwise it
        stays on the global config.
        """
        self.blockwise_members = list(set(
            field.name
            for block_config in self.block_configs + self.mtp_block_configs
            for field in dataclasses.fields(block_config)
            if getattr(block_config, field.name) is not None
        ))
        for member in self.blockwise_members:
            if hasattr(self, member):
                delattr(self, member)


def build_block_configs(block_config_dicts: list[dict]) -> tuple[list[dataclass], list[str]]:
    block_configs = []

    # build block config dataclasses
    for block_config in block_config_dicts:
        if isinstance(block_config, dict):
            config_class = BLOCK_TYPE_TO_CONFIG_CLASS[block_config["block_type"]]
            block_config = config_class(**block_config)
        block_configs.append(block_config)

    # build the explicit per-layer block-type list
    layers_block_type = [block_config.block_type for block_config in block_configs]

    return block_configs, layers_block_type