"""Task queue with atomic claiming for multi-agent work distribution.

This module implements a **file-system–based task queue** that allows an
orchestrator to enqueue work items and N competing agents to claim them
without any central coordinator or external service.

The fundamental invariant is::

    For any task, at most one agent holds an active claim at any time.

Atomic claiming
---------------
The atomic claiming primitive is ``open(path, O_CREAT | O_EXCL | O_WRONLY)``.
On POSIX-compliant file systems (including all common local filesystems and
most NFS mounts), this syscall either:

* **Succeeds** — creating the file and returning the fd (the callee won the
  race).
* **Fails with ENOENT/FileExistsError** — another agent created the file first
  (the callee lost the race).

No lock file, no mutex, no advisory lock, no external service is needed.

Re-claiming timed-out tasks
----------------------------
When an agent crashes or is killed while holding a claim, the claim's
``expires_at`` timestamp passes without a ``complete`` or ``fail`` call.
Such claims are called **timed-out** and the task becomes eligible for
re-claiming.

Re-claiming uses an **optimistic concurrency** approach:

1. Write a new claim via ``write_text_atomic`` (atomic rename over the
   existing claim file) embedding a unique ``claim_nonce``.
2. Read back the file and verify the ``claim_nonce`` matches what we wrote.
3. If the nonce matches → we won.  If not → another agent beat us; continue
   to the next candidate.

This gives at-most-once semantics for the common case (local filesystem,
sequential renames) with a tiny window of ambiguity under extreme concurrency.
The window closes immediately on read-back verification.

Directory layout
----------------
::

    .muse/coordination/
      tasks/     — task definitions (write-once, immutable after enqueue)
          <task-id>.json
      claims/    — claim records (O_CREAT|O_EXCL for initial claim;
                               write_text_atomic for updates)
          <task-id>.json

Derived task status (from file existence + claim.status)
---------------------------------------------------------
``"pending"``
    Task file exists, no claim file.

``"claimed"``
    Claim exists, ``status == "claimed"`` and ``expires_at > now``.

``"timed_out"``
    Claim exists, ``status == "claimed"`` and ``expires_at <= now``.
    The task is eligible for re-claiming.

``"completed"``
    Claim exists, ``status == "completed"``.

``"failed"``
    Claim exists, ``status == "failed"``.

``"cancelled"``
    Claim exists, ``status == "cancelled"``.

Security
--------
Every ``task_id`` is validated as a UUID before constructing any file path,
mirroring :func:`~muse.core.coordination._validate_reservation_id`.  A valid
UUID cannot contain ``/`` or ``..``, which prevents directory traversal
attacks.

Performance
-----------
``claim_next_task`` scans the ``tasks/`` directory once (O(n)) to build a
priority-sorted candidate list, then attempts O_EXCL claim in order.  On a
quiet queue the first attempt almost always succeeds.  On a hot queue
(many agents competing for few tasks) each agent tries at most *n* attempts,
but the total number of attempts across all agents is bounded by the number
of pending tasks × the number of competing agents — no spin loops.
"""

from __future__ import annotations

import dataclasses
import datetime
import json
import logging
import os
import pathlib
import re
import uuid as _uuid_mod

from muse._version import __version__
from muse.core.store import write_text_atomic


type _ClaimMap = dict[str, "ClaimRecord"]
# JSON-compatible value type for arbitrary task payloads and results.
JsonValue = str | int | float | bool | None | list["JsonValue"] | "JsonDict"
JsonDict = dict[str, JsonValue]

logger = logging.getLogger(__name__)

# ── Constants ─────────────────────────────────────────────────────────────────

# Valid task statuses stored in a claim record.
_CLAIM_STATUSES: frozenset[str] = frozenset(
    ("claimed", "completed", "failed", "cancelled", "timed_out")
)

# Derived task status values (not stored; computed from file state).
_TASK_STATUSES: frozenset[str] = frozenset(
    ("pending", "claimed", "timed_out", "completed", "failed", "cancelled")
)

# Maximum allowed tag count per task.
_MAX_TAGS = 32

# Maximum tag length in characters.
_MAX_TAG_LEN = 64

# Maximum title length in characters.
_MAX_TITLE_LEN = 256

