from typing import Literal


TaskStatus = Literal[
    "pending",
    "in_progress",
    "completed",
    "failed",
    "cancelled",
]

TASK_STATUSES = {"pending", "in_progress", "completed", "failed", "cancelled"}

create -> pending -> ready -> claim -> in_progress -> checkpoint -> completed
                                                                  -> 解锁下游任务

.agent/
└── tasks/
    ├── analyze.json
    ├── prepare-data.json
    ├── build-api.json
    ├── build-ui.json
    └── integration-test.json

{
  "schema_version": 1,
  "id": "build-api",
  "title": "实现搜索 API",
  "description": "实现搜索接口，并补充参数校验和单元测试。",
  "status": "in_progress",
  "dependencies": ["analyze"],
  "claimed_by": "agent-backend-01",
  "claim_token": "claim-7f3c2a",
  "lease_expires_at": "2026-07-09T10:35:00+00:00",
  "checkpoint": "接口已完成，正在补充分页测试。",
  "artifacts": ["src/search.py", "tests/test_search.py"],
  "result": null,
  "error": null,
  "created_at": "2026-07-09T10:00:00+00:00",
  "updated_at": "2026-07-09T10:30:00+00:00"
}

from dataclasses import asdict, dataclass, field


@dataclass
class Task:
    id: str
    title: str
    description: str
    status: TaskStatus = "pending"
    dependencies: list[str] = field(default_factory=list)
    claimed_by: str | None = None
    claim_token: str | None = None
    lease_expires_at: str | None = None
    checkpoint: str | None = None
    artifacts: list[str] = field(default_factory=list)
    result: str | None = None
    error: str | None = None
    created_at: str = ""
    updated_at: str = ""
    schema_version: int = 1

    @classmethod
    def from_dict(cls, data: dict) -> "Task":
        if not isinstance(data, dict):
            raise ValueError("task file must contain a JSON object")
        if data.get("schema_version") != 1:
            raise ValueError("unsupported task schema version")

        task = cls(**data)
        if task.status not in TASK_STATUSES:
            raise ValueError(f"invalid task status: {task.status}")
        if not isinstance(task.dependencies, list) or not all(
            isinstance(item, str) for item in task.dependencies
        ):
            raise ValueError("task dependencies must be a list of strings")
        return task

    def to_dict(self) -> dict:
        return asdict(self)

from datetime import datetime, timezone


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


def utc_now_text() -> str:
    return utc_now().isoformat(timespec="seconds")

import fcntl
import json
import os
import re
import tempfile
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator


TASK_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]{0,63}$")


class TaskStoreError(RuntimeError):
    pass


class TaskUnavailableError(TaskStoreError):
    pass


class TaskStore:
    def __init__(self, root: Path):
        self.root = root.resolve()
        self.root.mkdir(parents=True, exist_ok=True)
        self.lock_path = self.root / ".lock"

    def _task_path(self, task_id: str) -> Path:
        if not TASK_ID_PATTERN.fullmatch(task_id):
            raise TaskStoreError(f"invalid task id: {task_id}")
        return self.root / f"{task_id}.json"

    @contextmanager
    def locked(self) -> Iterator[None]:
        with self.lock_path.open("a+", encoding="utf-8") as lock_file:
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
            try:
                yield
            finally:
                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)

    def _load_path_unlocked(self, path: Path) -> Task:
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
            task = Task.from_dict(data)
            if self._task_path(task.id) != path.resolve():
                raise ValueError("task id does not match file name")
            return task
        except (OSError, TypeError, ValueError, TaskStoreError) as exc:
            raise TaskStoreError(f"cannot read task file {path.name}: {exc}") from exc

    def _read_unlocked(self, task_id: str) -> Task:
        path = self._task_path(task_id)
        if not path.exists():
            raise TaskStoreError(f"task does not exist: {task_id}")
        return self._load_path_unlocked(path)

    def _list_unlocked(self) -> list[Task]:
        return [
            self._load_path_unlocked(path)
            for path in sorted(self.root.glob("*.json"))
        ]

    def _write_unlocked(self, task: Task) -> None:
        task.updated_at = utc_now_text()
        target = self._task_path(task.id)
        temp_path: Path | None = None

        try:
            with tempfile.NamedTemporaryFile(
                mode="w",
                encoding="utf-8",
                dir=self.root,
                prefix=f".{task.id}.",
                suffix=".tmp",
                delete=False,
            ) as temp_file:
                json.dump(
                    task.to_dict(),
                    temp_file,
                    ensure_ascii=False,
                    indent=2,
                    sort_keys=True,
                )
                temp_file.write("\n")
                temp_file.flush()
                os.fsync(temp_file.fileno())
                temp_path = Path(temp_file.name)

            os.replace(temp_path, target)
        finally:
            if temp_path is not None and temp_path.exists():
                temp_path.unlink()

    def get(self, task_id: str) -> Task:
        with self.locked():
            return self._read_unlocked(task_id)

    def list(self) -> list[Task]:
        with self.locked():
            return self._list_unlocked()

