"""Tests for scheduler concurrency."""

import threading
import time
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest
import torch

from astrai.extension import CudaBackend, TorchNativeBackend, get_backend
from astrai.inference import GenerationResult, InferenceScheduler
from astrai.inference.metrics import MetricsCollector
from astrai.inference.runtime.executor import DecodeSteadyState, Executor
from astrai.inference.runtime.stepper import Stepper
from astrai.inference.task import Task
from astrai.model.transformer import AutoRegressiveLM
from tests.helpers import FakeTokenizer, make_rollout_config


@pytest.fixture
def mock_model_and_tokenizer():
    """Create mock model and tokenizer."""
    mock_model = MagicMock()
    mock_model.config = MagicMock()
    mock_model.config.num_key_value_heads = 8
    mock_model.config.num_attention_heads = 8
    mock_model.config.hidden_size = 128
    mock_model.config.num_hidden_layers = 2
    mock_model.config.max_position_embeddings = 100
    mock_model.parameters.return_value = iter(
        [MagicMock(dtype=torch.float32, device=torch.device("cpu"))]
    )

    mock_tokenizer = MagicMock()
    mock_tokenizer.encode.return_value = [1, 2, 3, 4, 5]
    mock_tokenizer.decode.return_value = "token"
    mock_tokenizer.stop_ids = [0]
    mock_tokenizer.pad_id = None

    return mock_model, mock_tokenizer


def _make_mock_scheduler(mock_model_and_tokenizer):
    """Build a CPU scheduler over mocks, patching scheduler-internal imports."""
    mock_model, mock_tokenizer = mock_model_and_tokenizer
    with (
        patch("astrai.inference.scheduler.AutoModel"),
        patch("astrai.inference.scheduler.AutoTokenizer"),
    ):
        return InferenceScheduler(
            model=mock_model,
            tokenizer=mock_tokenizer,
            max_batch_size=4,
            device="cpu",
        )


def _run_threads(*workers, timeout=10.0):
    threads = [threading.Thread(target=worker) for worker in workers]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=timeout)


def test_scheduler_concurrent_add_task(mock_model_and_tokenizer):
    """Test concurrent add_task operations."""
    scheduler = _make_mock_scheduler(mock_model_and_tokenizer)

    results = {"task_ids": [], "errors": []}
    lock = threading.Lock()

    def add_task_worker(worker_id):
        try:
            for i in range(10):
                task_id = scheduler.add_task(f"prompt from worker {worker_id}-{i}")
                with lock:
                    results["task_ids"].append(task_id)
        except Exception as e:
            results["errors"].append(str(e))

    _run_threads(*(lambda wid=i: add_task_worker(wid) for i in range(5)))

    scheduler.stop()

    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
    assert len(results["task_ids"]) == 50


def test_generation_loop_activates_backend_in_worker_thread():
    scheduler = object.__new__(InferenceScheduler)
    scheduler._backend = TorchNativeBackend()
    scheduler._stop_event = threading.Event()
    scheduler._task_cache = MagicMock()

    observed = []
    task_mgr = MagicMock()
    task_mgr.tokenizer.stop_ids = [0]
    task_mgr.remove_finished_tasks.return_value = []
    task_mgr.get_active_tasks.return_value = []
    task_mgr.max_batch_size = 1
    task_mgr.pull_candidates.return_value = []
    task_mgr.has_work.return_value = False

    def observe_backend(*args, **kwargs):
        observed.append(type(get_backend()))
        scheduler._stop_event.set()

    task_mgr.wait_for_tasks.side_effect = observe_backend
    scheduler._task_mgr = task_mgr

    thread = threading.Thread(target=scheduler._run_generation_loop)
    thread.start()
    thread.join(timeout=5)

    assert not thread.is_alive()
    assert observed == [TorchNativeBackend]