# Maximum queue name length in characters.
_MAX_QUEUE_LEN = 64

# UUID4 regex — same pattern as coordination.py's _UUID_RE.
_UUID_RE = re.compile(
    r"^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$",
    re.IGNORECASE,
)

# ── Directory helpers ─────────────────────────────────────────────────────────


def _coord_dir(root: pathlib.Path) -> pathlib.Path:
    """Return ``.muse/coordination/`` for *root*."""
    return root / ".muse" / "coordination"


def _tasks_dir(root: pathlib.Path) -> pathlib.Path:
    """Return ``.muse/coordination/tasks/`` for *root*."""
    return _coord_dir(root) / "tasks"


def _claims_dir(root: pathlib.Path) -> pathlib.Path:
    """Return ``.muse/coordination/claims/`` for *root*."""
    return _coord_dir(root) / "claims"


def ensure_task_dirs(root: pathlib.Path) -> None:
    """Create both task-queue subdirectories if they do not exist."""
    _tasks_dir(root).mkdir(parents=True, exist_ok=True)
    _claims_dir(root).mkdir(parents=True, exist_ok=True)


# ── Input validation ──────────────────────────────────────────────────────────


def _validate_task_id(task_id: str) -> None:
    """Raise :exc:`ValueError` if *task_id* is not a well-formed UUID4.

    Security note
    -------------
    ``task_id`` is used to construct file paths inside
    ``.muse/coordination/tasks/`` and ``.muse/coordination/claims/``.
    A valid UUID cannot contain ``/`` or ``..``, preventing directory
    traversal attacks where a crafted ID like ``../../etc/passwd`` would
    escape the coordination directory.
    """
    if not _UUID_RE.match(str(task_id)):
        raise ValueError(f"task_id must be a valid UUID4: {task_id!r}")


def _validate_queue_name(queue: str) -> None:
    """Raise :exc:`ValueError` if *queue* is not a safe queue name.

    Queue names must be non-empty, at most :data:`_MAX_QUEUE_LEN` characters,
    and contain only ASCII alphanumerics, hyphens, and underscores.
    """
    if not queue:
        raise ValueError("queue name must be non-empty")
    if len(queue) > _MAX_QUEUE_LEN:
        raise ValueError(
            f"queue name too long: {len(queue)} > {_MAX_QUEUE_LEN} chars"
        )
    if not re.match(r"^[a-zA-Z0-9_-]+$", queue):
        raise ValueError(
            f"queue name may only contain [a-zA-Z0-9_-]: {queue!r}"
        )


# ── Time helpers ──────────────────────────────────────────────────────────────


def _now_utc() -> datetime.datetime:
    return datetime.datetime.now(datetime.timezone.utc)


def _parse_dt(s: str) -> datetime.datetime:
    return datetime.datetime.fromisoformat(s)


# ── Data model ────────────────────────────────────────────────────────────────


