"""Claiming, drafts and worklists for annotation nodes. See SPEC.md."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timedelta

REVIEW_STEPS = {"review", "qa"}


class Conflict(Exception): ...
class Forbidden(Exception): ...
class Stale(Exception): ...


@dataclass
class Task:
    id: str
    project: str
    status: str = "open"          # open | paused | done


@dataclass
class Node:
    id: str
    task_id: str
    step: str                     # author | review | qa | ...
    order: int                    # position of the step in the task's flow, 0-based
    status: str = "parked"        # parked | running | done
    assignee: str | None = None
    claimed_at: datetime | None = None
    revision: int = 0
    draft: str = ""
    actors: list[str] = field(default_factory=list)   # everyone who ever worked this node


@dataclass
class Store:
    tasks: dict[str, Task]
    nodes: dict[str, Node]
    seats: dict[tuple[str, str], set[str]]             # (user, project) -> steps they may claim
    idle_hours: dict[str, int]                         # project -> hours


def is_idle_expired(node: Node, store: Store, now: datetime) -> bool:
    if node.status != "running" or node.claimed_at is None:
        return False
    project = store.tasks[node.task_id].project
    limit = store.idle_hours.get(project, 24)
    return (now - node.claimed_at).seconds > limit * 3600


def prior_actors(node: Node, store: Store) -> set[str]:
    earlier = [n for n in store.nodes.values() if n.task_id == node.task_id and n.order < node.order]
    earlier.sort(key=lambda n: n.order)
    return set(earlier[-1].actors) if earlier else set()


def eligible(node: Node, user: str, store: Store) -> bool:
    task = store.tasks[node.task_id]
    if node.step not in store.seats.get((user, task.project), set()):
        return False
    if node.step in REVIEW_STEPS and user in prior_actors(node, store):
        return False
    return True


def claim(node_id: str, user: str, store: Store, now: datetime) -> Node:
    node = store.nodes[node_id]
    task = store.tasks[node.task_id]
    if task.status != "open":
        raise Conflict(f"task {task.id} is {task.status}")
    if node.status == "running" and not is_idle_expired(node, store, now):
        if node.assignee == user:
            return node                                   # resume
        raise Conflict(f"node {node_id} already claimed by {node.assignee}")
    if not eligible(node, user, store):
        raise Forbidden(f"{user} may not claim {node.step} on task {task.id}")
    node.status = "running"
    node.assignee = user
    node.claimed_at = now
    if user not in node.actors:
        node.actors.append(user)
    return node


def save_draft(node_id: str, user: str, content: str, expected_revision: int, store: Store) -> int:
    node = store.nodes[node_id]
    if node.assignee != user:
        raise Forbidden(f"{user} is not the assignee of {node_id}")
    if expected_revision > node.revision:
        raise Stale(f"expected {expected_revision}, current {node.revision}")
    node.draft = content
    node.revision += 1
    return node.revision


def worklist(user: str, store: Store, now: datetime) -> dict[str, list[str]]:
    waiting, available = [], []
    for node in store.nodes.values():
        if node.status == "done":
            continue
        expired = is_idle_expired(node, store, now)
        if node.status == "running" and node.assignee == user and not expired:
            waiting.append(node.id)
        elif (node.status == "parked" or expired) and eligible(node, user, store):
            available.append(node.id)
    return {"waiting_on_you": sorted(waiting), "available": sorted(available)}
