"""Redis-backed broker that routes invocations to active device sessions.

The device CLI only ever talks to the web process (poll / SSE / output
POST), but an agent may dispatch a command from *any* process — notably a
Celery worker during a scheduled run. An in-memory broker can't bridge
that gap: the worker and the web process don't share Python memory, so a
worker-side dispatch never reaches the web-side SSE session and the tool
times out. Routing through Redis makes every hop cross-process.

What lives in Redis (shared) vs. in the process (ephemeral):

* ``dev:cmd:{device_id}``   — list of queued command envelopes (JSON).
* ``dev:ticket:{device_id}``— poll-issued SSE upgrade ticket (string, TTL).
* ``dev:inv:{invocation_id}``— invocation metadata hash (status, result).
* ``dev:out:{invocation_id}``— stream of stdout/stderr/control chunks.
* ``SessionState`` is per-connection state for the one SSE handler that
  owns the live socket; it stays in that process's memory.

Lifecycle:
1. Agent (any process) calls ``dispatch_invocation`` → metadata hash + an
   RPUSH onto the device's command list.
2. The web SSE handler blocks on that list (``next_command``) and emits
   each envelope to the wire; an offline device leaves the envelope queued
   until its next ``/poll`` issues a ticket and upgrades to SSE.
3. The CLI POSTs ack + chunked output back; ``submit_output_chunk`` XADDs
   chunks to the invocation's output stream and updates the hash.
4. A ``control`` chunk closes the invocation; ``drain_output`` (in the
   dispatching process) reads the stream from the start and stops on it.
"""

from __future__ import annotations

import json
import logging
import threading
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, Iterator, Optional

import anyio

from docsgpt.cache import get_redis_instance
from docsgpt.core.settings import settings
from docsgpt.streaming.async_redis import get_async_redis_instance


logger = logging.getLogger(__name__)


# Kept for backwards compatibility with callers/tests importing the
# sentinel; the Redis path signals end-of-output with a ``control`` chunk.
INVOCATION_DONE = object()


def _cmd_key(device_id: str) -> str:
    return f"dev:cmd:{device_id}"


def _ticket_key(device_id: str) -> str:
    return f"dev:ticket:{device_id}"


# Deletes the device's ticket iff it still holds the presented one. Runs
# server-side so two requests racing with one ticket can't both redeem it.
_REDEEM_TICKET_LUA = """
if redis.call('GET', KEYS[1]) == ARGV[1] then
    return redis.call('DEL', KEYS[1])
end
return 0
"""


def _inv_key(invocation_id: str) -> str:
    return f"dev:inv:{invocation_id}"


def _out_key(invocation_id: str) -> str:
    return f"dev:out:{invocation_id}"


def _as_str(value: Any) -> str:
    if isinstance(value, (bytes, bytearray)):
        return value.decode("utf-8", "replace")
    return str(value)


@dataclass
class Invocation:
    """Snapshot of an invocation's cross-process state.

    Returned by ``dispatch_invocation`` (fresh) and ``get_invocation`` (read
    from Redis). Carries enough for ownership checks and audit without the
    caller touching Redis directly.
    """

    invocation_id: str
    device_id: str
    completed: bool = False
    exit_code: Optional[int] = None
    duration_ms: Optional[int] = None
    error: Optional[str] = None
    started_at: Optional[float] = None
    finished_at: Optional[float] = None
    stdout_bytes: int = 0
    stderr_bytes: int = 0


@dataclass
class SessionState:
    """Per-connection state for the SSE handler that owns the live socket."""

    session_id: str
    device_id: str
    user_id: str
    last_event_id: int = 0
    last_activity_at: float = field(default_factory=time.time)
    closed: threading.Event = field(default_factory=threading.Event)


