"""
Live View event bus: thread-safe pub/sub for streaming agent progress to SSE clients.

LiveViewEventBus holds per-thread_id queues. The agent/main thread calls emit() (sync)
and the FastAPI SSE endpoint consumes via subscribe() (async generator).

Events are also stored in an in-memory history per thread_id so that switching
conversations or re-opening the Live View pane replays past activity.
History is persisted to disk on quit and restored on launch.

Preview images are stored as individually encrypted files in a separate directory
(~/.mugi/live_view_images/) keyed by content hash. The main history blob stores
only a ``preview_image_hash`` reference, keeping it small. Images are served on
demand via the /live_view/image/<hash> HTTP endpoint.
"""
from __future__ import annotations

import asyncio
import hashlib
import json
import os
import queue
import threading
import time
from datetime import datetime, timezone


import logging as _logging
_log = _logging.getLogger(__name__)

_SENTINEL = object()

# TTL (seconds) after a "done" event before the queue is removed from memory
_QUEUE_TTL_SECONDS = 120

# Max events stored per thread_id to bound memory usage
_MAX_HISTORY_PER_THREAD = 500

# Known event types emitted by the application. Plugins may use custom types;
# unrecognized types log a debug warning but are NOT blocked.
KNOWN_EVENT_TYPES = frozenset({
    "status", "progress", "preview", "done", "error",
})

# Cap the number of concurrent cleanup timer threads to prevent unbounded growth
# in high-throughput scenarios (many agent threads finishing simultaneously).
_CLEANUP_SEMAPHORE = threading.Semaphore(32)

# Default persistence directory (same as MUGI_DIR)
_DEFAULT_PERSIST_DIR = os.path.expanduser("~/.mugi")
_HISTORY_FILENAME = "live_view_history.enc"
_IMAGES_DIR = os.path.join(_DEFAULT_PERSIST_DIR, "live_view_images")


