"""The engine's work queue: what survives a crash.

Every external event — an MQTT message, a webhook, a cron tick, a connector
poll — becomes a work item before anything runs. The item is journaled first
and acknowledged only once the wave it started has quiesced, so an engine that
dies mid-cascade picks the work up again on the way back rather than losing it.

Two implementations: Redis Streams, which is what makes the above true, and an
in-memory one for tests and for running without Redis, where "durable" degrades
honestly to "not".
"""

from __future__ import annotations

import heapq
import json
import logging
import threading
import time
from abc import ABC, abstractmethod
from collections import deque
from dataclasses import dataclass, field
from typing import Any, cast

import redis

logger = logging.getLogger(__name__)

# The stream is a buffer, not an archive: anything this far behind is long
# superseded by a newer value of the same message.
STREAM_MAXLEN = 10_000
# An item redelivered this often is not going to succeed. Park it where someone
# can look rather than letting it loop forever.
MAX_DELIVERIES = 3
GROUP = "engine"


@dataclass
class WorkItem:
    """One unit of journaled work.

    :param kind: ``cascade`` replays a node's outputs and runs what is
        downstream; ``flush`` lets out what a node's rate limits held back;
        ``run`` is a whole batch run, and carries only its id.
    :param node: The node the item is about — the source for a cascade, the
        one whose rate limits are released for a flush.
    :param flow: The flow that node belongs to, so gating needs no lookup.
    :param outputs: What the source node emitted (cascade only).
    :param cause: Where the work came from, for logs and debugging.
    :param not_before: Epoch seconds before which the item must not run.
    :param guard: ``(key, value)`` the target node must still remember for
        this item to be worth running — how a rescheduled wait cancels the
        one it replaced.
    :param entry_id: Set by the queue on claim; stable across redeliveries,
        which is what makes it usable as an idempotency key.
    :param deliveries: How many times this item has been handed out.
    :param enqueued_at: When the item was made. Carried rather than read off
        the entry id, because only the Redis queue has timestamps in its ids.
    """

    kind: str
    node: str
    flow: str
    outputs: dict[str, Any] = field(default_factory=dict)
    cause: str = "system"
    not_before: float = 0.0
    guard_key: str = ""
    guard_value: str = ""
    entry_id: str = ""
    deliveries: int = 1
    enqueued_at: float = field(default_factory=time.time)
    #: The run this item is, on a ``run`` item. Everything else about a run is
    #: in its database row, so the journal only has to name it.
    run_id: str = ""

    def to_fields(self) -> dict[str, str]:
        return {
            "kind": self.kind,
            "node": self.node,
            "flow": self.flow,
            "outputs": json.dumps(self.outputs),
            "cause": self.cause,
            "not_before": str(self.not_before),
            "guard_key": self.guard_key,
            "guard_value": self.guard_value,
            "enqueued_at": str(self.enqueued_at),
            "run_id": self.run_id,
        }

    @classmethod
    def from_fields(
        cls, fields: dict[str, str], entry_id: str, deliveries: int = 1
    ) -> WorkItem:
        return cls(
            kind=fields.get("kind", "cascade"),
            node=fields.get("node", ""),
            flow=fields.get("flow", ""),
            outputs=json.loads(fields.get("outputs") or "{}"),
            cause=fields.get("cause", "system"),
            not_before=float(fields.get("not_before") or 0.0),
            guard_key=fields.get("guard_key", ""),
            guard_value=fields.get("guard_value", ""),
            entry_id=entry_id,
            deliveries=deliveries,
            enqueued_at=float(fields.get("enqueued_at") or time.time()),
            run_id=fields.get("run_id", ""),
        )