@dataclasses.dataclass
class TaskRecord:
    """An immutable task definition stored in ``tasks/<task-id>.json``.

    A :class:`TaskRecord` is written once (by the enqueuing agent) and never
    modified.  All mutable state lives in :class:`ClaimRecord`.

    Attributes
    ----------
    task_id:
        UUID4 that uniquely identifies this task.
    title:
        Short human-readable description (≤ :data:`_MAX_TITLE_LEN` chars).
    payload:
        Arbitrary JSON-serialisable dict.  Typically contains addresses,
        operation type, or other agent-specific parameters.
    priority:
        Integer priority — **higher values are processed first**.  Ties are
        broken by ``created_at`` (FIFO: earliest first).
    queue:
        Logical queue name.  Agents claiming with ``queue="billing"`` only see
        billing tasks.  Must match ``[a-zA-Z0-9_-]+``.
    created_at:
        UTC datetime when the task was enqueued.
    created_by:
        ``run_id`` of the agent or orchestrator that enqueued this task.
    ttl_seconds:
        How long (in seconds) from ``created_at`` the task may remain
        ``pending`` before it is considered expired and skipped during
        claiming.  Completed/failed tasks are not affected.
    tags:
        List of string labels for filtering or analytics.
    """

    task_id: str
    title: str
    payload: JsonDict
    priority: int
    queue: str
    created_at: datetime.datetime
    created_by: str
    ttl_seconds: int
    tags: list[str]

    def to_dict(self) -> JsonDict:
        """Serialise to a JSON-compatible dict."""
        return {
            "schema_version": __version__,
            "task_id": self.task_id,
            "title": self.title,
            "payload": self.payload,
            "priority": self.priority,
            "queue": self.queue,
            "created_at": self.created_at.isoformat(),
            "created_by": self.created_by,
            "ttl_seconds": self.ttl_seconds,
            "tags": self.tags,
        }

    @classmethod
    def from_dict(cls, d: JsonDict) -> "TaskRecord":
        """Deserialise from a JSON-compatible dict.

        Missing fields default to safe values so records written by older
        versions remain loadable.
        """
        return cls(
            task_id=str(d.get("task_id", "")),
            title=str(d.get("title", ""))[:_MAX_TITLE_LEN],
            payload=dict(d.get("payload") or {}),
            priority=int(d.get("priority", 0)),
            queue=str(d.get("queue", "default")),
            created_at=_parse_dt(str(d["created_at"])) if "created_at" in d else _now_utc(),
            created_by=str(d.get("created_by", "")),
            ttl_seconds=int(d.get("ttl_seconds", 3600)),
            tags=[str(t) for t in (d.get("tags") or [])],
        )

    def is_expired(self, now: datetime.datetime | None = None) -> bool:
        """Return True if the task's pending TTL has passed."""
        if now is None:
            now = _now_utc()
        return now >= self.created_at + datetime.timedelta(seconds=self.ttl_seconds)


@dataclasses.dataclass
class ClaimRecord:
    """Mutable claim state for a task stored in ``claims/<task-id>.json``.

    A :class:`ClaimRecord` is **created atomically** via
    ``open(O_CREAT | O_EXCL)`` by the claiming agent, then **updated** via
    :func:`~muse.core.store.write_text_atomic` (atomic rename) as the task
    progresses through its lifecycle.

    Attributes
    ----------
    task_id:
        UUID4 of the task being claimed.
    claimer_run_id:
        ``run_id`` of the agent that claimed the task.
    claimed_at:
        UTC datetime when the claim was created.
    expires_at:
        UTC datetime after which the claim is considered timed-out and the
        task becomes eligible for re-claiming.
    status:
        One of ``"claimed"``, ``"completed"``, ``"failed"``,
        ``"cancelled"``, ``"timed_out"``.
    heartbeat_at:
        UTC datetime of the most recent heartbeat.  Agents should call
        :func:`heartbeat_claim` every ``claim_ttl_seconds / 2`` to prevent
        expiry.
    claim_nonce:
        UUID4 generated at claim time.  Used for **optimistic re-claim
        verification**: after overwriting a timed-out claim, the new claimer
        reads back the file and checks that the nonce matches, confirming
        they won the race.
    result:
        Optional JSON-serialisable dict set by :func:`complete_task`.
    error:
        Optional human-readable error message set by :func:`fail_task`.
    """

    task_id: str
    claimer_run_id: str
    claimed_at: datetime.datetime
    expires_at: datetime.datetime
    status: str
    heartbeat_at: datetime.datetime
    claim_nonce: str
    result: JsonDict | None
    error: str | None

    def to_dict(self) -> JsonDict:
        """Serialise to a JSON-compatible dict."""
        return {
            "schema_version": __version__,
            "task_id": self.task_id,
            "claimer_run_id": self.claimer_run_id,
            "claimed_at": self.claimed_at.isoformat(),
            "expires_at": self.expires_at.isoformat(),
            "status": self.status,
            "heartbeat_at": self.heartbeat_at.isoformat(),
            "claim_nonce": self.claim_nonce,
            "result": self.result,
            "error": self.error,
        }

    @classmethod
    def from_dict(cls, d: JsonDict) -> "ClaimRecord":
        """Deserialise from a JSON-compatible dict."""
        return cls(
            task_id=str(d.get("task_id", "")),
            claimer_run_id=str(d.get("claimer_run_id", "")),
            claimed_at=_parse_dt(str(d["claimed_at"])),
            expires_at=_parse_dt(str(d["expires_at"])),
            status=str(d.get("status", "claimed")),
            heartbeat_at=_parse_dt(str(d.get("heartbeat_at", d["claimed_at"]))),
            claim_nonce=str(d.get("claim_nonce", str(_uuid_mod.uuid4()))),
            result=d.get("result"),
            error=d.get("error"),
        )

    def is_expired(self, now: datetime.datetime | None = None) -> bool:
        """Return True if this claim's TTL has passed without renewal."""
        if now is None:
            now = _now_utc()
        return now >= self.expires_at