class LiveViewEventBus:
    """Thread-safe event bus keyed by thread_id.

    - emit(thread_id, event_dict) can be called from any thread.
    - subscribe(thread_id) returns an async generator for use in a StreamingResponse.
    - get_history(thread_id) returns stored events for replay.
    """

    def __init__(self):
        self._lock = threading.Lock()
        # thread_id -> queue.Queue of event dicts (or _SENTINEL to close)
        self._queues: dict[str, queue.Queue] = {}
        # thread_id -> set of subscriber queues (for fan-out when multiple SSE clients connect)
        self._subscribers: dict[str, set[queue.Queue]] = {}
        # thread_id -> list of past event dicts (capped per thread)
        self._history: dict[str, list[dict]] = {}
        # Set when encrypted on-disk history cannot be decrypted (e.g. key rotation).
        self._persist_decrypt_failed: bool = False
        self._persist_load_error: str | None = None

    def _get_or_create_subs(self, thread_id: str) -> set[queue.Queue]:
        with self._lock:
            if thread_id not in self._subscribers:
                self._subscribers[thread_id] = set()
            return self._subscribers[thread_id]

    def emit(self, thread_id: str, event: dict) -> None:
        """Push an event to all subscribers for this thread_id. Safe to call from any thread."""
        if "timestamp" not in event:
            event["timestamp"] = datetime.now(timezone.utc).isoformat()

        etype = event.get("type")
        if etype not in KNOWN_EVENT_TYPES:
            _log.debug("emit: unrecognized event type %r for thread %r (plugin custom type?)", etype, thread_id)

        img_b64 = event.get("preview_image_base64")
        if img_b64 and isinstance(img_b64, str) and len(img_b64) > 100:
            img_hash = self._save_image(img_b64)
            if img_hash:
                event["preview_image_hash"] = img_hash
                event.pop("preview_image_base64", None)
            else:
                _log.debug(
                    "emit: preview_image_base64 present but image store failed for thread %r",
                    thread_id,
                )

        with self._lock:
            # Store in history
            hist = self._history.get(thread_id)
            if hist is None:
                hist = []
                self._history[thread_id] = hist
            hist.append(event)
            if len(hist) > _MAX_HISTORY_PER_THREAD:
                self._history[thread_id] = hist[-_MAX_HISTORY_PER_THREAD:]

            # Fan out to live subscribers
            subs = self._subscribers.get(thread_id)
            if subs:
                for q in list(subs):
                    try:
                        q.put_nowait(event)
                    except queue.Full:
                        _log.debug(
                            "emit: subscriber queue full for thread %r (event retained in history)",
                            thread_id,
                        )

        if event.get("type") == "done":
            self._schedule_cleanup(thread_id)

    def get_history(self, thread_id: str) -> list[dict]:
        """Return a copy of stored events for this thread_id (may be empty)."""
        with self._lock:
            return list(self._history.get(thread_id) or [])

    def get_persist_health(self) -> dict:
        """Health for encrypted on-disk history (DEF-097: surface key-rotation decrypt failure)."""
        with self._lock:
            return {
                "persist_decrypt_failed": bool(self._persist_decrypt_failed),
                "persist_load_error": self._persist_load_error,
            }

    def clear_history(self, thread_id: str) -> None:
        """Delete stored events for this thread_id and clean up orphaned images."""
        with self._lock:
            self._history.pop(thread_id, None)
        self._cleanup_orphan_images()

    def _schedule_cleanup(self, thread_id: str) -> None:
        """Remove subscriber set after a TTL so memory doesn't grow forever.

        Uses a semaphore to cap the number of concurrent sleeping cleanup threads.
        """
        def _cleanup():
            with _CLEANUP_SEMAPHORE:
                time.sleep(_QUEUE_TTL_SECONDS)
                with self._lock:
                    subs = self._subscribers.get(thread_id)
                    if subs is not None and len(subs) == 0:
                        del self._subscribers[thread_id]
        t = threading.Thread(target=_cleanup, daemon=True, name=f"lv-cleanup-{thread_id[:8]}")
        t.start()

    async def subscribe(self, thread_id: str):
        """Async generator that yields event dicts for this thread_id.

        Uses a short blocking timeout so asyncio cancellation and wait_for work properly.
        Exits when a _SENTINEL is received or the "done" event is yielded.
        """
        q: queue.Queue = queue.Queue(maxsize=512)
        subs = self._get_or_create_subs(thread_id)
        with self._lock:
            subs.add(q)

        def _get_with_timeout():
            try:
                return q.get(timeout=1.0)
            except queue.Empty:
                return None

        try:
            loop = asyncio.get_running_loop()
            while True:
                event = await loop.run_in_executor(None, _get_with_timeout)
                if event is None:
                    continue
                if event is _SENTINEL:
                    break
                yield event
                if event.get("type") == "done":
                    break
        finally:
            with self._lock:
                subs_set = self._subscribers.get(thread_id)
                if subs_set is not None:
                    subs_set.discard(q)

    def close(self, thread_id: str) -> None:
        """Signal all subscribers for this thread_id to stop (send sentinel)."""
        with self._lock:
            subs = self._subscribers.get(thread_id)
            if not subs:
                return
            for q in list(subs):
                try:
                    q.put_nowait(_SENTINEL)
                except queue.Full:
                    _log.debug("close: subscriber queue full for thread %r", thread_id)

    # ── Encrypted image store ────────────────────────────────────────────────

    def _save_image(self, b64: str) -> str | None:
        """Encrypt and write a base64 image to its own file. Returns the content hash, or None."""
        try:
            img_hash = hashlib.sha256(b64.encode("ascii", errors="replace")).hexdigest()[:32]
            path = os.path.join(_IMAGES_DIR, f"{img_hash}.enc")
            if os.path.isfile(path):
                return img_hash
            os.makedirs(_IMAGES_DIR, exist_ok=True)
            from chat_storage import _encrypt
            encrypted = _encrypt(b64)
            tmp = path + ".tmp"
            with open(tmp, "wb") as f:
                f.write(encrypted)
            os.replace(tmp, path)
            return img_hash
        except Exception:
            _log.debug("_save_image failed", exc_info=True)
            return None

    def load_image(self, img_hash: str) -> str | None:
        """Decrypt and return a base64 image string by its content hash, or None."""
        safe = "".join(c for c in (img_hash or "") if c in "0123456789abcdef")
        if len(safe) < 8:
            return None
        path = os.path.join(_IMAGES_DIR, f"{safe}.enc")
        if not os.path.isfile(path):
            return None
        try:
            from chat_storage import _decrypt
            with open(path, "rb") as f:
                return _decrypt(f.read())
        except Exception:
            return None

    def _cleanup_orphan_images(self) -> None:
        """Remove encrypted image files not referenced by any thread's history."""
        try:
            if not os.path.isdir(_IMAGES_DIR):
                return
            referenced: set[str] = set()
            with self._lock:
                for events in self._history.values():
                    for evt in events:
                        h = evt.get("preview_image_hash")
                        if h:
                            referenced.add(h)
            for fname in os.listdir(_IMAGES_DIR):
                if not fname.endswith(".enc"):
                    continue
                if fname[:-4] not in referenced:
                    try:
                        os.remove(os.path.join(_IMAGES_DIR, fname))
                    except OSError:
                        pass
        except Exception:
            pass

    def delete_image(self, img_hash: str) -> bool:
        """Remove an image from all history events and delete the encrypted file."""
        safe = "".join(c for c in (img_hash or "") if c in "0123456789abcdef")
        if len(safe) < 8:
            return False
        with self._lock:
            for events in self._history.values():
                for evt in events:
                    if evt.get("preview_image_hash") == safe:
                        evt.pop("preview_image_hash", None)
                        evt.pop("preview_image_base64", None)
        path = os.path.join(_IMAGES_DIR, f"{safe}.enc")
        try:
            os.remove(path)
        except OSError:
            pass
        return True

    # ── Encrypted disk persistence ──────────────────────────────────────────

    def save_to_disk(self, directory: str | None = None) -> None:
        """Persist history as an encrypted blob. Inline base64 image data is stripped;
        only the ``preview_image_hash`` reference is kept (images live in _IMAGES_DIR)."""
        dest_dir = directory or _DEFAULT_PERSIST_DIR
        os.makedirs(dest_dir, exist_ok=True)
        path = os.path.join(dest_dir, _HISTORY_FILENAME)
        with self._lock:
            serializable: dict[str, list[dict]] = {}
            for tid, events in self._history.items():
                stripped = []
                for evt in events:
                    e = {k: v for k, v in evt.items() if k != "preview_image_base64"}
                    stripped.append(e)
                serializable[tid] = stripped
        try:
            from chat_storage import _encrypt
            plaintext = json.dumps(serializable, separators=(",", ":"), default=str)
            encrypted = _encrypt(plaintext)
            tmp = path + ".tmp"
            with open(tmp, "wb") as f:
                f.write(encrypted)
            os.replace(tmp, path)
        except Exception as exc:
            print(f"[WARN] Failed to save Live View history: {exc}", flush=True)
        self._cleanup_orphan_images()

    def load_from_disk(self, directory: str | None = None) -> None:
        """Restore history from the encrypted file written by ``save_to_disk``."""
        src_dir = directory or _DEFAULT_PERSIST_DIR
        # Remove any old plaintext history left from before encryption was added
        _old_plain = os.path.join(src_dir, "live_view_history.json")
        if os.path.isfile(_old_plain):
            try:
                os.remove(_old_plain)
            except OSError:
                pass
        path = os.path.join(src_dir, _HISTORY_FILENAME)
        if not os.path.isfile(path):
            return
        try:
            from chat_storage import _decrypt
            with open(path, "rb") as f:
                encrypted = f.read()
            plaintext = _decrypt(encrypted)
            if not plaintext:
                return
            data = json.loads(plaintext)
            if not isinstance(data, dict):
                return
            with self._lock:
                for tid, events in data.items():
                    if isinstance(events, list):
                        self._history[tid] = events[-_MAX_HISTORY_PER_THREAD:]
                self._persist_decrypt_failed = False
                self._persist_load_error = None
        except Exception as exc:
            err = str(exc)
            with self._lock:
                self._persist_decrypt_failed = True
                self._persist_load_error = err[:240]
            _log.warning("Failed to load Live View history (decrypt/key mismatch?): %s", err)