class WorkQueue(ABC):
    """What the execution service pulls from."""

    @abstractmethod
    def add(self, item: WorkItem) -> None:
        """Journal an item for immediate execution."""

    @abstractmethod
    def add_delayed(self, item: WorkItem, not_before: float) -> None:
        """Journal an item that must not run before ``not_before``."""

    @abstractmethod
    def claim(self, count: int, block_ms: int) -> list[WorkItem]:
        """Take up to ``count`` items, waiting up to ``block_ms`` for one."""

    @abstractmethod
    def ack(self, item: WorkItem) -> None:
        """Mark an item done, so it is never redelivered."""

    def touch(self, entry_ids: list[str]) -> None:
        """Say these items are still being worked on, not abandoned.

        A node with no timeout may run far longer than the reclaim window, and
        nothing else distinguishes that from an engine that died holding the
        item. Concrete rather than abstract: a queue with no redelivery has
        nothing to answer here.
        """

    @abstractmethod
    def reclaim_stale(self, min_idle_ms: int) -> list[WorkItem]:
        """Take back items claimed by a consumer that never acknowledged them."""

    @abstractmethod
    def move_due(self, now: float) -> int:
        """Promote delayed items whose time has come. Returns how many."""

    @abstractmethod
    def dead_letter(self, item: WorkItem, reason: str) -> None:
        """Set an item aside after it has failed too often."""

    @abstractmethod
    def park(self, flow: str, item: WorkItem) -> None:
        """Hold an item while its flow is paused."""

    @abstractmethod
    def unpark(self, flow: str) -> list[WorkItem]:
        """Return a paused flow's held items to the queue, oldest first."""

    @abstractmethod
    def unpark_one(self, flow: str) -> WorkItem | None:
        """Take the oldest held item, leaving the rest parked. That is a step."""

    @abstractmethod
    def clear_flow(self, flow: str) -> None:
        """Forget anything held for a flow that no longer exists."""

    @abstractmethod
    def stats(self) -> dict[str, Any]:
        """In-flight, delayed and parked counts plus the oldest pending age,
        for the health endpoint."""

    @abstractmethod
    def dead_letters(self, count: int = 50) -> list[dict[str, Any]]:
        """What was given up on, newest first."""

    @abstractmethod
    def mark_done(self, entry_id: str, node: str) -> None:
        """Record that a side effect already happened for this delivery."""

    @abstractmethod
    def was_done(self, entry_id: str, node: str) -> bool:
        """Did this exact item already run that side-effecting node?"""

    @abstractmethod
    def close(self) -> None:
        """Release whatever the queue holds."""


class MemoryWorkQueue(WorkQueue):
    """In-process queue. No durability: a crash loses whatever is in flight."""

    def __init__(self) -> None:
        self._items: deque[WorkItem] = deque()
        self._delayed: list[tuple[float, int, WorkItem]] = []
        self._parked: dict[str, list[WorkItem]] = {}
        self._done: set[tuple[str, str]] = set()
        # Claimed and not yet acknowledged, which is what Redis's `pending` is.
        self._in_flight = 0
        self._counter = 0
        self._seq = 0
        self._lock = threading.Lock()
        self._wake = threading.Condition(self._lock)

    def _next_id(self) -> str:
        self._seq += 1
        return f"mem-{self._seq}"

    def add(self, item: WorkItem) -> None:
        with self._wake:
            if not item.entry_id:
                item.entry_id = self._next_id()
            self._items.append(item)
            self._wake.notify()

    def add_delayed(self, item: WorkItem, not_before: float) -> None:
        with self._wake:
            if not item.entry_id:
                item.entry_id = self._next_id()
            item.not_before = not_before
            self._counter += 1
            heapq.heappush(self._delayed, (not_before, self._counter, item))
            self._wake.notify()

    def claim(self, count: int, block_ms: int) -> list[WorkItem]:
        deadline = time.monotonic() + block_ms / 1000.0
        with self._wake:
            while not self._items:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    return []
                self._wake.wait(remaining)
            claimed = [
                self._items.popleft() for _ in range(min(count, len(self._items)))
            ]
            self._in_flight += len(claimed)
            return claimed

    def ack(self, item: WorkItem) -> None:
        """Nothing to take off the queue — claiming did that. Only counted."""
        with self._lock:
            self._in_flight = max(0, self._in_flight - 1)

    def reclaim_stale(self, min_idle_ms: int) -> list[WorkItem]:
        """Nothing to reclaim: an item lost here died with the process."""
        return []

    def move_due(self, now: float) -> int:
        moved = 0
        with self._wake:
            while self._delayed and self._delayed[0][0] <= now:
                self._items.append(heapq.heappop(self._delayed)[2])
                moved += 1
            if moved:
                self._wake.notify()
        return moved

    def dead_letter(self, item: WorkItem, reason: str) -> None:
        logger.error("Dropping work item for '%s': %s", item.node, reason)

    def park(self, flow: str, item: WorkItem) -> None:
        with self._lock:
            self._parked.setdefault(flow, []).append(item)

    def unpark(self, flow: str) -> list[WorkItem]:
        with self._lock:
            return self._parked.pop(flow, [])

    def unpark_one(self, flow: str) -> WorkItem | None:
        with self._lock:
            held = self._parked.get(flow)
            return held.pop(0) if held else None

    def clear_flow(self, flow: str) -> None:
        with self._lock:
            self._parked.pop(flow, None)

    def stats(self) -> dict[str, Any]:
        with self._lock:
            return {
                "pending": self._in_flight,
                "delayed": len(self._delayed),
                "parked": sum(len(v) for v in self._parked.values()),
                "oldest_pending_s": 0.0,
                "durable": False,
            }

    def dead_letters(self, count: int = 50) -> list[dict[str, Any]]:
        """Nothing is kept: a dropped item here died with the process."""
        return []

    def mark_done(self, entry_id: str, node: str) -> None:
        with self._lock:
            self._done.add((entry_id, node))

    def was_done(self, entry_id: str, node: str) -> bool:
        with self._lock:
            return (entry_id, node) in self._done

    def close(self) -> None:
        """Nothing is held outside the process."""


