import os
import torch
from torch import nn
from typing import Optional
from vllm.logger import init_logger
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.model_executor.model_loader.loader import DefaultModelLoader
from vllm.config import VllmConfig, ModelConfig, ParallelConfig
from vllm_mlu._mlu_utils import *
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp


logger = init_logger(__name__)


def get_parallel_num(
    model_config: ModelConfig,
    parallel_config: ParallelConfig
):
    attention_parallel_num = os.environ.get(ATTN_PARALLEL_NUM)
    ffn_parallel_num = os.environ.get(FFN_PARALLEL_NUM)
    if attention_parallel_num and attention_parallel_num.isdecimal():
        attention_parallel_num = int(attention_parallel_num)
    else:
        attention_parallel_num = 0
    if ffn_parallel_num and ffn_parallel_num.isdecimal():
        ffn_parallel_num = int(ffn_parallel_num)
    else:
        ffn_parallel_num = 0

    if parallel_config.tensor_parallel_size == 1:
        raise ValueError("Can not use context_comm_cmpt_parallel when tp num is 1.")
    if (attention_parallel_num <= 0 and ffn_parallel_num <= 0):
        raise ValueError("attention_parallel_num and ffn_parallel_num must be positive integers.")

    hidden_size = model_config.get_hidden_size()
    ffn_parallel_num = max(ffn_parallel_num, 1)
    if hidden_size % ffn_parallel_num != 0:
        raise ValueError(f"Hidden_size: {hidden_size} must be divisible by ffn_parallel_num: {ffn_parallel_num}")

    return attention_parallel_num, ffn_parallel_num


def get_attr_by_path(obj, path):
    # Split the path by dots to get individual attributes
    attributes = path.split('.')
    # Iterate through the attributes to access nested members
    for attr in attributes:
        if not hasattr(obj, attr):
            return None
        obj = getattr(obj, attr)
    return obj


def set_custom_attributes(model, model_config, parallel_config):
    attn_row_parallel_layers = []
    attn_weights = []
    ffn_row_parallel_layers = []
    ffn_weights = []
    sparse_moe_mlp_layers = []
    for module in model.modules():
        if module.__class__.__name__ == "FeedForward":
            ffn_weight = []
            if hasattr(module, "up_proj_name"):
                up_proj_name = getattr(module, "up_proj_name")
                up_proj = getattr(module, up_proj_name)
                if hasattr(up_proj, "weight"):
                    ffn_weight.append(up_proj.weight)
            if hasattr(module, "down_proj_name"):
                down_proj_name = getattr(module, "down_proj_name")
                down_proj = getattr(module, down_proj_name)
                if hasattr(down_proj, "weight"):
                    ffn_weight.append(down_proj.weight)
            if ffn_weight is not None:
                ffn_weights.append(ffn_weight)
            ffn_row_parallel_layers.append(module)
        for child_module in module.children():
            if child_module.__class__.__name__ == "Attention":
                for sibling_module in module.children():
                    if sibling_module.__class__.__name__ == "QKVParallelLinear":
                        if hasattr(sibling_module, "weight"):
                            weight = getattr(sibling_module, "weight")
                            attn_weights.append([weight])
                    if sibling_module.__class__.__name__ == "RowParallelLinear":
                        attn_row_parallel_layers.append(sibling_module)
        if module.__class__.__name__ == "SparseMoeMlp" or issubclass(module.__class__, SparseMoeMlp):
            sparse_moe_mlp_layers.append(module)

    if VLLM_PRELOAD_SIZE > 0:
        if (len(attn_row_parallel_layers) \
            == len(attn_weights) \
            == len(ffn_row_parallel_layers) \
            == len(ffn_weights)) and \
            len(attn_row_parallel_layers) != 0:

            for i in range(len(attn_row_parallel_layers)):
                attn_row_parallel_layers[i].preloaded_weights = ffn_weights[i]
                attn_row_parallel_layers[i].preload_size = VLLM_PRELOAD_SIZE
                if i < len(attn_row_parallel_layers) - 1:
                    ffn_row_parallel_layers[i].preloaded_weights = attn_weights[i+1]
                    ffn_row_parallel_layers[i].preload_size = VLLM_PRELOAD_SIZE
        else:
            logger.warning("%s does not support preload weight!", model.__class__.__name__)

    # context compute communication parallel
    if check_context_comm_cmpt_parallel():
        attention_parallel_num, ffn_parallel_num = get_parallel_num(model_config, parallel_config)
        for o_proj in attn_row_parallel_layers:
            setattr(o_proj.quant_method, 'parallel_num', attention_parallel_num)
        
        if len(sparse_moe_mlp_layers) != 0:
            for sparse_moe_mlp in sparse_moe_mlp_layers:
                setattr(sparse_moe_mlp, 'parallel_num', ffn_parallel_num)
        else:
            for ffn in ffn_row_parallel_layers:
                setattr(ffn, 'parallel_num', ffn_parallel_num)


vllm__model_executor__model_loader__loader__DefaultModelLoader__load_model__org = DefaultModelLoader.load_model


def vllm__model_executor__model_loader__loader__DefaultModelLoader__load_model(
        self, vllm_config: VllmConfig) -> nn.Module:
    model = vllm__model_executor__model_loader__loader__DefaultModelLoader__load_model__org(
                self, vllm_config=vllm_config)
    '''
    =============================
    Modify by custom vllm_mlu
    =============================
    @brief: According to the layer name in models, set custom optimize attributes.
    '''
    set_custom_attributes(model, vllm_config.model_config, vllm_config.parallel_config)
    '''
    =========================
    End of custom MLU Hijack
    =========================
    '''
    return model


MluHijackObject.apply_hijack(DefaultModelLoader,
                            DefaultModelLoader.load_model,
                            vllm__model_executor__model_loader__loader__DefaultModelLoader__load_model)