# ── Loading helpers ───────────────────────────────────────────────────────────


def load_all_tasks(root: pathlib.Path) -> list[TaskRecord]:
    """Return all task records from ``tasks/``.

    Corrupt or unreadable files are skipped with a warning.  Expired tasks
    are included — the caller decides whether to filter them.
    """
    tdir = _tasks_dir(root)
    if not tdir.exists():
        return []
    tasks: list[TaskRecord] = []
    for path in tdir.glob("*.json"):
        try:
            raw = json.loads(path.read_text())
            tasks.append(TaskRecord.from_dict(raw))
        except (json.JSONDecodeError, KeyError, ValueError) as exc:
            logger.warning("Corrupt task file %s: %s", path.name, exc)
    return tasks


def load_all_claims(root: pathlib.Path) -> _ClaimMap:
    """Return a mapping of ``task_id → ClaimRecord`` from ``claims/``.

    Corrupt or unreadable files are skipped with a warning.
    """
    cdir = _claims_dir(root)
    if not cdir.exists():
        return {}
    claims: _ClaimMap = {}
    for path in cdir.glob("*.json"):
        try:
            raw = json.loads(path.read_text())
            claim = ClaimRecord.from_dict(raw)
            claims[claim.task_id] = claim
        except (json.JSONDecodeError, KeyError, ValueError) as exc:
            logger.warning("Corrupt claim file %s: %s", path.name, exc)
    return claims


def load_task(root: pathlib.Path, task_id: str) -> TaskRecord | None:
    """Load and return the task with *task_id*, or ``None`` if not found."""
    _validate_task_id(task_id)
    path = _tasks_dir(root) / f"{task_id}.json"
    try:
        return TaskRecord.from_dict(json.loads(path.read_text()))
    except (OSError, json.JSONDecodeError, KeyError, ValueError):
        return None


def load_claim(root: pathlib.Path, task_id: str) -> ClaimRecord | None:
    """Load and return the claim for *task_id*, or ``None`` if unclaimed."""
    _validate_task_id(task_id)
    path = _claims_dir(root) / f"{task_id}.json"
    try:
        return ClaimRecord.from_dict(json.loads(path.read_text()))
    except (OSError, json.JSONDecodeError, KeyError, ValueError):
        return None


def get_task_status(
    task: TaskRecord,
    claim: ClaimRecord | None,
    now: datetime.datetime | None = None,
) -> str:
    """Derive the current status string for *task* given its *claim*.

    Parameters
    ----------
    task:
        The task definition.
    claim:
        The claim record if one exists, else ``None``.
    now:
        Current UTC time.  Defaults to :func:`_now_utc`.

    Returns
    -------
    str
        One of ``"pending"``, ``"claimed"``, ``"timed_out"``,
        ``"completed"``, ``"failed"``, ``"cancelled"``.
    """
    if now is None:
        now = _now_utc()
    if claim is None:
        return "pending"
    if claim.status == "claimed":
        return "timed_out" if claim.is_expired(now) else "claimed"
    return claim.status  # completed | failed | cancelled | timed_out


# ── Core operations ───────────────────────────────────────────────────────────


