from typing import Any, Optional, Union
import numpy as np
import torch
from vllm import envs
from vllm.distributed.kv_transfer.kv_transfer_state import get_kv_transfer_group, has_kv_transfer_group
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
from vllm.forward_context import set_forward_context
from vllm.sequence import IntermediateTensors
from vllm.two_batch_overlap.v1.two_batch_overlap_v1 import tbo_model_executable_v1
from vllm.utils import async_tensor_h2d
from vllm.v1.attention.backends.mla.common import MLACommonMetadataBuilder
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
from vllm.v1.outputs import ModelRunnerOutput
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
from vllm.v1.worker.block_table import BlockTable


class TBOModelInputSplit():
    def __init__(self):
        self.req_ids_left = []
        self.req_ids_right = []
        self.req_num_left = 0
        self.req_num_right = 0
        self.scheduler_output_left = None
        self.scheduler_output_right = None
        self.split_in_req = False

input_split = TBOModelInputSplit()

def split_scheduler_output(runner, scheduler_output: SchedulerOutput):
    """Split a step's scheduled tokens evenly into left/right halves.
    If a request crosses the split boundary, mark split_in_req=True and
    assign left/right token counts accordingly.
    """
    split_tokens = scheduler_output.total_num_scheduled_tokens // 2
    split_counter = 0
    num_scheduled_tokens_left: dict[int, int] = {}
    num_scheduled_tokens_right: dict[int, int] = {}
    input_split.req_ids_left.clear()
    input_split.req_ids_right.clear()
    total_num_scheduled_tokens_left = split_tokens
    total_num_scheduled_tokens_right = scheduler_output.total_num_scheduled_tokens - split_tokens

    req_splited = False
    input_split.split_in_req = False

    for key, value in scheduler_output.num_scheduled_tokens.items():
        split_counter += value
        if split_counter == split_tokens:
            req_splited = True
            num_scheduled_tokens_left[key] = value
            input_split.req_ids_left.append(key)
        elif split_counter > split_tokens:
            if req_splited:
                # boundary already hit earlier; entire req goes to right
                num_scheduled_tokens_right[key] = value
                input_split.req_ids_right.append(key)
            else:
                # The boundary falls inside this request -> split within req
                req_splited = True
                input_split.split_in_req = True
                right_tokens = split_counter - split_tokens
                left_tokens = value - right_tokens
                # right part
                num_scheduled_tokens_right[key] = right_tokens
                input_split.req_ids_right.append(key)
                # left part
                num_scheduled_tokens_left[key] = left_tokens
                input_split.req_ids_left.append(key)
        else:
            # before boundary, entire req goes to left
            num_scheduled_tokens_left[key] = value
            input_split.req_ids_left.append(key)

    input_split.req_num_left = len(input_split.req_ids_left)
    input_split.req_num_right = len(input_split.req_ids_right)

    input_split.scheduler_output_left = SchedulerOutput(
        scheduled_new_reqs=None,
        scheduled_cached_reqs=None,
        num_scheduled_tokens=num_scheduled_tokens_left,
        total_num_scheduled_tokens=total_num_scheduled_tokens_left,
        scheduled_spec_decode_tokens=scheduler_output.scheduled_spec_decode_tokens,
        scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs,  # unsupported yet
        num_common_prefix_blocks=scheduler_output.num_common_prefix_blocks,
        finished_req_ids=scheduler_output.finished_req_ids,
        free_encoder_input_ids=scheduler_output.free_encoder_input_ids,
        structured_output_request_ids=scheduler_output.structured_output_request_ids,
        grammar_bitmask=scheduler_output.grammar_bitmask,
    )
    input_split.scheduler_output_right = SchedulerOutput(
        scheduled_new_reqs=None,
        scheduled_cached_reqs=None,
        num_scheduled_tokens=num_scheduled_tokens_right,
        total_num_scheduled_tokens=total_num_scheduled_tokens_right,
        scheduled_spec_decode_tokens=scheduler_output.scheduled_spec_decode_tokens,
        scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs,  # unsupported yet
        num_common_prefix_blocks=scheduler_output.num_common_prefix_blocks,
        finished_req_ids=scheduler_output.finished_req_ids,
        free_encoder_input_ids=scheduler_output.free_encoder_input_ids,
        structured_output_request_ids=scheduler_output.structured_output_request_ids,
        grammar_bitmask=scheduler_output.grammar_bitmask,
    )