def test_step_splits_decode_batch_by_request_backend():
    scheduler = object.__new__(InferenceScheduler)
    scheduler._cache = SimpleNamespace(page_size=1)
    scheduler._task_cache = MagicMock()
    scheduler._task_cache.task_extend.return_value = True
    scheduler._metrics = MetricsCollector()
    scheduler._executor = MagicMock()
    scheduler._stepper = Stepper(
        scheduler._cache, scheduler._task_cache, scheduler._executor, scheduler._metrics
    )

    observed = []

    def execute(tasks, **kwargs):
        observed.append((type(get_backend()), [task.task_id for task in tasks]))
        return [1] * len(tasks)

    scheduler._executor.execute_decode.side_effect = execute

    torch_task = Task("torch", [1], backend=TorchNativeBackend())
    cuda_task = Task("cuda", [1], backend=CudaBackend())
    for task in (torch_task, cuda_task):
        task.input_tokens = 1
        task.output_ids = [1]
        task.mark_prefill_done()
        scheduler._metrics.register(task.task_id)

    produced, aborted = scheduler._step([torch_task, cuda_task])

    assert aborted == []
    assert produced == [torch_task, cuda_task]
    assert observed == [
        (TorchNativeBackend, ["torch"]),
        (CudaBackend, ["cuda"]),
    ]


def test_step_batches_ragged_prefill_with_shared_cache_start():
    scheduler = object.__new__(InferenceScheduler)
    scheduler._cache = SimpleNamespace(page_size=64)
    scheduler._task_cache = MagicMock()
    scheduler._task_cache.task_cached.return_value = 0
    scheduler._metrics = MetricsCollector()
    scheduler._executor = MagicMock()
    scheduler._stepper = Stepper(
        scheduler._cache, scheduler._task_cache, scheduler._executor, scheduler._metrics
    )

    short = Task("short", [1, 2, 3])
    long = Task("long", [4, 5, 6, 7, 8])
    for task in (short, long):
        scheduler._metrics.register(task.task_id)

    scheduler._executor.execute_prefill.return_value = (
        [long, short],
        [11, 12],
    )

    produced, aborted = scheduler._step([short, long])

    assert aborted == []
    assert produced == [long, short]
    scheduler._executor.execute_prefill.assert_called_once_with(
        [short, long], start_pos=0, return_logprobs=False
    )
    assert long.output_ids == [11]
    assert short.output_ids == [12]


def test_execute_prefill_packs_ragged_prompts_and_selects_last_logits():
    executor = object.__new__(Executor)
    executor.device = torch.device("cpu")
    executor.task_cache = MagicMock()
    executor.task_cache.bind.return_value = MagicMock()
    executor._workspace = MagicMock()
    executor._workspace.max_batch_size = 16  # Add max_batch_size for validation
    all_logits = torch.arange(42, dtype=torch.float32).reshape(6, 7)

    def fake_model(ids, *, position_ids, kv_cache, fwd, logits_positions):
        return {"logits": all_logits[logits_positions]}

    executor.model = MagicMock(side_effect=fake_model)
    executor._sample_logits = MagicMock(
        return_value=([101, 102], torch.tensor([101, 102]))
    )

    task_b = Task("b", [20, 21, 22, 23, 24])
    task_a = Task("a", [10, 11, 12])

    tasks, output = executor.execute_prefill([task_b, task_a], start_pos=1)

    assert tasks == [task_a, task_b]
    assert output == [101, 102]
    model_args, model_kwargs = executor.model.call_args
    assert model_args[0].tolist() == [11, 12, 21, 22, 23, 24]
    assert model_kwargs["position_ids"].tolist() == [1, 2, 1, 2, 3, 4]
    assert model_kwargs["logits_positions"].tolist() == [1, 5]
    executor.task_cache.bind.assert_called_once_with(
        ["a", "b"], executor._workspace, start_pos=1
    )
    sample_args, sample_kwargs = executor._sample_logits.call_args
    torch.testing.assert_close(sample_args[0], all_logits[[1, 5]])
    assert sample_args[1] == [task_a, task_b]
    assert sample_args[2] is False
    assert sample_kwargs == {}