def create_task(
    root: pathlib.Path,
    title: str,
    *,
    payload: JsonDict | None = None,
    priority: int = 0,
    queue: str = "default",
    ttl_seconds: int = 86400,
    created_by: str = "unknown",
    tags: list[str] | None = None,
) -> TaskRecord:
    """Enqueue a new task and return its :class:`TaskRecord`.

    Parameters
    ----------
    root:
        Repository root (directory containing ``.muse/``).
    title:
        Short human-readable description.  Truncated to
        :data:`_MAX_TITLE_LEN` characters.
    payload:
        Arbitrary JSON-serialisable dict attached to the task.  Agents use
        this to carry domain-specific parameters (addresses, operation type,
        branch names, etc.).
    priority:
        Integer priority.  **Higher = processed first** by
        :func:`claim_next_task`.  Default: ``0``.
    queue:
        Logical queue name.  Must match ``[a-zA-Z0-9_-]+``.  Default:
        ``"default"``.
    ttl_seconds:
        Seconds from ``created_at`` after which the task is no longer
        eligible for claiming.  Default: 86 400 s (24 h).
    created_by:
        ``run_id`` of the enqueuing agent.
    tags:
        Optional list of string labels.

    Returns
    -------
    TaskRecord
        The newly created task.

    Raises
    ------
    ValueError
        If *queue* is not a valid queue name or *title* is empty.
    """
    if not title:
        raise ValueError("task title must be non-empty")
    _validate_queue_name(queue)
    ensure_task_dirs(root)

    safe_tags = [str(t)[:_MAX_TAG_LEN] for t in (tags or [])][:_MAX_TAGS]

    task = TaskRecord(
        task_id=str(_uuid_mod.uuid4()),
        title=title[:_MAX_TITLE_LEN],
        payload=dict(payload or {}),
        priority=int(priority),
        queue=queue,
        created_at=_now_utc(),
        created_by=str(created_by),
        ttl_seconds=max(1, int(ttl_seconds)),
        tags=safe_tags,
    )
    path = _tasks_dir(root) / f"{task.task_id}.json"
    write_text_atomic(path, json.dumps(task.to_dict(), indent=2) + "\n")
    logger.debug(
        "Enqueued task %s (priority=%d, queue=%s, ttl=%ds)",
        task.task_id[:8], task.priority, task.queue, task.ttl_seconds,
    )
    return task


def claim_next_task(
    root: pathlib.Path,
    run_id: str,
    *,
    queue: str | None = None,
    claim_ttl_seconds: int = 3600,
) -> tuple[TaskRecord, ClaimRecord] | None:
    """Atomically claim the highest-priority pending task.

    Scans all tasks in *queue* (or all queues if ``None``), sorts by priority
    (descending) then by ``created_at`` (ascending, FIFO within same priority),
    and tries to claim each one in order.

    **Initial claim** (no existing claim file): uses ``O_CREAT | O_EXCL`` for
    a guaranteed-atomic, winner-takes-all claim.

    **Re-claim of timed-out task** (existing claim with ``expires_at`` in the
    past): uses ``write_text_atomic`` (atomic rename) followed by read-back
    nonce verification.  The agent whose nonce survives the rename wins.  In
    the rare case of a tie (two agents renaming within the same kernel
    operation), the losing agent detects the mismatch on read-back and moves
    to the next candidate.

    Parameters
    ----------
    root:
        Repository root.
    run_id:
        The claiming agent's identifier.
    queue:
        If set, only claim tasks from this queue.  If ``None``, claim from
        any queue (sorted globally by priority + FIFO).
    claim_ttl_seconds:
        How long (seconds) the claim is valid before another agent may
        re-claim the task.  Agents should call :func:`heartbeat_claim`
        every ``claim_ttl_seconds / 2`` to prevent expiry.

    Returns
    -------
    (TaskRecord, ClaimRecord) | None
        The claimed task and its claim record, or ``None`` if the queue
        contains no pending/claimable tasks.
    """
    ensure_task_dirs(root)
    now = _now_utc()
    all_tasks = load_all_tasks(root)
    all_claims = load_all_claims(root)

    # Build the candidate list: unclaimed OR timed-out, correct queue, not expired TTL.
    candidates: list[TaskRecord] = []
    for task in all_tasks:
        if queue is not None and task.queue != queue:
            continue
        if task.is_expired(now):
            continue  # Task TTL exhausted — skip.
        claim = all_claims.get(task.task_id)
        if claim is None:
            candidates.append(task)  # Unclaimed.
        elif claim.status == "claimed" and claim.is_expired(now):
            candidates.append(task)  # Timed-out — eligible for re-claim.

    # Sort: highest priority first, then FIFO (earliest created_at).
    candidates.sort(key=lambda t: (-t.priority, t.created_at))

    for task in candidates:
        existing = all_claims.get(task.task_id)
        if existing is None:
            # ── Initial claim via O_CREAT | O_EXCL ───────────────────────────
            result = _try_excl_claim(root, task.task_id, run_id, now, claim_ttl_seconds)
        else:
            # ── Re-claim of timed-out task via optimistic write + verify ─────
            result = _try_optimistic_reclaim(
                root, task.task_id, run_id, now, claim_ttl_seconds
            )
        if result is not None:
            logger.debug(
                "Agent %s claimed task %s (queue=%s, priority=%d)",
                run_id, task.task_id[:8], task.queue, task.priority,
            )
            return task, result

    return None


