import gc
import os
import torch
from typing import List, Optional, Set, Tuple, Type
from vllm.config import ParallelConfig
from vllm.distributed import init_distributed_environment, set_custom_all_reduce
from vllm.model_executor import set_random_seed
from vllm.worker.mlu_worker import MLUWorker, _check_if_gpu_supports_dtype
from vllm_mlu.worker.mlu_worker import MLUWorker_V2
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from ..distributed.parallel_state import ensure_model_parallel_initialized

import functools
from collections import defaultdict
from vllm.logger import init_logger
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding, ParallelLMHead
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear)
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
from vllm.distributed import get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size
from ..distributed.parallel_state import (get_moe_tensor_parallel_rank, get_moe_tensor_parallel_world_size,
                                          get_moe_expert_parallel_rank, get_moe_expert_parallel_world_size)


logger = init_logger(__name__)


def vllm__worker__mlu_worker__init_worker_distributed_environment(
    parallel_config: ParallelConfig,
    rank: int,
    distributed_init_method: Optional[str] = None,
    local_rank: int = -1,
) -> None:
    set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)

    init_distributed_environment(parallel_config.world_size, rank,
                                 distributed_init_method, local_rank,
                                 backend='cncl')

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add context_parallel_size, moe_tp_size, moe_ep_size
    '''
    ensure_model_parallel_initialized(parallel_config=parallel_config)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

def vllm__worker__mlu_worker__MLUWorker__init_device(self) -> None:
    if self.device_config.device.type == "mlu":
        # torch.distributed.all_reduce does not free the input tensor until
        # the synchronization point. This causes the memory usage to grow
        # as the number of all_reduce calls increases. This env var disables
        # this behavior.
        # Related issue:
        # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
        os.environ["TORCH_CNCL_AVOID_RECORD_STREAMS"] = "1"

        # This env var set by Ray causes exceptions with graph building.
        os.environ.pop("CNCL_ASYNC_ERROR_HANDLING", None)
        self.device = torch.device(f"mlu:{self.local_rank}")
        torch.mlu.set_device(self.device)

        _check_if_gpu_supports_dtype(self.model_config.dtype)
        gc.collect()
        torch.mlu.empty_cache()
        self.init_gpu_memory = torch.mlu.mem_get_info()[0]
    else:
        raise RuntimeError(
            f"Not support device type: {self.device_config.device}")
    # Initialize the distributed environment.
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: @brief: modify to vllm__worker__mlu_worker__init_worker_distributed_environment
    '''
    vllm__worker__mlu_worker__init_worker_distributed_environment(self.parallel_config, self.rank,
                                                                  self.distributed_init_method, self.local_rank)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    # Set random seed.
    set_random_seed(self.model_config.seed)


def default_act_range_value():
    return {
        "x": None,
        "split": None,
        "is_linear": False,
        "is_qkv": False,
        "q_proj_size": 0,
        "num_kv_head_replicas": 1,
        "is_merge": False,
        "input_id": [],
        "self_rank": 0,
        "rank": None,
        "tensor_rank": None,
        "tp_world_size": None,
        "moe_tp_rank": None,
        "moe_tp_world_size": None,
        "moe_ep_rank": None,
        "moe_ep_world_size": None,
        "weight": None,
    }

def vllm_mlu__worker__mlu_worker__MLUWorker_V2__setup_smooth_hook(self,
                                                       is_save_input_id: bool = False,
                                                       is_save_moe_info: bool = False):
    model = self.model_runner.model
    self.act_range = defaultdict(default_act_range_value)
    self.hooks = []
    linear_class_list = (ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear)
    other_class_list = (VocabParallelEmbedding, ParallelLMHead)
    class_list = linear_class_list + other_class_list
    row_class_list = (RowParallelLinear)

    for name, m in model.named_modules():
        if isinstance(m, FeedForward):
            m.use_bt_ffn = False
        if isinstance(m, SparseMoeMlp):
            m.is_use_fused_moe = False

        if isinstance(m, class_list):
            is_linear = True if isinstance(m, linear_class_list) else False
            split_type = "row" if isinstance(m, row_class_list) else "col"
            self.act_range[name]["split"] = split_type
            self.act_range[name]["is_linear"] = is_linear
            if isinstance(m, QKVParallelLinear):
                self.act_range[name]["is_qkv"] = True
                self.act_range[name]["q_proj_size"] = m.num_heads * m.head_size
                self.act_range[name]["num_kv_head_replicas"] = m.num_kv_head_replicas
            self.act_range[name]["is_merge"] = isinstance(m, MergedColumnParallelLinear)
            if is_save_moe_info:
                self.act_range[name]["rank"] = torch.distributed.get_rank()
                self.act_range[name]["tensor_rank"] = get_tensor_model_parallel_rank()
                self.act_range[name]["tp_world_size"] = get_tensor_model_parallel_world_size()
                self.act_range[name]["moe_tp_rank"] = get_moe_tensor_parallel_rank()
                self.act_range[name]["moe_tp_world_size"] = get_moe_tensor_parallel_world_size()
                self.act_range[name]["moe_ep_rank"] = get_moe_expert_parallel_rank()
                self.act_range[name]["moe_ep_world_size"] = get_moe_expert_parallel_world_size()
                if ".expert." in name:
                    self.act_range[name]["weight"] = m.weight
            logger.info(f"rank:{self.rank}, add hook to {name}, is_linear:{is_linear}, split_type:{split_type}")
            self.hooks.append(
                m.register_forward_hook(
                    functools.partial(self.stat_input_hook,
                                      name=name,
                                      act_range=self.act_range,
                                      is_linear=is_linear,
                                      is_save_input_id=is_save_input_id)))


def vllm_mlu__worker__mlu_worker__MLUWorker_V2__get_act_range(self):
    act_range = defaultdict(default_act_range_value)
    for layer_name, layer_range in self.act_range.items():
        for tensor_key, tensor_value in layer_range.items():
            if isinstance(tensor_value, torch.Tensor):
                act_range[layer_name][tensor_key] = tensor_value.to("cpu")
            elif tensor_key == "input_id" and isinstance(tensor_value, list):
                input_id_len = len(tensor_value)
                for i in range(input_id_len):
                    if isinstance(tensor_value[i], torch.Tensor):
                        act_range[layer_name][tensor_key].append(tensor_value[i].to("cpu"))
                    else:
                        act_range[layer_name][tensor_key].append(tensor_value[i])
            else:
                act_range[layer_name][tensor_key] = tensor_value

    return act_range


MluHijackObject.apply_hijack(MLUWorker,
                             MLUWorker.init_device,
                             vllm__worker__mlu_worker__MLUWorker__init_device)
MluHijackObject.apply_hijack(MLUWorker,
                             "setup_smooth_hook",
                             vllm_mlu__worker__mlu_worker__MLUWorker_V2__setup_smooth_hook)
MluHijackObject.apply_hijack(MLUWorker,
                             "get_act_range",
                             vllm_mlu__worker__mlu_worker__MLUWorker_V2__get_act_range)