def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
    """Test concurrent add and remove task operations."""
    scheduler = _make_mock_scheduler(mock_model_and_tokenizer)

    results = {"added": [], "removed": [], "errors": []}
    add_ready = threading.Event()

    def add_worker():
        try:
            for i in range(20):
                task_id = scheduler.add_task(f"prompt {i}")
                results["added"].append(task_id)
                if len(results["added"]) >= 10:
                    add_ready.set()
        except Exception as e:
            results["errors"].append(f"Add: {str(e)}")

    def remove_worker():
        try:
            add_ready.wait(timeout=5.0)
            for task_id in results["added"][:10]:
                scheduler.remove_task(task_id)
                results["removed"].append(task_id)
        except Exception as e:
            results["errors"].append(f"Remove: {str(e)}")

    _run_threads(add_worker, remove_worker)
    scheduler.stop()

    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
    assert len(results["added"]) == 20


def test_scheduler_concurrent_get_stats(mock_model_and_tokenizer):
    """Test concurrent get_stats operations."""
    scheduler = _make_mock_scheduler(mock_model_and_tokenizer)

    results = {"stats": [], "errors": []}
    started = threading.Event()
    stats_done = threading.Event()

    def add_tasks():
        try:
            for i in range(20):
                scheduler.add_task(f"prompt {i}")
                started.set()
        except Exception as e:
            results["errors"].append(f"Add: {str(e)}")

    def get_stats():
        try:
            started.wait(timeout=5.0)
            for _ in range(50):
                stats = scheduler.get_stats()
                results["stats"].append(stats)
            stats_done.set()
        except Exception as e:
            results["errors"].append(f"Get stats: {str(e)}")

    _run_threads(add_tasks, get_stats)
    scheduler.stop()
    stats_done.wait(timeout=5.0)

    assert len(results["errors"]) == 0, f"Errors: {results['errors']}"
    assert len(results["stats"]) == 50

    for stats in results["stats"]:
        assert "total_tasks" in stats
        assert stats["total_tasks"] >= 0


def _make_real_scheduler(device):
    """Build a scheduler backed by a tiny real model for run_batch tests."""
    cfg = make_rollout_config(max_position_embeddings=64)
    model = AutoRegressiveLM(cfg).to(device=device, dtype=torch.bfloat16).eval()
    tokenizer = FakeTokenizer()
    scheduler = InferenceScheduler(
        model=model,
        tokenizer=tokenizer,
        max_batch_size=8,
        max_seq_len=64,
    )
    return scheduler, tokenizer, model