def _try_excl_claim(
    root: pathlib.Path,
    task_id: str,
    run_id: str,
    now: datetime.datetime,
    claim_ttl_seconds: int,
) -> ClaimRecord | None:
    """Attempt to claim *task_id* via ``O_CREAT | O_EXCL``.

    Returns the :class:`ClaimRecord` on success, or ``None`` if another agent
    claimed the task first (:exc:`FileExistsError`).

    The claim file is written with an fsync before close to ensure durability
    even if the process is killed immediately after claiming.
    """
    claim_path = _claims_dir(root) / f"{task_id}.json"
    nonce = str(_uuid_mod.uuid4())
    claim = ClaimRecord(
        task_id=task_id,
        claimer_run_id=run_id,
        claimed_at=now,
        expires_at=now + datetime.timedelta(seconds=claim_ttl_seconds),
        status="claimed",
        heartbeat_at=now,
        claim_nonce=nonce,
        result=None,
        error=None,
    )
    content = (json.dumps(claim.to_dict(), indent=2) + "\n").encode()
    try:
        fd = os.open(
            str(claim_path),
            os.O_CREAT | os.O_EXCL | os.O_WRONLY,
            0o644,
        )
        try:
            os.write(fd, content)
            try:
                os.fsync(fd)
            except OSError:
                pass  # Best-effort durability; atomicity still holds.
        finally:
            os.close(fd)
        return claim
    except FileExistsError:
        return None  # Another agent got there first.


def _try_optimistic_reclaim(
    root: pathlib.Path,
    task_id: str,
    run_id: str,
    now: datetime.datetime,
    claim_ttl_seconds: int,
) -> ClaimRecord | None:
    """Re-claim a timed-out task using optimistic concurrency.

    Overwrites the existing (expired) claim file via ``write_text_atomic``
    (atomic rename), then reads it back to verify the ``claim_nonce`` matches
    what we wrote.  The agent whose write survives the atomic rename wins.

    Returns the :class:`ClaimRecord` if we won the race, or ``None`` if
    another agent's write arrived later.

    Security note
    -------------
    This function does NOT use ``O_EXCL`` — it is possible (though unlikely on
    a local filesystem) for two agents to both believe they won in a very tight
    race window.  The read-back verification closes this window to the duration
    of a single file stat + read, which is microseconds on local storage.
    Callers should treat this as best-effort at-most-once (not strict exactly-
    once) for the re-claim case.
    """
    claim_path = _claims_dir(root) / f"{task_id}.json"
    nonce = str(_uuid_mod.uuid4())
    claim = ClaimRecord(
        task_id=task_id,
        claimer_run_id=run_id,
        claimed_at=now,
        expires_at=now + datetime.timedelta(seconds=claim_ttl_seconds),
        status="claimed",
        heartbeat_at=now,
        claim_nonce=nonce,
        result=None,
        error=None,
    )
    write_text_atomic(claim_path, json.dumps(claim.to_dict(), indent=2) + "\n")
    # Read back and verify nonce.
    try:
        actual = json.loads(claim_path.read_text())
        if actual.get("claim_nonce") == nonce:
            return claim
    except (OSError, json.JSONDecodeError):
        pass
    return None


def complete_task(
    root: pathlib.Path,
    task_id: str,
    run_id: str,
    *,
    result: JsonDict | None = None,
) -> ClaimRecord:
    """Mark a claimed task as completed.

    Parameters
    ----------
    root:
        Repository root.
    task_id:
        UUID of the task to complete.
    run_id:
        The claiming agent's ``run_id``.  Must match ``claim.claimer_run_id``.
    result:
        Optional JSON-serialisable dict describing the outcome (e.g. proposal URL,
        output summary, created artefacts).

    Returns
    -------
    ClaimRecord
        The updated claim record with ``status="completed"``.

    Raises
    ------
    ValueError
        If *task_id* is not a valid UUID.
    FileNotFoundError
        If the task or claim does not exist.
    PermissionError
        If *run_id* does not match the claimer.
    RuntimeError
        If the task is not in a claimable state (already completed, failed, etc.).
    """
    _validate_task_id(task_id)
    return _update_claim_status(
        root, task_id, run_id, "completed",
        result=result, error=None,
    )


