"""CPU fixtures for the native scheduler; no model weights or device execution."""

from collections.abc import Callable, Iterator
from itertools import count
from types import SimpleNamespace
from typing import TYPE_CHECKING

import numpy as np
import pytest

if TYPE_CHECKING:
    from torch.distributed import ProcessGroup

    from deepinfer._C.engine import BlockManager, Scheduler, Sequence, SequenceMetadataPool


@pytest.fixture(scope="module")
def scheduler_process_group() -> Iterator["ProcessGroup"]:
    import torch.distributed as dist

    owns_group = not dist.is_initialized()
    if owns_group:
        dist.init_process_group("gloo", store=dist.HashStore(), rank=0, world_size=1)
    elif dist.get_world_size() != 1 or dist.get_backend() != "gloo":
        pytest.skip("Native scheduler unit tests require a single-rank CPU Gloo group")
    try:
        yield dist.group.WORLD
    finally:
        if owns_group:
            dist.destroy_process_group()


@pytest.fixture
def scheduler_factory(
    scheduler_process_group: "ProcessGroup", monkeypatch: pytest.MonkeyPatch
) -> Callable[..., tuple["Scheduler", "BlockManager", "BlockManager | None"]]:
    from deepinfer._C.engine import BlockManager, Scheduler

    def make(
        num_blocks: int = 33,
        *,
        num_states: int | None = None,
        stop_preempted_request: bool = True,
        nextn: int = 0,
        max_num_batched_tokens: int = 64,
        max_num_seqs: int = 4,
        max_model_len: int = 256,
    ) -> tuple["Scheduler", "BlockManager", "BlockManager | None"]:
        monkeypatch.setenv("DEEPINFER_USE_NEXTN_DECODE", str(nextn))
        blocks = BlockManager(num_blocks)
        states = BlockManager(num_states) if num_states is not None else None
        config = SimpleNamespace(
            max_model_len=max_model_len,
            runtime_max_model_len=max_model_len + 64,
            max_num_batched_tokens=max_num_batched_tokens,
            max_num_seqs=max_num_seqs,
            stop_preempted_request=stop_preempted_request,
            chunked_prefill=True,
            engine_type=0,
            prefill_attempt_interval=2,
            prefill_any_attempt_interval=100,
        )
        scheduler = Scheduler(
            block_size=16,
            block_manager=blocks,
            state_manager=states,
            eos_token_id=2,
            dp_group=scheduler_process_group,
            scheduler_config=config,
            truncate_len=0,
            lookahead_tokens=0,
            prefill_max_sum_seqlen_k=4096,
            flops_estimator=None,
            engram_lookup_table=[],
            cp_size=1,
            vocab_size=100,
            log_stats=False,
        )
        return scheduler, blocks, states

    return make


@pytest.fixture
def sequence_factory() -> Callable[..., "Sequence"]:
    from deepinfer._C.engine import Sequence
    from deepinfer.sampling_params import SamplingParams

    ids = count()

    def make(length: int = 128, *, request_id: str = "request", block_size: int = 16) -> "Sequence":
        tokens = np.ones(length, dtype=np.int32)
        return Sequence(
            seq_id=next(ids),
            prompt_token_ids=tokens,
            compressed_prompt_token_ids=tokens,
            multi_modal_data=None,
            block_size=block_size,
            request_id=request_id,
            user_id="test",
            sampling_params=SamplingParams(max_tokens=4, ignore_eos=True, eos_id=2, seq_seed=0),
            arrival_time=0.0,
            golden_token_ids=None,
        )

    return make


@pytest.fixture
def metadata_pool() -> "SequenceMetadataPool":
    from deepinfer._C.engine import SequenceMetadataPool

    return SequenceMetadataPool()


@pytest.fixture
def unscheduled_outputs() -> Callable:
    import torch

    from deepinfer._C.engine import RequestOutput, RequestOutputPool, build_request_outputs

    def build(outputs: object) -> list[object]:
        pool = RequestOutputPool()
        pointers = build_request_outputs(
            outputs,
            64,
            torch.empty(0, dtype=torch.int32),
            torch.empty(0, dtype=torch.float32),
            top_logprobs_=None,
            top_token_ids_=None,
            sum_smaller_probs_=None,
            accepted_token_ids_=None,
            accepted_token_logprobs_=None,
            accepted_top_logprobs_=None,
            accepted_top_token_ids_=None,
            accepted_sum_smaller_probs_=None,
            accepted_token_last_logprobs_=None,
            pushed_tokens_before_finish=np.empty(0, dtype=np.int32),
            prompt_logprobs=None,
            request_output_pool=pool,
        )
        result = [RequestOutput.clone(int(pointer)) for pointer in pointers]
        pool.release(pointers)
        return result

    return build