app/queue_store.py

"""The discovery queue (reading frontier) — process state in the central bucket.

One file per source under ``queue/`` (frontmatter carries status + lease). All
mutations run under one process-wide lock so an atomic ``claim`` can't hand the
same source to two agents (the Space is a single uvicorn worker; the lock
serialises the rare claim/skip/processed writes — same reasoning as
``DurableJobQuota``). Reads for ``GET /v1/queue`` go straight through the read
model. Priority for claiming is ``ref_count`` desc (most-cited-unread first).
"""
from __future__ import annotations

import threading
from dataclasses import dataclass
from datetime import datetime, timedelta

from app.config import Settings
from app.frontmatter import serialise
from app.hub import HubClient
from app.naming import stamp_iso, utc_now
from app.read_model import ReadModel
from app.wiki import sanitize_id


def _parse_iso(ts: str) -> datetime | None:
    if not ts:
        return None
    try:
        return datetime.fromisoformat(ts.replace("Z", "+00:00"))
    except ValueError:
        return None


@dataclass
class ClaimedItem:
    id: str
    title: str
    type: str
    topic_hint: str
    ref_count: int


class AtomicQueue:
    def __init__(self, settings: Settings, hub: HubClient, read_model: ReadModel):
        self._settings = settings
        self._hub = hub
        self._rm = read_model
        self._prefix = settings.queue_prefix if settings.queue_prefix.endswith("/") else settings.queue_prefix + "/"
        self._folder = self._prefix.rstrip("/")
        self._lease = timedelta(seconds=settings.queue_claim_lease_s)
        self._lock = threading.Lock()

    # ── storage helpers ──────────────────────────────────────────────
    def _path(self, source_id: str) -> str:
        return f"{self._prefix}{sanitize_id(source_id)}.md"

    def _items(self) -> dict[str, dict]:
        """id -> {fm, body} for every well-formed queue file."""
        out: dict[str, dict] = {}
        for r in self._rm.records(self._folder):
            if r.parse_error:
                continue
            sid = str(r.frontmatter.get("id") or "")
            if sid:
                out[sid] = {"fm": dict(r.frontmatter), "body": r.body}
        return out

    def _write(self, source_id: str, fm: dict, body: str) -> None:
        content = serialise(fm, body)
        path = self._path(source_id)
        self._hub.write_text_central(path, content)
        self._rm.write_through(path, fm, body, len(content.encode("utf-8")))

    @staticmethod
    def _new_item(source_id: str, *, title="", type_="paper", topic_hint="", discovered_from=None) -> dict:
        return {
            "id": source_id,
            "status": "queued",
            "type": type_,
            "title": title,
            "topic_hint": topic_hint,
            "ref_count": 0,
            "priority": "discovered",
            "discovered_from": list(discovered_from or []),
            "claimed_by": "",
            "claim_expires": "",
            "processed_source": "",
        }

    def _reap_locked(self, items: dict[str, dict]) -> int:
        """Return expired claims to the frontier. Assumes the lock is held."""
        now = utc_now()
        reaped = 0
        for sid, it in items.items():
            fm = it["fm"]
            if fm.get("status") != "claimed":
                continue
            exp = _parse_iso(str(fm.get("claim_expires") or ""))
            if exp is not None and exp < now:
                fm["status"] = "queued"
                fm["claimed_by"] = ""
                fm["claim_expires"] = ""
                self._write(sid, fm, it["body"])
                reaped += 1
        return reaped

    # ── public API ───────────────────────────────────────────────────
    def reap(self) -> int:
        with self._lock:
            return self._reap_locked(self._items())

    def add(self, items: list[dict]) -> dict[str, int]:
        """Enqueue discovered sources. Dedup by id: a repeat bumps ref_count and
        unions discovered_from instead of duplicating. Returns counts."""
        added = updated = 0
        with self._lock:
            existing = self._items()
            for spec in items:
                sid = str(spec.get("id") or "").strip()
                if not sid:
                    continue
                origin = [o for o in (spec.get("discovered_from") or []) if o]
                if sid in existing:
                    fm = existing[sid]["fm"]
                    fm["ref_count"] = int(fm.get("ref_count") or 0) + 1
                    merged = list(dict.fromkeys([*(fm.get("discovered_from") or []), *origin]))
                    fm["discovered_from"] = merged
                    if not fm.get("title") and spec.get("title"):
                        fm["title"] = str(spec["title"])
                    self._write(sid, fm, existing[sid]["body"])
                    updated += 1
                else:
                    fm = self._new_item(
                        sid,
                        title=str(spec.get("title") or ""),
                        type_=str(spec.get("type") or "paper"),
                        topic_hint=str(spec.get("topic_hint") or ""),
                        discovered_from=origin,
                    )
                    existing[sid] = {"fm": fm, "body": "Discovered source.\n"}
                    self._write(sid, fm, existing[sid]["body"])
                    added += 1
        return {"added": added, "updated": updated}

    def claim(self, agent_id: str, topic: str | None = None) -> ClaimedItem | None:
        """Atomically lease the highest-priority unclaimed source to ``agent_id``."""
        with self._lock:
            items = self._items()
            self._reap_locked(items)
            candidates = [
                (sid, it["fm"]) for sid, it in items.items()
                if it["fm"].get("status") == "queued"
                and (not topic or str(it["fm"].get("topic_hint") or "") == topic)
            ]
            if not candidates:
                return None
            candidates.sort(key=lambda kv: (-int(kv[1].get("ref_count") or 0), kv[0]))
            sid, fm = candidates[0]
            fm["status"] = "claimed"
            fm["claimed_by"] = agent_id
            fm["claim_expires"] = stamp_iso(utc_now() + self._lease)
            self._write(sid, fm, items[sid]["body"])
            return ClaimedItem(
                id=sid,
                title=str(fm.get("title") or ""),
                type=str(fm.get("type") or "paper"),
                topic_hint=str(fm.get("topic_hint") or ""),
                ref_count=int(fm.get("ref_count") or 0),
            )

    def skip(self, agent_id: str, source_id: str, reason: str = "") -> bool:
        with self._lock:
            items = self._items()
            if source_id not in items:
                return False
            fm = items[source_id]["fm"]
            fm["status"] = "skipped"
            fm["skip_reason"] = reason
            fm["skipped_by"] = agent_id
            fm["claimed_by"] = ""
            fm["claim_expires"] = ""
            self._write(source_id, fm, items[source_id]["body"])
            return True

    def mark_processed(self, source_id: str, processed_source: str = "") -> None:
        """Flip a source to processed (called by the merge-bot on a merged
        sources/ PR). Creates the item if it wasn't on the frontier."""
        with self._lock:
            items = self._items()
            if source_id in items:
                fm = items[source_id]["fm"]
                body = items[source_id]["body"]
            else:
                fm = self._new_item(source_id)
                body = "Processed source.\n"
            fm["status"] = "processed"
            fm["processed_source"] = processed_source
            fm["claimed_by"] = ""
            fm["claim_expires"] = ""
            self._write(source_id, fm, body)