class DeviceBroker:
    """Cross-process device registry backed by Redis.

    Redis is the source of truth for queued commands, output, and tickets,
    so dispatch and drain work regardless of which process they run in. A
    small in-memory map of live SSE sessions stays local to the web process
    that holds each socket (sessions are inherently per-connection).
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._sessions_by_id: Dict[str, SessionState] = {}
        self._sessions_by_device: Dict[str, SessionState] = {}

    # ------------------------------------------------------------------
    # Session lifecycle (web process / CLI side)
    # ------------------------------------------------------------------
    def register_session(self, device_id: str, user_id: str) -> SessionState:
        """Open the SSE session for ``device_id``, adopting its poll ticket.

        The poll-issued ticket becomes the ``session_id`` so the URL the CLI
        opens matches the live session. A previous local session for the
        same device is closed and replaced (the CLI reconnected).
        """
        redis = get_redis_instance()
        issued = None
        if redis is not None:
            try:
                issued = redis.get(_ticket_key(device_id))
                if issued is not None:
                    redis.delete(_ticket_key(device_id))
            except Exception:
                logger.exception("ticket lookup failed for %s", device_id)
        session_id = _as_str(issued) if issued else f"st_{uuid.uuid4().hex}"
        return self._adopt_session(session_id, device_id, user_id)

    def redeem_ticket(self, device_id: str, user_id: str, ticket: str) -> Optional[SessionState]:
        """Consume the device's poll-issued ``ticket`` and open its session under it.

        Returns ``None``, and opens nothing, unless ``ticket`` is still the
        device's unexpired ticket. The compare and delete are one Redis step,
        so a replay racing the first redeem can't replace its session.
        """
        if not ticket:
            return None
        redis = get_redis_instance()
        if redis is None:
            return None
        try:
            redeemed = redis.eval(_REDEEM_TICKET_LUA, 1, _ticket_key(device_id), ticket)
        except Exception:
            logger.exception("ticket redeem failed for %s", device_id)
            return None
        if not redeemed:
            return None
        return self._adopt_session(ticket, device_id, user_id)

    def _adopt_session(self, session_id: str, device_id: str, user_id: str) -> SessionState:
        """Make a new session the device's live one, closing the one it replaces."""
        sess = SessionState(
            session_id=session_id, device_id=device_id, user_id=user_id
        )
        with self._lock:
            prior = self._sessions_by_device.get(device_id)
            if prior is not None:
                prior.closed.set()
                self._sessions_by_id.pop(prior.session_id, None)
            self._sessions_by_device[device_id] = sess
            self._sessions_by_id[session_id] = sess
        return sess

    def close_session(self, session_id: str, *, reason: str = "closed") -> None:
        """Close the local session by id.

        Queued-but-undelivered commands stay on the device's Redis list and
        are picked up by the next session; an in-flight command whose socket
        drops falls back to the tool's own drain deadline.
        """
        with self._lock:
            sess = self._sessions_by_id.pop(session_id, None)
            if sess is None:
                return
            sess.closed.set()
            if self._sessions_by_device.get(sess.device_id) is sess:
                self._sessions_by_device.pop(sess.device_id, None)
        logger.debug("device session closed: %s (%s)", session_id, reason)

    def get_session(self, session_id: str) -> Optional[SessionState]:
        with self._lock:
            return self._sessions_by_id.get(session_id)

    def next_command(
        self, session: SessionState, timeout: float = 1.0
    ) -> Optional[Dict[str, Any]]:
        """Block up to ``timeout`` for the next queued command envelope.

        Returns the decoded envelope, or ``None`` on timeout so the SSE
        handler can emit a keepalive and re-check the session lifecycle.
        """
        redis = get_redis_instance()
        if redis is None:
            time.sleep(timeout)
            return None
        try:
            popped = redis.blpop(_cmd_key(session.device_id), timeout=timeout)
        except Exception:
            logger.exception("blpop failed for %s", session.device_id)
            time.sleep(timeout)
            return None
        if not popped:
            return None
        envelope = self._decode_envelope(popped[1])
        if envelope is None:
            return None
        # Drop an envelope whose invocation was already reaped (timed out /
        # cleaned up) after it was queued, so a command the user already saw
        # fail can't still run on the device. Best-effort — it narrows but does
        # not fully close the BLPOP-vs-cleanup window (see cleanup_invocation).
        inv_id = envelope.get("invocation_id")
        if inv_id and self.get_invocation(inv_id) is None:
            logger.debug("dropping reaped invocation %s", inv_id)
            return None
        return envelope

    async def next_command_async(
        self, session: SessionState, timeout: float = 1.0
    ) -> Optional[Dict[str, Any]]:
        """Await up to ``timeout`` for the next queued command envelope.

        Event-loop twin of :meth:`next_command` for the native-async SSE
        stream: the same envelope checks over the async Redis client, so an
        idle device session holds no thread.
        """
        redis = await get_async_redis_instance()
        if redis is None:
            await anyio.sleep(timeout)
            return None
        try:
            popped = await redis.blpop(_cmd_key(session.device_id), timeout=timeout)
        except Exception:
            logger.exception("async blpop failed for %s", session.device_id)
            await anyio.sleep(timeout)
            return None
        if not popped:
            return None
        envelope = self._decode_envelope(popped[1])
        if envelope is None:
            return None
        # Same reaped-invocation guard as ``next_command``; a failed lookup
        # drops the envelope too, as ``get_invocation`` returning None does.
        inv_id = envelope.get("invocation_id")
        if inv_id:
            try:
                reaped = not await redis.exists(_inv_key(inv_id))
            except Exception:
                logger.exception("invocation lookup failed for %s", inv_id)
                reaped = True
            if reaped:
                logger.debug("dropping reaped invocation %s", inv_id)
                return None
        return envelope

    @staticmethod
    def _decode_envelope(raw: Any) -> Optional[Dict[str, Any]]:
        """Parse a queued command envelope; ``None`` for anything malformed."""
        try:
            envelope = json.loads(_as_str(raw))
        except (TypeError, ValueError):
            logger.warning("dropping malformed command envelope")
            return None
        return envelope if isinstance(envelope, dict) else None

    # ------------------------------------------------------------------
    # Polling / tickets
    # ------------------------------------------------------------------
    def claim_ticket(self, device_id: str, ttl_seconds: float) -> Optional[str]:
        """Return an SSE upgrade ticket iff the device has queued work.

        Reuses an unexpired ticket so repeated polls don't churn it; the
        ticket's Redis TTL doubles as the advertised ``expires_in`` window.
        """
        redis = get_redis_instance()
        if redis is None:
            return None
        try:
            if redis.llen(_cmd_key(device_id)) <= 0:
                return None
            existing = redis.get(_ticket_key(device_id))
            if existing:
                return _as_str(existing)
            ticket = f"st_{uuid.uuid4().hex}"
            redis.set(_ticket_key(device_id), ticket, ex=int(ttl_seconds))
            return ticket
        except Exception:
            logger.exception("claim_ticket failed for %s", device_id)
            return None

    def validate_ticket(self, device_id: str, session_id: str) -> bool:
        """True iff ``session_id`` is the unexpired ticket issued to the device.

        An absent, mismatched, or expired ticket is rejected; expiry is
        enforced by Redis's TTL (a ``GET`` of an expired key returns nil).
        """
        if not session_id:
            return False
        redis = get_redis_instance()
        if redis is None:
            return False
        try:
            issued = redis.get(_ticket_key(device_id))
        except Exception:
            logger.exception("validate_ticket failed for %s", device_id)
            return False
        return issued is not None and _as_str(issued) == session_id

    # ------------------------------------------------------------------
    # Dispatch (server-issued, any process)
    # ------------------------------------------------------------------
    def dispatch_invocation(
        self,
        device_id: str,
        user_id: str,
        envelope: Dict[str, Any],
    ) -> Invocation:
        """Queue an invocation for ``device_id`` and record its metadata.

        Writes the metadata hash, then RPUSHes the envelope onto the
        device's command list. A live SSE session draining the list picks it
        up immediately; otherwise it waits for the next poll-issued ticket.
        """
        invocation_id = envelope["invocation_id"]
        inv = Invocation(invocation_id=invocation_id, device_id=device_id)
        redis = get_redis_instance()
        if redis is None:
            inv.error = "device broker unavailable"
            inv.completed = True
            return inv
        envelope_json = json.dumps(envelope)
        try:
            redis.hset(
                _inv_key(invocation_id),
                mapping={
                    "device_id": device_id,
                    "user_id": user_id,
                    "envelope": envelope_json,
                    "completed": "0",
                    "stdout_bytes": "0",
                    "stderr_bytes": "0",
                },
            )
            redis.expire(_inv_key(invocation_id), self._inv_ttl())
            redis.rpush(_cmd_key(device_id), envelope_json)
            redis.expire(_cmd_key(device_id), self._cmd_ttl())
        except Exception:
            logger.exception("dispatch_invocation failed for %s", invocation_id)
            # Don't strand the metadata hash (it holds the plaintext command)
            # if the queue write failed partway through.
            try:
                redis.delete(_inv_key(invocation_id))
            except Exception:
                logger.debug(
                    "cleanup after failed dispatch failed for %s", invocation_id
                )
            inv.error = "device broker dispatch failed"
            inv.completed = True
        return inv

    def get_invocation(self, invocation_id: str) -> Optional[Invocation]:
        """Read an invocation's current metadata snapshot from Redis."""
        redis = get_redis_instance()
        if redis is None:
            return None
        try:
            raw = redis.hgetall(_inv_key(invocation_id))
        except Exception:
            logger.exception("get_invocation failed for %s", invocation_id)
            return None
        if not raw:
            return None
        h = {_as_str(k): _as_str(v) for k, v in raw.items()}
        return Invocation(
            invocation_id=invocation_id,
            device_id=h.get("device_id", ""),
            completed=h.get("completed") == "1",
            exit_code=_to_int(h.get("exit_code")),
            duration_ms=_to_int(h.get("duration_ms")),
            error=h.get("error") or None,
            started_at=_to_float(h.get("started_at")),
            finished_at=_to_float(h.get("finished_at")),
            stdout_bytes=_to_int(h.get("stdout_bytes")) or 0,
            stderr_bytes=_to_int(h.get("stderr_bytes")) or 0,
        )

    # ------------------------------------------------------------------
    # Output streaming (web process / CLI side)
    # ------------------------------------------------------------------
    def submit_output_chunk(
        self, invocation_id: str, chunk: Dict[str, Any]
    ) -> bool:
        """Forward one CLI output chunk to the dispatching process's drain.

        Updates the metadata hash (byte counts; result fields on the closing
        ``control`` chunk) and XADDs the chunk to the invocation's output
        stream. Returns ``False`` for an unknown invocation.
        """
        redis = get_redis_instance()
        if redis is None:
            return False
        key = _inv_key(invocation_id)
        try:
            device_id = redis.hget(key, "device_id")
        except Exception:
            logger.exception("submit_output_chunk read failed for %s", invocation_id)
            return False
        if device_id is None:
            return False
        device_id = _as_str(device_id)
        now = time.time()
        stream = chunk.get("stream")
        try:
            # Append to the output stream BEFORE flipping completion state, so a
            # reader that observes completed=1 is guaranteed the chunk (and every
            # earlier one) is already on the stream — otherwise drain could see
            # completion and stop before the control chunk lands.
            redis.xadd(
                _out_key(invocation_id),
                {"c": json.dumps(chunk)},
                maxlen=self._out_maxlen(),
                approximate=True,
            )
            redis.expire(_out_key(invocation_id), self._inv_ttl())
            if stream in ("stdout", "stderr"):
                text = chunk.get("chunk")
                if isinstance(text, str):
                    field = "stdout_bytes" if stream == "stdout" else "stderr_bytes"
                    redis.hincrby(key, field, len(text.encode("utf-8")))
            elif stream == "control":
                redis.hsetnx(key, "started_at", repr(now))
                mapping = {"completed": "1", "finished_at": repr(now)}
                if chunk.get("exit_code") is not None:
                    mapping["exit_code"] = str(int(chunk["exit_code"]))
                if chunk.get("duration_ms") is not None:
                    mapping["duration_ms"] = str(int(chunk["duration_ms"]))
                if chunk.get("error"):
                    mapping["error"] = _as_str(chunk["error"])
                redis.hset(key, mapping=mapping)
            redis.expire(key, self._inv_ttl())
        except Exception:
            logger.exception("submit_output_chunk failed for %s", invocation_id)
            return False
        # Keep a co-located SSE session alive while output flows (single-worker
        # web tier); a cross-worker session relies on its own keepalive.
        with self._lock:
            sess = self._sessions_by_device.get(device_id)
        if sess is not None:
            sess.last_activity_at = now
        return True

    def submit_ack(
        self,
        invocation_id: str,
        decision: str,
        reason: Optional[str] = None,
    ) -> bool:
        """Record the CLI's accept/deny decision for an invocation."""
        redis = get_redis_instance()
        if redis is None:
            return False
        key = _inv_key(invocation_id)
        try:
            if not redis.exists(key):
                return False
            now = time.time()
            mapping = {"decision": decision}
            if reason:
                mapping["decision_reason"] = reason
            redis.hset(key, mapping=mapping)
            redis.hsetnx(key, "started_at", repr(now))
            redis.expire(key, self._inv_ttl())
            if decision == "denied":
                # XADD the synthetic control chunk BEFORE marking completed, so a
                # racing drain that observes completed=1 always finds the chunk
                # and reports "denied" rather than a false timeout.
                redis.xadd(
                    _out_key(invocation_id),
                    {"c": json.dumps(
                        {"stream": "control", "exit_code": None, "error": "denied"}
                    )},
                    maxlen=self._out_maxlen(),
                    approximate=True,
                )
                redis.expire(_out_key(invocation_id), self._inv_ttl())
                redis.hset(
                    key,
                    mapping={
                        "completed": "1",
                        "error": "denied",
                        "finished_at": repr(now),
                    },
                )
        except Exception:
            logger.exception("submit_ack failed for %s", invocation_id)
            return False
        return True

    def drain_output(
        self,
        invocation_id: str,
        timeout: float = 0.5,
        deadline: Optional[float] = None,
    ) -> Iterator[Dict[str, Any]]:
        """Yield queued output chunks for an invocation; stop on ``control``.

        Reads the invocation's output stream from the start (so chunks
        XADDed before draining began are never missed) and blocks up to
        ``timeout`` per read. ``deadline`` is an absolute ``time.time()``;
        once it passes with no closing ``control`` chunk the generator
        returns so a device that never responds can't loop forever.
        """
        redis = get_redis_instance()
        if redis is None:
            yield {
                "stream": "control",
                "exit_code": None,
                "error": "device broker unavailable",
            }
            return
        last_id = "0-0"
        block_ms = max(1, int(timeout * 1000))
        out_key = _out_key(invocation_id)
        # ``resp`` is parsed as the RESP2 list shape ``[[key, [(id, fields)]]]``;
        # the client must stay protocol=2 (redis-py default — see cache.py).
        tail_flush = False
        while True:
            try:
                # On the final sweep, read non-blocking (block=None) so any
                # entries that landed after the prior empty read are drained
                # before returning — closing the completion-vs-control race.
                resp = redis.xread(
                    {out_key: last_id},
                    count=200,
                    block=None if tail_flush else block_ms,
                )
            except Exception:
                logger.exception("xread failed for %s", invocation_id)
                return
            if not resp:
                if tail_flush:
                    return
                done = self._is_completed(redis, invocation_id)
                timed_out = deadline is not None and time.time() >= deadline
                if done or timed_out:
                    # One final non-blocking sweep from last_id, then stop.
                    tail_flush = True
                continue
            for _stream_key, entries in resp:
                for entry_id, fields in entries:
                    last_id = _as_str(entry_id)
                    raw = fields.get(b"c")
                    if raw is None:
                        raw = fields.get("c")
                    if raw is None:
                        continue
                    try:
                        chunk = json.loads(_as_str(raw))
                    except (TypeError, ValueError):
                        continue
                    if not isinstance(chunk, dict):
                        continue
                    if chunk.get("stream") in ("stdout", "stderr", "control"):
                        yield chunk
                    if chunk.get("stream") == "control":
                        return

    def cleanup_invocation(self, invocation_id: str) -> None:
        """Drop an invocation's Redis state, including any undelivered command.

        Deletes the metadata hash first so a concurrent ``next_command`` that
        just BLPOPped this envelope re-checks ``get_invocation``, sees it gone,
        and drops it instead of delivering a command the user already saw fail;
        the ``LREM`` then clears it for the still-queued (offline-device) case.
        Best-effort: a delivery that wins a tight race with the delete can still
        reach the device once, and that run's output is discarded.
        """
        redis = get_redis_instance()
        if redis is None:
            return
        try:
            raw = redis.hgetall(_inv_key(invocation_id))
            device_id = envelope_json = None
            if raw:
                h = {_as_str(k): _as_str(v) for k, v in raw.items()}
                device_id = h.get("device_id")
                envelope_json = h.get("envelope")
            redis.delete(_inv_key(invocation_id))
            redis.delete(_out_key(invocation_id))
            if device_id and envelope_json:
                redis.lrem(_cmd_key(device_id), 0, envelope_json)
        except Exception:
            logger.exception("cleanup_invocation failed for %s", invocation_id)

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------
    @staticmethod
    def _is_completed(redis, invocation_id: str) -> bool:
        try:
            return _as_str(redis.hget(_inv_key(invocation_id), "completed")) == "1"
        except Exception:
            return False

    @staticmethod
    def _inv_ttl() -> int:
        return int(settings.REMOTE_DEVICE_INVOCATION_TTL_SECONDS)

    @staticmethod
    def _cmd_ttl() -> int:
        return int(settings.REMOTE_DEVICE_CMD_QUEUE_TTL_SECONDS)

    @staticmethod
    def _out_maxlen() -> int:
        return int(settings.REMOTE_DEVICE_OUTPUT_STREAM_MAXLEN)


def _to_int(value: Optional[str]) -> Optional[int]:
    if value is None or value == "":
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _to_float(value: Optional[str]) -> Optional[float]:
    if value is None or value == "":
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


_broker_instance: Optional[DeviceBroker] = None
_broker_lock = threading.Lock()


def get_broker() -> DeviceBroker:
    """Return the process-wide ``DeviceBroker`` instance."""
    global _broker_instance
    if _broker_instance is None:
        with _broker_lock:
            if _broker_instance is None:
                _broker_instance = DeviceBroker()
    return _broker_instance