def create_task(
    store: TaskStore,
    task_id: str,
    title: str,
    description: str,
    dependencies: list[str] | None = None,
) -> Task:
    dependency_ids = sorted(set(dependencies or []))
    normalized_title = title.strip()
    normalized_description = description.strip()

    if not normalized_title:
        raise TaskStoreError("task title is required")
    if not normalized_description:
        raise TaskStoreError("task description is required")

    with store.locked():
        target = store._task_path(task_id)
        if target.exists():
            existing_task = store._read_unlocked(task_id)
            same_definition = (
                existing_task.title == normalized_title
                and existing_task.description == normalized_description
                and existing_task.dependencies == dependency_ids
            )
            if same_definition:
                return existing_task
            raise TaskStoreError(f"task id conflicts with existing task: {task_id}")

        existing = {task.id: task for task in store._list_unlocked()}
        missing = [item for item in dependency_ids if item not in existing]
        if missing:
            raise TaskStoreError(f"unknown dependencies: {missing}")

        if task_id in dependency_ids:
            raise TaskStoreError("task cannot depend on itself")

        now = utc_now_text()
        task = Task(
            id=task_id,
            title=normalized_title,
            description=normalized_description,
            dependencies=dependency_ids,
            created_at=now,
            updated_at=now,
        )
        store._write_unlocked(task)
        return task

def unresolved_dependencies(
    task: Task,
    tasks: dict[str, Task],
) -> list[str]:
    return [
        dependency_id
        for dependency_id in task.dependencies
        if tasks[dependency_id].status != "completed"
    ]


def task_view_state(task: Task, tasks: dict[str, Task]) -> str:
    if task.status == "pending":
        blocked_by = unresolved_dependencies(task, tasks)
        return "blocked" if blocked_by else "ready"

    if task.status == "in_progress" and not lease_is_active(task):
        return "reclaimable"

    return task.status

def list_ready_tasks(store: TaskStore) -> list[Task]:
    tasks = store.list()
    task_map = {task.id: task for task in tasks}

    return [
        task
        for task in tasks
        if task_view_state(task, task_map) in {"ready", "reclaimable"}
    ]

from datetime import timedelta


def lease_is_active(task: Task, now: datetime | None = None) -> bool:
    if not task.lease_expires_at:
        return False

    current = now or utc_now()
    expires_at = datetime.fromisoformat(task.lease_expires_at)
    return expires_at > current

def claim_task(
    store: TaskStore,
    task_id: str,
    agent_id: str,
    claim_token: str,
    lease_seconds: int = 300,
) -> Task:
    claim_token = claim_token.strip()
    if not claim_token:
        raise TaskStoreError("claim token is required")

    lease_seconds = max(30, min(lease_seconds, 3600))

    with store.locked():
        tasks = {task.id: task for task in store._list_unlocked()}
        task = tasks.get(task_id)
        if task is None:
            raise TaskUnavailableError(f"task does not exist: {task_id}")

        if task.status in {"completed", "failed", "cancelled"}:
            raise TaskUnavailableError(f"task cannot be claimed: {task.status}")

        if task.status == "in_progress" and lease_is_active(task):
            same_claim = (
                task.claimed_by == agent_id and task.claim_token == claim_token
            )
            if same_claim:
                return task
            raise TaskUnavailableError(f"task is claimed by {task.claimed_by}")

        blocked_by = unresolved_dependencies(task, tasks)
        if blocked_by:
            raise TaskUnavailableError(f"task is blocked by: {blocked_by}")

        now = utc_now()
        task.status = "in_progress"
        task.claimed_by = agent_id
        task.claim_token = claim_token
        task.lease_expires_at = (
            now + timedelta(seconds=lease_seconds)
        ).isoformat(timespec="seconds")
        store._write_unlocked(task)
        return task

def ensure_task_owner(task: Task, agent_id: str, claim_token: str) -> None:
    if task.status != "in_progress":
        raise TaskStoreError("task is not in progress")
    if task.claimed_by != agent_id:
        raise TaskStoreError(f"task is claimed by {task.claimed_by}")
    if task.claim_token != claim_token:
        raise TaskStoreError("task claim token does not match")
    if not lease_is_active(task):
        raise TaskStoreError("task lease has expired")


