"""Encrypted shared context for AI agent coordination.

Provides a key-value store backed by MLS-encrypted channel messages.
Agents read and write shared state; late joiners get a snapshot.
LWW (last-writer-wins) conflict resolution uses Hybrid Logical Clocks.

Usage::

    from skytale_sdk import SkytaleChannelManager
    from skytale_sdk.context import SharedContext, ContextType

    mgr = SkytaleChannelManager(identity=b"agent", api_key="sk_live_...")
    ctx = SharedContext(mgr, "acme/research/results")
    ctx.set("task_status", {"phase": "running"})
    current = ctx.get("task_status")
"""

from __future__ import annotations

import json
import logging
import threading
import time
from typing import Any, Callable, Dict, List, Optional, Union

from skytale_sdk.context_types import (
    AccessScope,
    ContextEntry,
    ContextType,
    HLC,
)
from skytale_sdk.envelope import Envelope, Protocol

logger = logging.getLogger(__name__)


class SharedContext:
    """Encrypted shared context backed by an MLS channel.

    All context operations are sent as ``STATE`` protocol envelopes through
    the channel's existing MLS encryption. The relay routes them as opaque
    ciphertext — it never sees keys or values.

    Args:
        manager: A :class:`~skytale_sdk.channels.SkytaleChannelManager`
            instance with the channel already created or joined.
        channel: Channel name (must already exist in the manager).
        agent_id: Optional agent DID for access-scope enforcement.

    Example::

        ctx = SharedContext(mgr, "acme/research/results")
        ctx.set("status", {"phase": "running"})
        print(ctx.get("status"))  # {"phase": "running"}
    """

    def __init__(
        self,
        manager: Any,
        channel: str,
        agent_id: Optional[str] = None,
    ):
        self._manager = manager
        self._channel = channel
        self._agent_id = agent_id or getattr(manager, "_did_uri", None) or ""
        self._store: Dict[str, ContextEntry] = {}
        self._hlc = HLC()
        self._lock = threading.Lock()
        self._subscribers: List[_Subscription] = []
        self._receive_thread: Optional[threading.Thread] = None
        self._running = False

    def set(
        self,
        key: str,
        value: Any,
        *,
        type: ContextType = ContextType.STATE_UPDATE,
        visible_to: Optional[List[str]] = None,
        ttl_ms: Optional[int] = None,
    ) -> None:
        """Set a context entry and broadcast to channel members.

        Args:
            key: Entry key.
            value: Any JSON-serializable value.
            type: Semantic type for the entry.
            visible_to: List of DIDs that can see this entry. If ``None``,
                visible to all channel members.
            ttl_ms: Optional time-to-live in milliseconds. Entry auto-expires
                after this duration.

        Example::

            ctx.set("task_status", {"phase": "running"})
            ctx.set("private", data, visible_to=["did:key:z6MkAgentB..."])
            ctx.set("temp", "busy", ttl_ms=60000)
        """
        scope = (
            AccessScope.specific(visible_to)
            if visible_to
            else AccessScope.all_members()
        )

        with self._lock:
            self._hlc.tick()
            entry = ContextEntry(
                key=key,
                value=value,
                context_type=type,
                scope=scope,
                hlc=HLC(
                    wall_ms=self._hlc.wall_ms,
                    counter=self._hlc.counter,
                    node_id=self._hlc.node_id,
                ),
                ttl_ms=ttl_ms,
            )
            self._store[key] = entry

        payload = _serialize_op("set", key, value, type, scope, entry.hlc, ttl_ms=ttl_ms)
        self._send(payload)

    def get(self, key: str, default: Any = None) -> Any:
        """Get a context entry value.

        Expired entries (past TTL) are removed and return ``default``.

        Args:
            key: Entry key.
            default: Value to return if key is missing or expired.

        Returns:
            The entry value, or *default*.

        Example::

            >>> ctx.set("x", 42)
            >>> ctx.get("x")
            42
            >>> ctx.get("missing", "nope")
            'nope'
        """
        with self._lock:
            entry = self._store.get(key)
            if entry is None:
                return default
            if entry.is_expired():
                del self._store[key]
                return default
            return entry.value

    def delete(self, key: str) -> None:
        """Delete a context entry and broadcast the deletion.

        Args:
            key: Entry key to delete.

        Example::

            ctx.set("temp", "value")
            ctx.delete("temp")
            assert ctx.get("temp") is None
        """
        with self._lock:
            self._hlc.tick()
            self._store.pop(key, None)

        payload = _serialize_op(
            "delete", key, None, ContextType.STATE_UPDATE,
            AccessScope.all_members(), self._hlc,
        )
        self._send(payload)

    def snapshot(self) -> Dict[str, Any]:
        """Return a snapshot of all non-expired context entries.

        Returns:
            Dict mapping keys to their current values.

        Example::

            ctx.set("a", 1)
            ctx.set("b", 2)
            snap = ctx.snapshot()  # {"a": 1, "b": 2}
        """
        with self._lock:
            result = {}
            expired_keys = []
            for key, entry in self._store.items():
                if entry.is_expired():
                    expired_keys.append(key)
                else:
                    result[key] = entry.value
            for key in expired_keys:
                del self._store[key]
            return result

    def list(
        self, *, type_filter: Optional[ContextType] = None
    ) -> List[str]:
        """List context entry keys, optionally filtered by type.

        Args:
            type_filter: If provided, only return keys of this type.

        Returns:
            List of matching key strings.

        Example::

            ctx.set("a", 1, type=ContextType.ARTIFACT)
            ctx.set("b", 2, type=ContextType.STATE_UPDATE)
            ctx.list(type_filter=ContextType.ARTIFACT)  # ["a"]
        """
        with self._lock:
            keys = []
            for key, entry in self._store.items():
                if entry.is_expired():
                    continue
                if type_filter is not None and entry.context_type != type_filter:
                    continue
                keys.append(key)
            return keys

    def handoff(
        self,
        target: str,
        context: Dict[str, Any],
        *,
        tried: Optional[List[str]] = None,
        rejected: Optional[List[str]] = None,
    ) -> None:
        """Send a structured handoff to a specific agent.

        Handoffs carry what was tried and rejected so the receiving agent
        can avoid repeating failed approaches.

        Args:
            target: DID URI of the target agent.
            context: Context data to hand off.
            tried: List of approaches already attempted.
            rejected: List of approaches rejected and why.

        Example::

            ctx.handoff(
                target="did:key:z6MkAgentB...",
                context={"task": "summarize", "docs": ["doc1"]},
                tried=["extractive_summary"],
                rejected=["too_slow"],
            )
        """
        scope = AccessScope.specific([target])
        meta = {}
        if tried:
            meta["tried"] = tried
        if rejected:
            meta["rejected"] = rejected

        with self._lock:
            self._hlc.tick()

        payload = _serialize_op(
            "handoff", f"_handoff_{target}", context,
            ContextType.HANDOFF, scope, self._hlc, meta=meta,
        )
        self._send(payload)

    def subscribe(
        self,
        callback: Callable[[str, Any], None],
        *,
        type_filter: Optional[ContextType] = None,
    ) -> None:
        """Register a callback for context changes.

        The callback receives ``(key, value)`` for each incoming update.
        Optionally filter by :class:`ContextType`.

        Args:
            callback: Function called with ``(key, value)`` on updates.
            type_filter: If set, only fire for entries of this type.

        Example::

            ctx.subscribe(
                lambda key, val: print(f"{key} = {val}"),
                type_filter=ContextType.HANDOFF,
            )
        """
        self._subscribers.append(_Subscription(callback, type_filter))

    def start_receiving(self) -> None:
        """Start a background thread that processes incoming context updates.

        Call this after creating/joining the channel. The thread polls for
        STATE envelopes and applies them to the local store.

        Example::

            ctx = SharedContext(mgr, "acme/research/results")
            ctx.start_receiving()
        """
        if self._running:
            return
        self._running = True
        self._receive_thread = threading.Thread(
            target=self._receive_loop, daemon=True, name="skytale-context-recv"
        )
        self._receive_thread.start()

    def stop_receiving(self) -> None:
        """Stop the background receive thread.

        Example::

            ctx.stop_receiving()
        """
        self._running = False

    def request_snapshot(self) -> None:
        """Request a full state snapshot from other channel members.

        Useful for late-joining agents. Members with state will respond
        with their current snapshot.

        Example::

            ctx.request_snapshot()
        """
        with self._lock:
            self._hlc.tick()
        payload = _serialize_op(
            "snapshot_req", "", None, ContextType.STATE_UPDATE,
            AccessScope.all_members(), self._hlc,
        )
        self._send(payload)

    def _send(self, payload: dict) -> None:
        """Send a STATE envelope through the MLS channel."""
        envelope = Envelope(
            protocol=Protocol.STATE,
            content_type="application/json",
            payload=json.dumps(payload, separators=(",", ":")).encode("utf-8"),
        )
        self._manager.send_envelope(self._channel, envelope)

    def _receive_loop(self) -> None:
        """Background loop: poll for STATE envelopes and apply."""
        while self._running:
            try:
                envelopes = self._manager.receive_envelopes(
                    self._channel, timeout=1.0
                )
                for env in envelopes:
                    if env.protocol == Protocol.STATE:
                        self._apply(env.payload)
            except Exception:
                logger.debug("context receive error", exc_info=True)
            time.sleep(0.05)

    def _apply(self, payload_bytes: bytes) -> None:
        """Apply a received STATE operation to the local store."""
        try:
            data = json.loads(payload_bytes)
        except (json.JSONDecodeError, UnicodeDecodeError):
            return

        op = data.get("op")
        key = data.get("key", "")
        value = data.get("value")
        ctx_type_str = data.get("type", ContextType.STATE_UPDATE.value)
        scope_dict = data.get("scope", {"access": "all"})
        hlc_list = data.get("hlc")
        ttl_ms = data.get("ttl_ms")
        meta = data.get("meta")

        try:
            ctx_type = ContextType(ctx_type_str)
        except ValueError:
            ctx_type = ContextType.STATE_UPDATE

        scope = AccessScope.from_dict(scope_dict)

        # Access control: skip if scoped to specific members and we're not in the list
        if scope.access == "members" and self._agent_id and self._agent_id not in scope.ids:
            return

        remote_hlc = HLC.from_list(hlc_list) if hlc_list else HLC()

        if op == "set":
            with self._lock:
                self._hlc.update(remote_hlc.wall_ms, remote_hlc.counter)
                existing = self._store.get(key)
                # LWW: keep entry with higher HLC
                if existing is not None and not (remote_hlc > existing.hlc):
                    return
                self._store[key] = ContextEntry(
                    key=key,
                    value=value,
                    context_type=ctx_type,
                    scope=scope,
                    hlc=remote_hlc,
                    ttl_ms=ttl_ms,
                )
            self._notify(key, value, ctx_type)

        elif op == "delete":
            with self._lock:
                self._hlc.update(remote_hlc.wall_ms, remote_hlc.counter)
                self._store.pop(key, None)
            self._notify(key, None, ctx_type)

        elif op == "handoff":
            with self._lock:
                self._hlc.update(remote_hlc.wall_ms, remote_hlc.counter)
                self._store[key] = ContextEntry(
                    key=key,
                    value={"context": value, "meta": meta or {}},
                    context_type=ContextType.HANDOFF,
                    scope=scope,
                    hlc=remote_hlc,
                )
            self._notify(key, value, ContextType.HANDOFF)

        elif op == "snapshot_req":
            self._send_snapshot()

        elif op == "snapshot_res":
            entries = data.get("entries", {})
            with self._lock:
                self._hlc.update(remote_hlc.wall_ms, remote_hlc.counter)
                for k, entry_data in entries.items():
                    entry_hlc = HLC.from_list(entry_data.get("hlc", [0, 0, ""]))
                    existing = self._store.get(k)
                    if existing is not None and not (entry_hlc > existing.hlc):
                        continue
                    entry_scope = AccessScope.from_dict(entry_data.get("scope", {"access": "all"}))
                    if entry_scope.access == "members" and self._agent_id and self._agent_id not in entry_scope.ids:
                        continue
                    try:
                        entry_type = ContextType(entry_data.get("type", "state_update"))
                    except ValueError:
                        entry_type = ContextType.STATE_UPDATE
                    self._store[k] = ContextEntry(
                        key=k,
                        value=entry_data.get("value"),
                        context_type=entry_type,
                        scope=entry_scope,
                        hlc=entry_hlc,
                        ttl_ms=entry_data.get("ttl_ms"),
                    )

    def _send_snapshot(self) -> None:
        """Send current state as a snapshot response."""
        with self._lock:
            entries = {}
            for key, entry in self._store.items():
                if entry.is_expired():
                    continue
                entries[key] = {
                    "value": entry.value,
                    "type": entry.context_type.value,
                    "scope": entry.scope.to_dict(),
                    "hlc": entry.hlc.to_list(),
                    "ttl_ms": entry.ttl_ms,
                }
            self._hlc.tick()

        payload = {
            "op": "snapshot_res",
            "entries": entries,
            "hlc": self._hlc.to_list(),
        }
        self._send(payload)

    def _notify(self, key: str, value: Any, ctx_type: ContextType) -> None:
        """Fire subscriber callbacks."""
        for sub in self._subscribers:
            if sub.type_filter is not None and sub.type_filter != ctx_type:
                continue
            try:
                sub.callback(key, value)
            except Exception:
                logger.debug("subscriber error for key %s", key, exc_info=True)


class _Subscription:
    """Internal: holds a subscriber callback and optional type filter."""

    __slots__ = ("callback", "type_filter")

    def __init__(
        self,
        callback: Callable[[str, Any], None],
        type_filter: Optional[ContextType],
    ):
        self.callback = callback
        self.type_filter = type_filter


def _serialize_op(
    op: str,
    key: str,
    value: Any,
    ctx_type: ContextType,
    scope: AccessScope,
    hlc: HLC,
    *,
    ttl_ms: Optional[int] = None,
    meta: Optional[dict] = None,
) -> dict:
    """Build a STATE operation payload dict."""
    payload: dict = {
        "op": op,
        "key": key,
        "type": ctx_type.value,
        "scope": scope.to_dict(),
        "hlc": hlc.to_list(),
    }
    if value is not None:
        payload["value"] = value
    if ttl_ms is not None:
        payload["ttl_ms"] = ttl_ms
    if meta:
        payload["meta"] = meta
    return payload