def prepare_tbo_atten_metadata(
    runner,
    scheduler_output: "SchedulerOutput",
    req_ids,
    req_offset: int,
) -> dict[str, Any]:  # (attn_metadata)
    """Prepare attention metadata for one half (left/right).
    Key fixes for correctness when a request is split:
      - Align seq_len_offset / query_start_offset with block_table slicing.
      - For the right half, if a request was split, make seq_lens[0]
        = (history + left-prefix + right-half tokens).
      - Pass cloned slices to CommonAttentionMetadata to avoid aliasing.
    """
    total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
    assert total_num_scheduled_tokens > 0
    num_reqs = len(req_ids)
    assert num_reqs > 0

    # Tokens per req in THIS half
    tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
    num_scheduled_tokens = np.array(tokens, dtype=np.int32)
    max_num_scheduled_tokens = max(tokens)

    # Request indices (relative to the WHOLE step), used by kernels
    req_indices = np.repeat(runner.arange_np[:num_reqs],
                            num_scheduled_tokens) + req_offset

    # Cumulative tokens within this half
    cu_num_tokens, arange = runner._get_cumsum_and_arange(num_scheduled_tokens)

    # --- query_start_loc (within this half) ---
    runner.query_start_loc_np[0] = 0
    runner.query_start_loc_np[1:num_reqs + 1] = cu_num_tokens

    # --- seq_lens (absolute context length per-req row) ---
    # Default (no split across req boundary)
    # Maps rows [req_offset ... req_offset+num_reqs-1]
    default_seq_lens = (
        runner.input_batch.num_computed_tokens_cpu[req_offset : req_offset + num_reqs]
        + num_scheduled_tokens
    )

    # Offsets for copying into the *global* GPU buffers
    # Left-half writes at the natural position; right-half depends on split.
    if req_offset == 0:
        # LEFT
        seq_len_offset = 0
        query_start_offset = 0
        seq_lens_cpu_local = torch.as_tensor(default_seq_lens, device=runner.seq_lens_cpu.device)
    else:
        # RIGHT
        if input_split.split_in_req:
            # The block_table for RIGHT starts from (req_offset-1).
            # Align both offsets to that, and re-build the seq_lens for row-0.
            seq_len_offset = req_offset - 1
            query_start_offset = req_offset

            # row-0 is the split request (global row index = req_offset-1):
            base_hist = runner.input_batch.num_computed_tokens_cpu[req_offset - 1].item()
            left_prefix = input_split.scheduler_output_left.num_scheduled_tokens[req_ids[0]]
            right_tokens0 = scheduler_output.num_scheduled_tokens[req_ids[0]]
            first_row = base_hist + left_prefix + right_tokens0

            if num_reqs > 1:
                # rows 1.. map to global rows [req_offset .. req_offset+num_reqs-2]
                tail_base = runner.input_batch.num_computed_tokens_cpu[req_offset : req_offset + num_reqs - 1]
                tail_tokens = num_scheduled_tokens[1:]
                tail = tail_base + tail_tokens
                seq_lens_cpu_local = torch.empty(num_reqs, dtype=runner.seq_lens_cpu.dtype, device=runner.seq_lens_cpu.device)
                seq_lens_cpu_local[0] = first_row
                seq_lens_cpu_local[1:] = torch.as_tensor(tail, device=runner.seq_lens_cpu.device)
            else:
                seq_lens_cpu_local = torch.tensor([first_row], dtype=runner.seq_lens_cpu.dtype, device=runner.seq_lens_cpu.device)
        else:
            # RIGHT without split-in-req: natural positions
            seq_len_offset = req_offset
            query_start_offset = req_offset + 1
            seq_lens_cpu_local = torch.as_tensor(default_seq_lens, device=runner.seq_lens_cpu.device)

    # Copy query_start_loc into global GPU buffer window
    runner.query_start_loc[query_start_offset: query_start_offset + num_reqs + 1].copy_(
        runner.query_start_loc_cpu[:num_reqs + 1], non_blocking=True
    )
    # Pad tail (FlashAttn requires non-decreasing)
    if req_offset > 0:
        runner.query_start_loc[query_start_offset + num_reqs + 1:].fill_(
            runner.query_start_loc_cpu[num_reqs].item()
        )

    # Copy seq_lens into the aligned window; zero out the remainder on RIGHT
    runner.seq_lens[seq_len_offset : seq_len_offset + num_reqs].copy_(
        seq_lens_cpu_local, non_blocking=True
    )
    if req_offset > 0:
        runner.seq_lens[seq_len_offset + num_reqs:].fill_(0)

    # Build common metadata (pass CLONES to avoid aliasing between threads)
    # query_start_loc = runner.query_start_loc[query_start_offset: query_start_offset + num_reqs + 1].clone()
    # seq_lens = runner.seq_lens[seq_len_offset : seq_len_offset + num_reqs].clone()
    query_start_loc = runner.query_start_loc[query_start_offset: query_start_offset + num_reqs + 1]
    seq_lens = runner.seq_lens[seq_len_offset : seq_len_offset + num_reqs]
    common_attn_metadata = CommonAttentionMetadata(
        query_start_loc=query_start_loc,
        seq_lens=seq_lens,
        num_reqs=num_reqs,
        num_actual_tokens=total_num_scheduled_tokens,
        max_query_len=max_num_scheduled_tokens,
    )

    # Prepare attention metadata for each KV cache group
    attn_metadata: dict[str, Any] = {}
    for kv_cache_group_id, kv_cache_group_spec in enumerate(runner.kv_cache_config.kv_cache_groups):
        common_prefix_len = 0
        metadata_builder = runner.attn_metadata_builders[kv_cache_group_id]
        if runner.cascade_attn_enabled:
            common_prefix_len = runner._compute_cascade_attn_prefix_len(
                num_scheduled_tokens,
                scheduler_output.num_common_prefix_blocks[kv_cache_group_id],
                kv_cache_group_spec.kv_cache_spec,
                metadata_builder,
            )

        # Slice block_table / slot_mapping for RIGHT half
        if req_offset > 0:
            origin_block_table = metadata_builder.block_table.block_table
            if input_split.split_in_req:
                metadata_builder.block_table.block_table = origin_block_table[req_offset - 1:, :]
            else:
                metadata_builder.block_table.block_table = origin_block_table[req_offset:, :]

            origin_slot_mapping = metadata_builder.block_table.slot_mapping
            origin_slot_mapping_cpu = metadata_builder.block_table.slot_mapping_cpu
            left_tokens = input_split.scheduler_output_left.total_num_scheduled_tokens
            metadata_builder.block_table.slot_mapping = origin_slot_mapping[left_tokens:]
            metadata_builder.block_table.slot_mapping_cpu = origin_slot_mapping_cpu[left_tokens:]

        # MLA-specific counters (safe to ignore for Qwen/FA paths)
        if isinstance(metadata_builder, MLACommonMetadataBuilder):
            _num_decodes_record = metadata_builder._num_decodes
            _num_prefills_record = metadata_builder._num_prefills
            _num_decode_tokens_record = metadata_builder._num_decode_tokens
            _num_prefill_tokens_record = metadata_builder._num_prefill_tokens
            metadata_builder._num_decodes = 0
            metadata_builder._num_prefills = num_reqs
            metadata_builder._num_decode_tokens = 0
            metadata_builder._num_prefill_tokens = total_num_scheduled_tokens

        attn_metadata_i = metadata_builder.build(
            common_prefix_len=common_prefix_len,
            common_attn_metadata=common_attn_metadata,
        )

        # Restore tables
        if req_offset > 0:
            metadata_builder.block_table.block_table = origin_block_table
            metadata_builder.block_table.slot_mapping = origin_slot_mapping
            metadata_builder.block_table.slot_mapping_cpu = origin_slot_mapping_cpu

        if isinstance(metadata_builder, MLACommonMetadataBuilder):
            metadata_builder._num_decodes = _num_decodes_record
            metadata_builder._num_prefills = _num_prefills_record
            metadata_builder._num_decode_tokens = _num_decode_tokens_record
            metadata_builder._num_prefill_tokens = _num_prefill_tokens_record

        for layer_name in kv_cache_group_spec.layer_names:
            attn_metadata[layer_name] = attn_metadata_i

    return attn_metadata


