#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import torch

import vllm.distributed.elastic_ep.elastic_execute as elastic_execute_module
from vllm.distributed.elastic_ep.elastic_execute import ElasticEPScalingExecutor


class _FakeWorker:
    def __init__(self, model_runner):
        self.model_runner = model_runner
        self.compile_calls = 0

    def compile_or_warm_up_model(self):
        self.compile_calls += 1


class _FakeV2Row:
    def __init__(self, values):
        self.gpu = torch.tensor(values, dtype=torch.int32)
        self.clear_staged_writes_calls = 0

    def clear_staged_writes(self):
        self.clear_staged_writes_calls += 1


class _FakeNumBlocks:
    def __init__(self, values):
        self.cpu = torch.tensor(values, dtype=torch.int32)
        self.copy_to_uva_calls = 0

    def copy_to_uva(self):
        self.copy_to_uva_calls += 1


class _FakeV2BlockTables:
    def __init__(self):
        self.block_tables = [
            _FakeV2Row([[1, 2], [3, 4]]),
            _FakeV2Row([[5, 6], [7, 8]]),
        ]
        self.input_block_tables = [
            torch.tensor([[9, 10], [11, 12]], dtype=torch.int32),
            torch.tensor([[13, 14], [15, 16]], dtype=torch.int32),
        ]
        self.num_blocks = _FakeNumBlocks([[2, 1], [1, 2]])


class _FakeV2Runner:
    def __init__(self):
        self.block_tables = _FakeV2BlockTables()
        self.max_num_tokens = 32
        self.dummy_run_calls = []

    def _dummy_run(self, num_tokens, *, is_profile, skip_eplb):
        self.dummy_run_calls.append((num_tokens, is_profile, skip_eplb))


class _FakeCpuGpuBuffer:
    def __init__(self, gpu_values, cpu_values):
        self.gpu = torch.tensor(gpu_values, dtype=torch.int32)
        self.cpu = torch.tensor(cpu_values, dtype=torch.int32)


class _FakeV1Row:
    def __init__(self, gpu_values, cpu_values):
        self.block_table = _FakeCpuGpuBuffer(gpu_values, cpu_values)


class _FakeMultiBlockTable:
    def __init__(self):
        self.block_tables = [
            _FakeV1Row([[1, 2], [3, 4]], [[5, 6], [7, 8]]),
            _FakeV1Row([[9, 10], [11, 12]], [[13, 14], [15, 16]]),
        ]
        self.clear_calls = 0

    def clear(self):
        self.clear_calls += 1
        for row in self.block_tables:
            row.block_table.gpu.zero_()
            row.block_table.cpu.zero_()


class _FakeInputBatch:
    def __init__(self):
        self.block_table = _FakeMultiBlockTable()


class _FakeV1Runner:
    def __init__(self):
        self.input_batch = _FakeInputBatch()
        self.max_num_tokens = 16
        self.dummy_run_calls = []

    def _dummy_run(self, num_tokens, *, is_profile, skip_eplb):
        self.dummy_run_calls.append((num_tokens, is_profile, skip_eplb))


def _patch_rewarm_dependencies(monkeypatch):
    monkeypatch.setattr(elastic_execute_module, "lock_workspace", lambda: None)
    monkeypatch.setattr(elastic_execute_module, "unlock_workspace", lambda: None)
    monkeypatch.setattr(
        ElasticEPScalingExecutor, "_release_cuda_graphs", lambda self: None
    )


def test_rewarm_workspace_uses_v2_block_tables(monkeypatch):
    _patch_rewarm_dependencies(monkeypatch)
    worker = _FakeWorker(_FakeV2Runner())
    executor = ElasticEPScalingExecutor(worker)
    block_tables = worker.model_runner.block_tables
    original_rows = [row.gpu.clone() for row in block_tables.block_tables]
    original_input_rows = [row.clone() for row in block_tables.input_block_tables]
    original_num_blocks = block_tables.num_blocks.cpu.clone()

    executor.rewarm_workspace()

    assert worker.model_runner.dummy_run_calls == [(32, True, True)]
    assert worker.compile_calls == 1
    assert [row.gpu.tolist() for row in block_tables.block_tables] == [
        row.tolist() for row in original_rows
    ]
    assert [row.tolist() for row in block_tables.input_block_tables] == [
        row.tolist() for row in original_input_rows
    ]
    assert torch.equal(block_tables.num_blocks.cpu, original_num_blocks)
    assert block_tables.num_blocks.copy_to_uva_calls == 2
    assert [row.clear_staged_writes_calls for row in block_tables.block_tables] == [2, 2]


def test_rewarm_workspace_preserves_v1_block_tables(monkeypatch):
    _patch_rewarm_dependencies(monkeypatch)
    worker = _FakeWorker(_FakeV1Runner())
    executor = ElasticEPScalingExecutor(worker)
    multi_block_table = worker.model_runner.input_batch.block_table
    original_gpu = [
        row.block_table.gpu.clone() for row in multi_block_table.block_tables
    ]
    original_cpu = [
        row.block_table.cpu.clone() for row in multi_block_table.block_tables
    ]

    executor.rewarm_workspace()

    assert worker.model_runner.dummy_run_calls == [(16, True, True)]
    assert worker.compile_calls == 1
    assert multi_block_table.clear_calls == 1
    assert [row.block_table.gpu.tolist() for row in multi_block_table.block_tables] == [
        row.tolist() for row in original_gpu
    ]
    assert [row.block_table.cpu.tolist() for row in multi_block_table.block_tables] == [
        row.tolist() for row in original_cpu
    ]