def checkpoint_task(
    store: TaskStore,
    task_id: str,
    agent_id: str,
    claim_token: str,
    summary: str,
    artifacts: list[str] | None = None,
    renew_seconds: int = 300,
) -> Task:
    summary = summary.strip()
    renew_seconds = max(30, min(renew_seconds, 3600))
    if not summary:
        raise TaskStoreError("checkpoint summary is required")
    if len(summary) > 4000:
        raise TaskStoreError("checkpoint summary is too long")
    if artifacts is not None and not all(
        isinstance(item, str) and item.strip() for item in artifacts
    ):
        raise TaskStoreError("artifacts must be non-empty strings")

    with store.locked():
        task = store._read_unlocked(task_id)
        ensure_task_owner(task, agent_id, claim_token)

        task.checkpoint = summary
        new_artifacts = [item.strip() for item in artifacts or []]
        task.artifacts = list(
            dict.fromkeys([*task.artifacts, *new_artifacts])
        )[:20]
        task.lease_expires_at = (
            utc_now() + timedelta(seconds=renew_seconds)
        ).isoformat(timespec="seconds")
        store._write_unlocked(task)
        return task

TASK_TOOLS = [
    TASK_CREATE_TOOL,
    TASK_LIST_TOOL,
    TASK_GET_TOOL,
    TASK_CLAIM_TOOL,
    TASK_CHECKPOINT_TOOL,
    TASK_RELEASE_TOOL,
    TASK_COMPLETE_TOOL,
    TASK_FAIL_TOOL,
    TASK_RETRY_TOOL,
]

TOOLS.extend(TASK_TOOLS)
TOOL_HANDLERS.update(
    {
        "task_create": run_task_create,
        "task_list": run_task_list,
        "task_get": run_task_get,
        "task_claim": run_task_claim,
        "task_checkpoint": run_task_checkpoint,
        "task_release": run_task_release,
        "task_complete": run_task_complete,
        "task_fail": run_task_fail,
        "task_retry": run_task_retry,
    }
)

TASK_ROOT = (WORKDIR / ".agent" / "tasks").resolve()


def ensure_not_task_store_path(path: Path) -> None:
    target = path.resolve()
    if target.is_relative_to(TASK_ROOT):
        raise PermissionError("task files must be changed through Task System tools")

def recover_agent_task(
    store: TaskStore,
    agent_id: str,
    claim_token: str,
) -> Task | None:
    tasks = store.list()
    task_map = {task.id: task for task in tasks}

    active = [
        task
        for task in tasks
        if task.status == "in_progress"
        and task.claimed_by == agent_id
        and task.claim_token == claim_token
        and lease_is_active(task)
    ]
    if active:
        return sorted(active, key=lambda task: task.updated_at)[-1]

    own_expired = sorted(
        (
            task
            for task in tasks
            if task.status == "in_progress"
            and task.claimed_by == agent_id
            and not lease_is_active(task)
        ),
        key=lambda task: task.updated_at,
        reverse=True,
    )
    own_expired_ids = {task.id for task in own_expired}

    ready = sorted(
        (
            task
            for task in tasks
            if task_view_state(task, task_map) == "ready"
        ),
        key=lambda task: task.id,
    )
    other_expired = [
        task
        for task in tasks
        if task_view_state(task, task_map) == "reclaimable"
        and task.id not in own_expired_ids
    ]

    for task in [*own_expired, *ready, *other_expired]:
        try:
            return claim_task(
                store=store,
                task_id=task.id,
                agent_id=agent_id,
                claim_token=claim_token,
            )
        except TaskUnavailableError:
            # 状态可能在 list 和 claim 之间被其它 Worker 改变。
            continue

    return None

def render_task_state(
    store: TaskStore,
    agent_id: str,
    claim_token: str,
) -> str:
    tasks = store.list()
    task_map = {task.id: task for task in tasks}

    owned = [
        task
        for task in tasks
        if task.status == "in_progress"
        and task.claimed_by == agent_id
        and task.claim_token == claim_token
        and lease_is_active(task)
    ]
    ready = [
        task.id
        for task in tasks
        if task_view_state(task, task_map) in {"ready", "reclaimable"}
    ][:5]

    lines = [f"Agent id: {agent_id}"]
    if owned:
        current = owned[0]
        lines.append(f"Current task: {current.id} - {current.title}")
        lines.append(f"Checkpoint: {current.checkpoint or 'none'}")
    else:
        lines.append("Current task: none")

    lines.append("Ready tasks: " + (", ".join(ready) or "none"))
    lines.append("Use task_get before starting an unfamiliar task.")
    return "\n".join(lines)