def fail_task(
    root: pathlib.Path,
    task_id: str,
    run_id: str,
    *,
    error: str = "",
) -> ClaimRecord:
    """Mark a claimed task as failed.

    Parameters
    ----------
    root:
        Repository root.
    task_id:
        UUID of the task to fail.
    run_id:
        The claiming agent's ``run_id``.  Must match ``claim.claimer_run_id``.
    error:
        Human-readable error message describing why the task failed.

    Returns
    -------
    ClaimRecord
        The updated claim record with ``status="failed"``.

    Raises
    ------
    ValueError
        If *task_id* is not a valid UUID.
    FileNotFoundError
        If the task or claim does not exist.
    PermissionError
        If *run_id* does not match the claimer.
    RuntimeError
        If the task is not in ``"claimed"`` status.
    """
    _validate_task_id(task_id)
    return _update_claim_status(
        root, task_id, run_id, "failed",
        result=None, error=str(error),
    )


def cancel_task(
    root: pathlib.Path,
    task_id: str,
    run_id: str,
    *,
    force: bool = False,
) -> ClaimRecord:
    """Cancel a pending or claimed task.

    A **pending** task (no claim file) is cancelled by creating a new claim
    record directly with ``status="cancelled"`` via ``O_CREAT | O_EXCL``.
    If another agent claims the task between the pending-check and the cancel
    write, the cancel fails with :exc:`FileExistsError` (the task is now
    claimed and must be cancelled via the claimer or ``force=True``).

    A **claimed** task may only be cancelled by its claimer (``run_id ==
    claim.claimer_run_id``) unless ``force=True``.

    Parameters
    ----------
    root:
        Repository root.
    task_id:
        UUID of the task to cancel.
    run_id:
        The calling agent's ``run_id``.
    force:
        If ``True``, cancel even if the task is claimed by a different agent.
        Use with caution — this aborts in-flight work without notifying the
        claimer.

    Returns
    -------
    ClaimRecord
        The resulting claim record with ``status="cancelled"``.

    Raises
    ------
    ValueError
        If *task_id* is not a valid UUID.
    FileNotFoundError
        If the task does not exist.
    RuntimeError
        If the task is already in a terminal state (completed, failed, cancelled).
    PermissionError
        If the task is claimed by a different agent and ``force=False``.
    """
    _validate_task_id(task_id)
    ensure_task_dirs(root)

    task_path = _tasks_dir(root) / f"{task_id}.json"
    if not task_path.exists():
        raise FileNotFoundError(f"task not found: {task_id}")

    claim = load_claim(root, task_id)
    now = _now_utc()

    if claim is None:
        # Pending task — cancel it via O_EXCL (prevent race with a claiming agent).
        cancel_claim = ClaimRecord(
            task_id=task_id,
            claimer_run_id=run_id,
            claimed_at=now,
            expires_at=now + datetime.timedelta(hours=24),
            status="cancelled",
            heartbeat_at=now,
            claim_nonce=str(_uuid_mod.uuid4()),
            result=None,
            error="cancelled before claiming",
        )
        claim_path = _claims_dir(root) / f"{task_id}.json"
        try:
            fd = os.open(str(claim_path), os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644)
            try:
                os.write(fd, (json.dumps(cancel_claim.to_dict(), indent=2) + "\n").encode())
                try:
                    os.fsync(fd)
                except OSError:
                    pass
            finally:
                os.close(fd)
        except FileExistsError as exc:
            raise FileExistsError(
                f"task {task_id} was claimed by another agent before cancel; "
                f"use force=True to cancel a claimed task"
            ) from exc
        return cancel_claim

    # Claimed / terminal task.
    if claim.status in ("completed", "failed", "cancelled"):
        raise RuntimeError(
            f"task {task_id} is already in terminal state: {claim.status!r}"
        )
    if not force and claim.claimer_run_id != run_id:
        raise PermissionError(
            f"task {task_id} is claimed by {claim.claimer_run_id!r}; "
            f"only the claimer or force=True can cancel it"
        )

    return _update_claim_status(
        root, task_id, run_id if force else claim.claimer_run_id,
        "cancelled", result=None, error="cancelled",
        _skip_ownership_check=True,
    )