def pad_num_input_tokens(self, scheduler_output):
    num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
    if (self.use_cuda_graph and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
        # CUDA graphs (piecewise). Add padding to batch size.
        num_input_tokens = self.vllm_config.pad_for_cudagraph(num_scheduled_tokens)
    else:
        # Eager mode: pad to TP multiple for SP+collective fusion
        tp_size = self.vllm_config.parallel_config.tensor_parallel_size
        if self.vllm_config.compilation_config.pass_config.enable_sequence_parallelism and tp_size > 1:
            from vllm.utils import round_up
            num_input_tokens = round_up(num_scheduled_tokens, tp_size)
        else:
            num_input_tokens = num_scheduled_tokens

    # DP padding
    num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
    num_input_tokens += num_pad
    return num_input_tokens, num_tokens_across_dp


def tbo_split_and_execute_model(
    runner,
    attn_metadata,
    num_input_tokens,
    num_tokens_across_dp,
    input_ids,
    positions,
    inputs_embeds,
    scheduler_output: "SchedulerOutput",
    intermediate_tensors: Optional[IntermediateTensors],
    skip_cuda_graphs: bool,
) -> Union[ModelRunnerOutput, IntermediateTensors]:
    # If below TBO threshold, run the normal single-batch path (supports decode/prefill as-is).
    # Two-batch overlap path
    split_scheduler_output(runner, scheduler_output)
    num_input_tokens_left = input_split.scheduler_output_left.total_num_scheduled_tokens
    num_input_tokens_right = input_split.scheduler_output_right.total_num_scheduled_tokens

    attn_metadata_left = prepare_tbo_atten_metadata(
        runner, input_split.scheduler_output_left, input_split.req_ids_left, 0
    )
    attn_metadata_right = prepare_tbo_atten_metadata(
        runner, input_split.scheduler_output_right, input_split.req_ids_right, input_split.req_num_left
    )

    # === Added: split inputs_embeds & intermediate_tensors per half; setup KV connector ===
    # real token nums
    num_tokens_left = int(input_split.scheduler_output_left.total_num_scheduled_tokens)
    num_tokens_right = int(input_split.scheduler_output_right.total_num_scheduled_tokens)
    
    # split intermediate tensors
    def _split_intermediate_tensors(it, l, r):
        if it is None: return None, None
        left_tensor_map, right_tensor_map = {}, {}
        for name, tensor in it.tensors.items():
            vl, vr = torch.split(tensor[:l + r], [l, r], dim=0)
            left_tensor_map[name], right_tensor_map[name] = vl, vr
        return IntermediateTensors(left_tensor_map), IntermediateTensors(right_tensor_map)

    intermediate_tensors_left, intermediate_tensors_right = _split_intermediate_tensors(
        intermediate_tensors, num_tokens_left, num_tokens_right
    )

    with set_forward_context(attn_metadata,
                                runner.vllm_config,
                                num_tokens=num_input_tokens,
                                num_tokens_across_dp=num_tokens_across_dp,
                                skip_cuda_graphs=True):
            runner.maybe_setup_kv_connector(scheduler_output)
            
            model_output = tbo_model_executable_v1(
                runner,
                attn_metadata_left, 
                attn_metadata_right, 
                num_input_tokens_left,
                num_input_tokens_right,
                num_tokens_across_dp,
                input_ids,
                positions,
                (intermediate_tensors_left, intermediate_tensors_right),
                inputs_embeds) 
            
            runner.maybe_wait_for_kv_save()
            finished_sending, finished_recving = (
                runner.get_finished_kv_transfers(scheduler_output))

    return model_output, finished_sending, finished_recving