class RedisWorkQueue(WorkQueue):
    """Redis Streams queue: the one that survives the process.

    Items live in a stream read through a consumer group, so an item handed to
    a consumer that dies before acknowledging it stays pending and is reclaimed
    by whoever comes next.
    """

    def __init__(
        self,
        host: str,
        port: int = 6379,
        # A root of its own: the state backend owns `pipeline:`, and sharing the
        # prefix meant `RedisState.clear()` could DEL the queue stream.
        namespace: str = "queue",
        consumer: str | None = None,
    ) -> None:
        self._redis = redis.Redis(host=host, port=port, decode_responses=True)
        self._ns = namespace
        self._consumer = consumer or f"engine-{int(time.time() * 1000) % 1_000_000}"
        self._stream = f"{namespace}:__queue__"
        self._delayed_key = f"{namespace}:__delayed__"
        self._dead_key = f"{namespace}:__dead__"
        self._ensure_group()

    def _ensure_group(self) -> None:
        try:
            self._redis.xgroup_create(self._stream, GROUP, id="0", mkstream=True)
        except redis.ResponseError as exc:
            if "BUSYGROUP" not in str(exc):
                raise

    def _parked_key(self, flow: str) -> str:
        return f"{self._ns}:__parked__:{flow}"

    def _done_key(self, entry_id: str, node: str) -> str:
        return f"{self._ns}:__done__:{entry_id}:{node}"

    def add(self, item: WorkItem) -> None:
        self._redis.xadd(
            self._stream,
            cast(Any, item.to_fields()),
            maxlen=STREAM_MAXLEN,
            approximate=True,
        )

    def add_delayed(self, item: WorkItem, not_before: float) -> None:
        item.not_before = not_before
        self._redis.zadd(self._delayed_key, {json.dumps(item.to_fields()): not_before})

    def claim(self, count: int, block_ms: int) -> list[WorkItem]:
        response = cast(
            list[Any],
            self._redis.xreadgroup(
                GROUP,
                self._consumer,
                {self._stream: ">"},
                count=count,
                block=block_ms,
            ),
        )
        items: list[WorkItem] = []
        for _stream, entries in response or []:
            for entry_id, fields in entries:
                items.append(WorkItem.from_fields(fields, entry_id))
        return items

    def ack(self, item: WorkItem) -> None:
        if item.entry_id:
            self._redis.xack(self._stream, GROUP, item.entry_id)

    def touch(self, entry_ids: list[str]) -> None:
        if not entry_ids:
            return
        # Claiming an entry we already hold resets how long it has been idle,
        # which is the only thing `reclaim_stale` reads. `justid` keeps the
        # delivery count where it is, so a long node neither redelivers nor
        # spends its way towards the dead-letter cap. An entry already
        # acknowledged is simply not there, and this is a no-op for it.
        self._redis.xclaim(
            self._stream,
            GROUP,
            self._consumer,
            min_idle_time=0,
            message_ids=entry_ids,
            justid=True,
        )

    def reclaim_stale(self, min_idle_ms: int) -> list[WorkItem]:
        """Take over entries a dead consumer never acknowledged."""
        _cursor, entries, _deleted = cast(
            tuple[Any, list[Any], Any],
            self._redis.xautoclaim(
                self._stream,
                GROUP,
                self._consumer,
                min_idle_time=min_idle_ms,
                count=32,
            ),
        )
        if not entries:
            return []

        # xautoclaim does not report delivery counts, so ask xpending for them.
        counts: dict[str, int] = {}
        for record in cast(
            list[dict[str, Any]],
            self._redis.xpending_range(self._stream, GROUP, min="-", max="+", count=64),
        ):
            counts[record["message_id"]] = record["times_delivered"]

        return [
            WorkItem.from_fields(fields, entry_id, counts.get(entry_id, 1))
            for entry_id, fields in entries
            if fields
        ]

    def move_due(self, now: float) -> int:
        due = cast(
            list[str],
            self._redis.zrangebyscore(self._delayed_key, "-inf", now, start=0, num=100),
        )
        moved = 0
        for raw in due:
            # Whoever removes it owns it: a second engine would get 0 here.
            if self._redis.zrem(self._delayed_key, raw):
                self._redis.xadd(
                    self._stream,
                    cast(Any, json.loads(raw)),
                    maxlen=STREAM_MAXLEN,
                    approximate=True,
                )
                moved += 1
        return moved

    def dead_letter(self, item: WorkItem, reason: str) -> None:
        fields = item.to_fields()
        fields["reason"] = reason
        self._redis.xadd(
            self._dead_key, cast(Any, fields), maxlen=1000, approximate=True
        )
        logger.error("Dead-lettered work item for '%s': %s", item.node, reason)

    def park(self, flow: str, item: WorkItem) -> None:
        self._redis.rpush(self._parked_key(flow), json.dumps(item.to_fields()))

    def unpark(self, flow: str) -> list[WorkItem]:
        key = self._parked_key(flow)
        raw = cast(list[str], self._redis.lrange(key, 0, -1))
        self._redis.delete(key)
        return [WorkItem.from_fields(json.loads(r), "") for r in raw]

    def unpark_one(self, flow: str) -> WorkItem | None:
        raw = cast("str | None", self._redis.lpop(self._parked_key(flow)))
        return WorkItem.from_fields(json.loads(raw), "") if raw else None

    def clear_flow(self, flow: str) -> None:
        self._redis.delete(self._parked_key(flow))

    def stats(self) -> dict[str, Any]:
        pending = cast(dict[str, Any], self._redis.xpending(self._stream, GROUP))
        oldest = 0.0
        count = pending.get("pending", 0) if isinstance(pending, dict) else 0
        if count:
            records = cast(
                list[dict[str, Any]],
                self._redis.xpending_range(
                    self._stream, GROUP, min="-", max="+", count=1
                ),
            )
            if records:
                oldest = records[0]["time_since_delivered"] / 1000.0
        parked = sum(
            cast(int, self._redis.llen(key))
            for key in self._redis.scan_iter(f"{self._ns}:__parked__:*")
        )
        return {
            "pending": count,
            "delayed": cast(int, self._redis.zcard(self._delayed_key)),
            "parked": parked,
            "oldest_pending_s": round(oldest, 1),
            "durable": True,
        }

    def dead_letters(self, count: int = 50) -> list[dict[str, Any]]:
        """What was set aside, newest first. The entry id carries the time."""
        entries = cast(
            list[tuple[str, dict[str, str]]],
            self._redis.xrevrange(self._dead_key, count=count),
        )
        return [
            {
                "id": entry_id,
                "ts": int(entry_id.split("-")[0]) / 1000.0,
                "flow": fields.get("flow", ""),
                "node": fields.get("node", ""),
                "cause": fields.get("cause", ""),
                "reason": fields.get("reason", ""),
            }
            for entry_id, fields in entries
        ]

    def mark_done(self, entry_id: str, node: str) -> None:
        # An hour outlives any redelivery; after that the marker is noise.
        self._redis.set(self._done_key(entry_id, node), "1", ex=3600)

    def was_done(self, entry_id: str, node: str) -> bool:
        return bool(self._redis.exists(self._done_key(entry_id, node)))

    def close(self) -> None:
        self._redis.close()