from typing import (Any, Dict)

from vllm.model_executor.sampling_metadata import SequenceGroupToSample
from vllm.worker import model_runner_base
from vllm_mlu.mlu_hijack_utils import MluHijackObject

def vllm__worker__model_runner_base___init_sampling_metadata_from_tensor_dict(  # type: ignore
        tensor_dict: Dict[str, Any]) -> Dict[str, Any]:
    """
    Helper method to initialize SamplingMetadata based on broadcastable
    SamplingMetadata fields.
    """
    from vllm.model_executor import SamplingMetadata

    selected_token_indices = tensor_dict.pop("selected_token_indices", None)
    if selected_token_indices is not None:
        if 'seq_group_metadata' in tensor_dict.keys() and len(tensor_dict['seq_group_metadata']) > 0:
            '''
            ==========================
            Modify by Context Parallel
            ==========================
            @brief: construct sampling metadata.
            '''
            sequence_group_to_sample_list = []
            for seq_group_metadata in tensor_dict['seq_group_metadata']:
                seq_ids = list(seq_group_metadata.seq_data.keys())
                sampling_params = seq_group_metadata.sampling_params
                seq_data = seq_group_metadata.seq_data
                is_prompt = seq_group_metadata.is_prompt
                if is_prompt:
                    seq_len = query_len = list(seq_data.values())[0].get_prompt_len()
                else:
                    seq_len = None
                    query_len = 1
                prompt_logprob_indices = []
                sample_indices = seq_ids
                sequence_group_to_sample = SequenceGroupToSample(seq_ids,
                                                                 sampling_params,
                                                                 seq_data,
                                                                 seq_len,
                                                                 query_len,
                                                                 None, # Generator
                                                                 is_prompt,
                                                                 prompt_logprob_indices,
                                                                 sample_indices)
                sequence_group_to_sample_list.append(sequence_group_to_sample)
            tensor_dict["sampling_metadata"] = SamplingMetadata(
                seq_groups=sequence_group_to_sample_list,
                selected_token_indices=selected_token_indices,
                categorized_sample_indices=None,
                num_prompts=len(sequence_group_to_sample_list),
            )
            del tensor_dict['seq_group_metadata']
            '''
            =======================
            End of Context Parallel
            =======================
            '''
        else:
            # An empty SamplingMetadata to signal that the worker should skip
            # sampling.
            tensor_dict["sampling_metadata"] = SamplingMetadata(
                seq_groups=None,
                selected_token_indices=selected_token_indices,
                categorized_sample_indices=None,
                num_prompts=0,
            )
    return tensor_dict

MluHijackObject.apply_hijack(
    model_runner_base,
    model_runner_base._init_sampling_metadata_from_tensor_dict,
    vllm__worker__model_runner_base___init_sampling_metadata_from_tensor_dict
)