def test_cancel_waiting_task_storm_returns_to_baseline(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        task_ids = [
            scheduler.add_task(f"waiting-{index}", max_tokens=32) for index in range(32)
        ]

        assert all(scheduler.cancel_task(task_id) for task_id in task_ids)
        stats = scheduler.get_stats()
        assert stats["active_tasks"] == 0
        assert stats["waiting_tasks"] == 0
        assert stats["in_flight_tasks"] == 0
        assert stats["kv_cache_tasks"] == 0
        assert stats["cancelled_total"] == len(task_ids)
    finally:
        scheduler.stop()


def test_cancel_active_task_releases_metrics_and_kv(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        task_id = scheduler.add_task("active", max_tokens=32)
        task = scheduler._task_mgr.pull_candidates(1)[0]
        assert scheduler._task_cache.task_alloc(task.task_id, task.prompt_ids)
        assert scheduler._task_mgr.activate(task)

        before = scheduler.get_stats()
        assert before["active_tasks"] == 1
        assert before["in_flight_tasks"] == 1
        assert before["kv_cache_tasks"] == 1

        assert scheduler.cancel_task(task_id)
        scheduler.start()
        deadline = time.monotonic() + 5
        while time.monotonic() < deadline:
            after = scheduler.get_stats()
            if (
                after["active_tasks"] == 0
                and after["in_flight_tasks"] == 0
                and after["kv_cache_tasks"] == 0
            ):
                break
            time.sleep(0.01)

        assert after["active_tasks"] == 0
        assert after["waiting_tasks"] == 0
        assert after["in_flight_tasks"] == 0
        assert after["kv_cache_tasks"] == 0
        assert after["cancelled_total"] == 1
    finally:
        scheduler.stop()


def test_cancel_during_kv_allocation_releases_metrics_and_kv(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    allocation_started = threading.Event()
    continue_allocation = threading.Event()
    original_alloc = scheduler._task_cache.task_alloc

    def blocking_alloc(*args, **kwargs):
        allocation_started.set()
        assert continue_allocation.wait(timeout=5)
        return original_alloc(*args, **kwargs)

    try:
        with patch.object(
            scheduler._task_cache,
            "task_alloc",
            side_effect=blocking_alloc,
        ):
            scheduler.start()
            task_id = scheduler.add_task("allocation-race", max_tokens=32)
            assert allocation_started.wait(timeout=5)
            assert scheduler.cancel_task(task_id)
            continue_allocation.set()

            deadline = time.monotonic() + 5
            while time.monotonic() < deadline:
                stats = scheduler.get_stats()
                if (
                    stats["active_tasks"] == 0
                    and stats["waiting_tasks"] == 0
                    and stats["in_flight_tasks"] == 0
                    and stats["kv_cache_tasks"] == 0
                ):
                    break
                time.sleep(0.01)

        assert stats["active_tasks"] == 0
        assert stats["waiting_tasks"] == 0
        assert stats["in_flight_tasks"] == 0
        assert stats["kv_cache_tasks"] == 0
        assert stats["cancelled_total"] == 1
    finally:
        continue_allocation.set()
        scheduler.stop()


def test_run_batch_returns_token_sequences(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        prompts = [[10, 20, 30], [5, 6, 7, 8]]
        results = scheduler.run_batch(prompts, max_tokens=4, temperature=1.0)
        assert len(results) == 2
        for ids in results:
            assert isinstance(ids, list)
            assert len(ids) <= 4
            assert all(0 <= i < 200 for i in ids)
    finally:
        scheduler.stop()


def test_run_batch_tokens_match_full_sequence_forward(device):
    scheduler, _tok, model = _make_real_scheduler(device)
    prompt = [10, 20, 30, 40]
    try:
        expected = []
        sequence = list(prompt)
        for _ in range(2):
            input_ids = torch.tensor([sequence], dtype=torch.long, device=device)
            position_ids = torch.arange(len(sequence), device=device).unsqueeze(0)
            input_mask = torch.ones(
                1, len(sequence), len(sequence), dtype=torch.bool, device=device
            ).tril()
            with torch.inference_mode():
                logits = model(
                    input_ids,
                    input_mask=input_mask,
                    position_ids=position_ids,
                )["logits"][:, -1, :]
            token = logits.argmax(dim=-1).item()
            expected.append(token)
            sequence.append(token)

        result = scheduler.run_batch(
            prompt_ids_list=[prompt], max_tokens=2, temperature=0
        )
        assert result == [expected]
    finally:
        scheduler.stop()


def test_run_batch_return_logprobs_aligned(device):
    """return_logprobs=True gives (token_ids, logprobs) tuples with equal len."""
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        prompts = [[10, 20, 30, 40]]
        results = scheduler.run_batch(
            prompts, max_tokens=5, temperature=1.0, return_logprobs=True
        )
        assert len(results) == 1
        token_ids, logprobs = results[0]
        assert len(token_ids) == len(logprobs)
        assert all(lp <= 1e-5 for lp in logprobs)  # logprobs ≤ 0
    finally:
        scheduler.stop()


def test_ragged_prefill_matches_sequential_greedy_tokens_and_logprobs(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    prompts = [
        [10, 20, 30],
        [5, 6, 7, 8],
        [40, 41, 42, 43, 44],
    ]
    try:
        ragged = scheduler.run_batch(
            prompts, max_tokens=1, temperature=0, return_logprobs=True
        )
        sequential = [
            scheduler.run_batch(
                [prompt], max_tokens=1, temperature=0, return_logprobs=True
            )[0]
            for prompt in prompts
        ]

        assert [result[0] for result in ragged] == [result[0] for result in sequential]
        for ragged_result, sequential_result in zip(ragged, sequential):
            assert ragged_result[1] == pytest.approx(sequential_result[1], abs=1e-6)
    finally:
        scheduler.stop()


def test_run_batch_zero_max_tokens_returns_empty(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        assert scheduler.run_batch([[10, 20, 30]], max_tokens=0) == [[]]
    finally:
        scheduler.stop()


def test_run_batch_stop_id_terminates(device):
    """A token matching stop_ids terminates generation for that prompt."""
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        # Make every token a stop id: generation must end after exactly
        # one token (the stop token itself) instead of running to max_tokens.
        scheduler._task_mgr.tokenizer.stop_ids = list(range(200))
        prompts = [[10, 20, 30]]
        results = scheduler.run_batch(prompts, max_tokens=32, temperature=1.0)
        assert len(results[0]) == 1
    finally:
        scheduler.stop()


def test_run_batch_empty_prompts(device):
    """Empty prompt list yields empty result list."""
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        assert scheduler.run_batch([], max_tokens=4) == []
    finally:
        scheduler.stop()


def test_scheduler_weight_versions_are_monotonic_and_acknowledged(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        assert scheduler.policy_version == 0
        assert scheduler.update_weights(1) == 1
        assert scheduler.policy_version == 1
        assert scheduler.get_stats()["policy_version"] == 1
        assert scheduler.update_weights(1) == 1
        with pytest.raises(ValueError, match="cannot move backwards"):
            scheduler.update_weights(0)
        with pytest.raises(ValueError, match="non-negative integer"):
            scheduler.update_weights(True)
    finally:
        scheduler.stop()


def test_scheduler_applies_weight_mutation_and_version_atomically(device):
    scheduler, _tok, model = _make_real_scheduler(device)
    before = next(model.parameters()).detach().clone()

    def mutate():
        with torch.no_grad():
            next(model.parameters()).add_(1)
        return "updated"

    try:
        assert scheduler.apply_weight_update(1, mutate) == "updated"
        assert scheduler.policy_version == 1
        assert not torch.equal(next(model.parameters()), before)
        with pytest.raises(ValueError, match="must advance"):
            scheduler.apply_weight_update(1, mutate)

        def failed_mutation():
            raise RuntimeError("optimizer failed")

        with pytest.raises(RuntimeError, match="optimizer failed"):
            scheduler.apply_weight_update(2, failed_mutation)
        assert scheduler.policy_version == 1

        # None derives live+1 under the lock: no read-compute-write race
        # on the current version for advance-by-one callers.
        assert scheduler.apply_weight_update(None, mutate) == "updated"
        assert scheduler.policy_version == 2
    finally:
        scheduler.stop()


def test_scheduler_atomic_advance_survives_interleaved_publish(device):
    """A concurrent publish between reading the live version and applying
    the update must not fail ``require_advance`` (regression: callers
    computed live+1 outside the lock, a TOCTOU that raised spuriously)."""
    scheduler, _tok, _model = _make_real_scheduler(device)

    try:
        # Simulate the race directly: a version read that goes stale before
        # apply_weight_update acquires the lock. With None the scheduler
        # re-derives live+1 inside the critical section.
        stale_read = scheduler.policy_version + 1
        scheduler.update_weights(1)
        assert stale_read == 1  # now equals live -> explicit form would raise
        with pytest.raises(ValueError, match="must advance"):
            scheduler.apply_weight_update(stale_read, lambda: "ok")
        assert scheduler.apply_weight_update(None, lambda: "ok") == "ok"
        assert scheduler.policy_version == 2
    finally:
        scheduler.stop()


def test_scheduler_serializes_policy_snapshot_and_direct_update(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    snapshot_started = threading.Event()
    release_snapshot = threading.Event()
    update_finished = threading.Event()
    errors = []

    def inspect(version):
        assert version == 0
        snapshot_started.set()
        assert release_snapshot.wait(timeout=5)

    def take_snapshot():
        try:
            scheduler.with_policy_snapshot(inspect)
        except BaseException as exc:
            errors.append(exc)

    def update():
        try:
            scheduler.update_weights(1)
            update_finished.set()
        except BaseException as exc:
            errors.append(exc)

    snapshot_thread = threading.Thread(target=take_snapshot)
    update_thread = threading.Thread(target=update)
    try:
        snapshot_thread.start()
        assert snapshot_started.wait(timeout=5)
        update_thread.start()
        assert not update_finished.wait(timeout=0.1)
        release_snapshot.set()
        snapshot_thread.join(timeout=5)
        update_thread.join(timeout=5)
        assert not snapshot_thread.is_alive()
        assert not update_thread.is_alive()
        assert errors == []
        assert scheduler.policy_version == 1
    finally:
        release_snapshot.set()
        snapshot_thread.join(timeout=5)
        update_thread.join(timeout=5)
        scheduler.stop()


def test_scheduler_rejects_weight_update_with_queued_tasks(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    task_id = scheduler.add_task("queued")
    try:
        with pytest.raises(RuntimeError, match="while tasks are queued"):
            scheduler.update_weights(1)
        scheduler.remove_task(task_id)
        assert scheduler.update_weights(1) == 1
    finally:
        scheduler.stop()


def test_run_batch_too_long_prompt_skipped(device):
    """A prompt longer than max_seq_len yields an empty result slot."""
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        long = list(range(100))  # > max_seq_len=64
        results = scheduler.run_batch([long, [10, 20]], max_tokens=2)
        assert results[0] == []
        assert len(results[1]) <= 2
    finally:
        scheduler.stop()


def test_run_batch_details_distinguish_rejection_from_success(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        long_prompt = list(range(100))
        results = scheduler.run_batch(
            [long_prompt, [10, 20]],
            max_tokens=2,
            temperature=0,
            return_logprobs=True,
            return_details=True,
        )

        assert results[0] == GenerationResult(
            token_ids=[],
            logprobs=[],
            finish_reason="rejected",
            error_reason="prompt_too_long",
        )
        assert results[1].finish_reason in ("stop", "length")
        assert results[1].error_reason is None
        assert len(results[1].token_ids) == len(results[1].logprobs)
    finally:
        scheduler.stop()


def test_run_batch_details_report_non_positive_max_tokens(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        result = scheduler.run_batch([[10, 20]], max_tokens=0, return_details=True)[0]
        assert result.finish_reason == "rejected"
        assert result.error_reason == "max_tokens_non_positive"
    finally:
        scheduler.stop()


def test_run_batch_details_report_allocation_failure(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        with patch.object(scheduler._task_cache, "task_alloc", return_value=False):
            result = scheduler.run_batch([[10, 20]], max_tokens=2, return_details=True)[
                0
            ]
        assert result.finish_reason == "rejected"
        assert result.error_reason == "kv_cache_allocation_failed"
    finally:
        scheduler.stop()


def test_run_batch_details_report_extension_failure_and_cleanup(device):
    scheduler, _tok, _model = _make_real_scheduler(device)
    try:
        with patch.object(
            scheduler._stepper,
            "step",
            side_effect=lambda tasks, **_kwargs: ([], list(tasks)),
        ):
            result = scheduler.run_batch([[10, 20]], max_tokens=2, return_details=True)[
                0
            ]

        assert result.finish_reason == "rejected"
        assert result.error_reason == "kv_cache_extension_failed"
        assert scheduler._task_cache._states == {}
        assert scheduler._metrics._timings == {}
    finally:
        scheduler.stop()


def test_decode_does_not_reuse_previous_batch_state():
    executor = object.__new__(Executor)
    executor.device = torch.device("cpu")
    executor.task_cache = MagicMock()
    executor.task_cache.bind_was_steady = True
    executor.task_cache.bind.return_value = MagicMock()
    executor._graph_supported = False
    executor._graph_ctx = SimpleNamespace(enabled=False)

    workspace = MagicMock()
    workspace.max_batch_size = 16
    workspace.position_ids = torch.tensor([2], dtype=torch.long)
    workspace.fill_input_ids.return_value = torch.tensor([7], dtype=torch.long)
    workspace.decode_mask.return_value = torch.ones(1, 1, 9, dtype=torch.bool)
    executor._workspace = workspace
    executor.model = MagicMock(
        return_value={"logits": torch.zeros(1, 1, 10, dtype=torch.float32)}
    )

    old_info = object()
    new_info = object()
    executor._decode_cache = DecodeSteadyState(("old",), [2], old_info)
    executor._sample_logits = MagicMock(
        return_value=([3], torch.tensor([3], dtype=torch.long))
    )

    task = Task("new", list(range(8)), temperature=0)
    task.input_tokens = 8
    task.output_ids = [7]
    task.mark_prefill_done()

    with patch(
        "astrai.inference.runtime.executor._build_sampling_batch_info",
        return_value=new_info,
    ):
        assert executor.execute_decode([task]) == [3]

    assert workspace.position_ids.tolist() == [8]
    assert executor._decode_cache.task_sig == ("new",)
    executor._sample_logits.assert_called_once()
    args, kwargs = executor._sample_logits.call_args
    assert args[1:] == ([task], False)
    assert kwargs["info"] is new_info


def test_decode_fills_input_ids_from_device_on_matching_signature():
    """Steady-state decode copies cached device tokens, skipping the host."""
    executor = object.__new__(Executor)
    executor.device = torch.device("cpu")
    executor.task_cache = MagicMock()
    executor.task_cache.bind_was_steady = True
    executor.task_cache.bind.return_value = MagicMock()
    executor._graph_supported = False
    executor._graph_ctx = SimpleNamespace(enabled=False)

    workspace = MagicMock()
    workspace.max_batch_size = 16
    workspace.position_ids = torch.tensor([2], dtype=torch.long)
    workspace.fill_input_ids_from_device.return_value = torch.tensor(
        [9], dtype=torch.long
    )
    executor._workspace = workspace
    executor.model = MagicMock(
        return_value={"logits": torch.zeros(1, 1, 10, dtype=torch.float32)}
    )

    info = object()
    tokens = torch.tensor([3], dtype=torch.long)
    executor._decode_cache = DecodeSteadyState(("t1",), [2], info, last_tokens=tokens)
    executor._sample_logits = MagicMock(return_value=([3], tokens))

    task = Task("t1", list(range(8)), temperature=0)
    task.input_tokens = 8
    task.output_ids = [7]
    task.mark_prefill_done()

    with patch(
        "astrai.inference.runtime.executor._build_sampling_batch_info",
        return_value=info,
    ):
        assert executor.execute_decode([task]) == [3]

    workspace.fill_input_ids.assert_not_called()
    workspace.fill_input_ids_from_device.assert_called_once_with(tokens)
    assert workspace.position_ids.tolist() == [3]
    assert executor._decode_cache.task_sig == ("t1",)
    assert executor._decode_cache.last_tokens is tokens