def heartbeat_claim(
    root: pathlib.Path,
    task_id: str,
    run_id: str,
    *,
    extension_seconds: int = 3600,
) -> ClaimRecord:
    """Extend the claim TTL by writing a fresh heartbeat.

    Agents should call this every ``claim_ttl_seconds / 2`` to prevent their
    claim from timing out.  Only the current claimer may heartbeat a claim.

    Parameters
    ----------
    root:
        Repository root.
    task_id:
        UUID of the claimed task.
    run_id:
        The claiming agent's ``run_id``.  Must match ``claim.claimer_run_id``.
    extension_seconds:
        How far into the future to set the new ``expires_at``.
        Default: 3600 s (1 h).

    Returns
    -------
    ClaimRecord
        The updated claim record with refreshed ``heartbeat_at`` and
        ``expires_at``.

    Raises
    ------
    ValueError
        If *task_id* is not a valid UUID.
    FileNotFoundError
        If no claim exists for *task_id*.
    PermissionError
        If *run_id* does not match the claimer.
    RuntimeError
        If the task is not in ``"claimed"`` status.
    """
    _validate_task_id(task_id)
    claim = load_claim(root, task_id)
    if claim is None:
        raise FileNotFoundError(f"no claim found for task {task_id}")
    if claim.claimer_run_id != run_id:
        raise PermissionError(
            f"task {task_id} is claimed by {claim.claimer_run_id!r}, not {run_id!r}"
        )
    if claim.status != "claimed":
        raise RuntimeError(
            f"cannot heartbeat task {task_id} with status {claim.status!r}"
        )
    now = _now_utc()
    claim.heartbeat_at = now
    claim.expires_at = now + datetime.timedelta(seconds=extension_seconds)
    claim_path = _claims_dir(root) / f"{task_id}.json"
    write_text_atomic(claim_path, json.dumps(claim.to_dict(), indent=2) + "\n")
    return claim


# ── Private helpers ───────────────────────────────────────────────────────────


def _update_claim_status(
    root: pathlib.Path,
    task_id: str,
    run_id: str,
    new_status: str,
    *,
    result: JsonDict | None,
    error: str | None,
    _skip_ownership_check: bool = False,
) -> ClaimRecord:
    """Update a claim's status field via atomic rename.

    Internal helper used by :func:`complete_task`, :func:`fail_task`, and
    :func:`cancel_task`.

    Raises
    ------
    FileNotFoundError
        If the task file or claim file does not exist.
    PermissionError
        If *run_id* does not match the claimer (unless *_skip_ownership_check*).
    RuntimeError
        If the claim is not in a state that permits the transition.
    """
    task_path = _tasks_dir(root) / f"{task_id}.json"
    if not task_path.exists():
        raise FileNotFoundError(f"task not found: {task_id}")

    claim = load_claim(root, task_id)
    if claim is None:
        raise FileNotFoundError(f"no claim found for task {task_id}")

    if not _skip_ownership_check and claim.claimer_run_id != run_id:
        raise PermissionError(
            f"task {task_id} is claimed by {claim.claimer_run_id!r}, not {run_id!r}"
        )

    # Only "claimed" tasks can transition to completed/failed/cancelled.
    # (A timed_out task is still status="claimed" on disk — it's just expired.)
    if claim.status not in ("claimed",):
        raise RuntimeError(
            f"task {task_id} cannot transition from {claim.status!r} to {new_status!r}"
        )

    claim.status = new_status
    claim.result = result
    claim.error = error

    claim_path = _claims_dir(root) / f"{task_id}.json"
    write_text_atomic(claim_path, json.dumps(claim.to_dict(), indent=2) + "\n")
    logger.debug("Task %s → %s (by %s)", task_id[:8], new_status, run_id)
    return claim