from __future__ import annotations

from unittest import mock

import pytest

from task_plane.api import TaskPlaneService
from task_plane.errors import ConflictError, LockedError
from task_plane.policy import ALL_ACTION_CAPABILITIES, PolicyEngine
from task_plane.store import TaskPlaneStore


IDENTITY = "spiffe://aios/agent/heartbeat"
WORKER_ID = "heartbeat-worker"


def claimed_task():
    store = TaskPlaneStore(":memory:")
    service = TaskPlaneService(
        store,
        policy=PolicyEngine(
            execution_enabled=True,
            identity_capabilities={"*": set(ALL_ACTION_CAPABILITIES)},
        ),
    )
    service.register_worker(
        {
            "worker_id": WORKER_ID,
            "worker_type": "test",
            "identity": IDENTITY,
            "capabilities": ["code.read"],
            "trust_zone": "test",
            "max_concurrency": 1,
            "heartbeat_interval_seconds": 15,
            "supported_protocols": ["aios-task-v1"],
        }
    )
    task = service.create_task(
        {"task_id": "task-heartbeat", "required_capabilities": ["code.read"]}
    )
    claim = service.claim(
        task["task_id"],
        {
            "worker_id": WORKER_ID,
            "actor_identity": IDENTITY,
            "lease_seconds": 30,
        },
    )
    return store, task, claim


def test_cancelled_task_cannot_renew_a_still_active_claim():
    store, task, claim = claimed_task()
    original_lease = store.get_claim(claim["claim_id"]).lease_until
    current = store.get_task(task["task_id"])
    store.cancel_task_with_event(task["task_id"], current.revision, "owner cancelled")

    with pytest.raises(ConflictError, match="task is terminal"):
        store.heartbeat_claim(task["task_id"], claim["claim_id"], 60)

    assert store.get_claim(claim["claim_id"]).lease_until == original_lease
    assert store.get_task(task["task_id"]).status == "CANCELLED"


def test_lease_at_current_time_is_expired_and_cannot_be_renewed():
    store, task, claim = claimed_task()
    boundary = store.get_claim(claim["claim_id"]).lease_until

    with mock.patch("task_plane.store._now", return_value=boundary):
        with pytest.raises(LockedError, match="lease expired"):
            store.heartbeat_claim(task["task_id"], claim["claim_id"], 30)

    assert store.get_claim(claim["claim_id"]).lease_until == boundary