import asyncio
from types import SimpleNamespace
from unittest.mock import Mock

import pytest

from artemis.exploration.budget import ExplorationBudget
from artemis.runtime import task_control


def test_budget_excludes_demonstration_and_waiting_and_resets_only_on_continue(monkeypatch):
    now = [0.0]
    monkeypatch.setattr("artemis.runtime.task_limits.time.monotonic", lambda: now[0])
    pause = Mock()
    monkeypatch.setattr("artemis.exploration.budget.request_pause", pause)
    repository = Mock()
    budget = ExplorationBudget("run", 30, repository)
    now[0] = 10
    assert budget.update({"status": "demonstrating"}) == 20
    now[0] = 500
    assert budget.update({"status": "task_paused"}) == 20
    now[0] = 1000
    assert budget.update({"status": "running"}) == 20
    now[0] = 1020
    assert budget.update({"status": "running"}) == 0
    pause.assert_called_once_with("run", reason="execution_budget")
    now[0] = 1030
    budget.update({"status": "task_paused"})
    budget.update({"status": "task_paused"})
    assert pause.call_count == 1
    now[0] = 2000
    assert budget.update({"status": "running", "budget_epoch": 1}) == 30
    assert budget.pause_requested is False
    assert budget.previous_elapsed == 40
    now[0] = 2010
    assert budget.update({"status": "running", "budget_epoch": 1}) == 20


@pytest.mark.asyncio
async def test_answer_received_before_pause_acknowledgement_is_not_lost(tmp_path, monkeypatch):
    monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)

    class Lock:
        device_id = "test-device"
        description = "test"
        max_concurrency = None
        concurrency_mode = "per_device"
        ingress = "test"
        lock_scope = None

        def __init__(self, *args, **kwargs):
            self.released = False

        def acquire(self, **kwargs):
            return None

        def release(self):
            self.released = True

    monkeypatch.setattr(task_control, "DeviceExecutionLock", Lock)
    control = task_control.TaskControl("run", Lock())
    task_control.request_pause("run", reason="question")
    task_control.add_instruction("run", "用户已回答", instruction_id="answer")
    assert await asyncio.wait_for(control.checkpoint(), 1)
    assert task_control.read_state("run")["status"] == "running"
    assert control.budget_epoch == 0


@pytest.mark.asyncio
async def test_budget_resume_increments_epoch_after_lease_is_reacquired(tmp_path, monkeypatch):
    monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
    old = SimpleNamespace(
        device_id="test",
        description="test",
        max_concurrency=None,
        concurrency_mode="per_device",
        ingress="test",
        lock_scope=None,
        release=Mock(),
    )
    lock = SimpleNamespace(acquire=Mock(), release=Mock())
    monkeypatch.setattr(task_control, "DeviceExecutionLock", Mock(return_value=lock))
    control = task_control.TaskControl("run", old)
    task_control.request_pause("run", reason="execution_budget")
    task_control.request_resume("run")
    assert await control.checkpoint()
    lock.acquire.assert_called_once()
    assert task_control.read_state("run")["budget_epoch"] == 1