"""Auto-heartbeat — background loop that extends claim TTLs for this MCP session."""

from __future__ import annotations

import asyncio
import logging
from datetime import datetime, timedelta, timezone

import asyncpg
from redis.asyncio import Redis

from loom.graph import cache, store

logger = logging.getLogger(__name__)


class AutoHeartbeat:
    """Tracks claimed tasks and auto-extends their TTLs in a background loop.

    Lifecycle:
    1. loom_claim / loom_batch_claim → register(task_id, ttl)
    2. Background loop runs every `interval` seconds, extends claim_expires_at
    3. loom_done / loom_fail / loom_reset → unregister(task_id)
    4. Process exit → loop dies, claims expire naturally via TTL
    """

    def __init__(self, pool: asyncpg.Pool, redis: Redis, interval: int = 90) -> None:
        self._pool = pool
        self._redis = redis
        self._interval = interval
        self._claims: dict[str, int] = {}  # task_id -> ttl_seconds
        self._task: asyncio.Task | None = None

    def register(self, task_id: str, ttl_seconds: int) -> None:
        """Start tracking a claimed task. Starts the loop if not running."""
        self._claims[task_id] = ttl_seconds
        if self._task is None or self._task.done():
            self._task = asyncio.create_task(self._loop())

    def unregister(self, task_id: str) -> None:
        """Stop tracking a task. Cancels the loop if no tasks remain."""
        self._claims.pop(task_id, None)
        if not self._claims and self._task is not None and not self._task.done():
            self._task.cancel()
            self._task = None

    async def _loop(self) -> None:
        """Periodically extend claim TTLs for all tracked tasks."""
        try:
            while self._claims:
                await asyncio.sleep(self._interval)
                # Snapshot to avoid dict-changed-during-iteration
                snapshot = dict(self._claims)
                for task_id, ttl in snapshot.items():
                    if task_id not in self._claims:
                        continue  # unregistered while iterating
                    try:
                        expires_at = datetime.now(timezone.utc) + timedelta(seconds=ttl)
                        task = await store.set_claim_expiry(self._pool, task_id, expires_at)
                        await cache.sync_task(self._redis, task)
                        logger.debug("Auto-heartbeat extended %s to %s", task_id, expires_at)
                    except Exception:
                        logger.warning(
                            "Auto-heartbeat failed to extend %s, will retry next cycle",
                            task_id,
                            exc_info=True,
                        )
        except asyncio.CancelledError:
            return

    async def shutdown(self) -> None:
        """Cancel the background loop cleanly."""
        if self._task is not None and not self._task.done():
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        self._task = None
        self._claims.clear()
