"""High-level channel manager for AI agent frameworks.

Wraps ``SkytaleClient`` with background message buffering so that LLM tool
calls (which need synchronous request-response semantics) can read messages
without blocking on the ``messages()`` iterator.

Supports both raw string messages (backward compatible) and structured
:class:`~skytale_sdk.envelope.Envelope` messages for multi-protocol use.

Usage::

    from skytale_sdk.channels import SkytaleChannelManager

    mgr = SkytaleChannelManager(identity=b"my-agent")
    mgr.create("acme/research/results")
    mgr.send("acme/research/results", "hello")
    msgs = mgr.receive("acme/research/results")
"""

from __future__ import annotations

import base64
import logging
import os
import threading
import time
from typing import Dict, List, Optional, Union

logger = logging.getLogger(__name__)

from skytale_sdk._native import PySkytaleClient as SkytaleClient
from skytale_sdk.envelope import Envelope
from skytale_sdk.errors import AuthError, SkytaleError, TransportError


_MAX_BUFFER_SIZE = 10000


def _listener_died_error(channel_name: str, error: str) -> TransportError:
    """Build a consistent error for dead listener threads."""
    return TransportError(
        f"unable to receive messages: listener for '{channel_name}' "
        f"died — {error}",
        code="listener_died",
    )


class ChannelHandle:
    """Holds a ``Channel`` and a daemon thread that buffers incoming messages.

    Internally buffers raw bytes. Use :meth:`drain_bytes` for raw access or
    :meth:`drain` for backward-compatible UTF-8 strings.
    """

    def __init__(self, channel):
        self.channel = channel
        self._buffer: List[bytes] = []
        self._lock = threading.Lock()
        self._running = True
        self._error: Optional[str] = None
        self._thread = threading.Thread(target=self._listen, daemon=True)
        self._thread.start()

    def _listen(self):
        try:
            for msg in self.channel.messages():
                if not self._running:
                    break
                with self._lock:
                    if len(self._buffer) < _MAX_BUFFER_SIZE:
                        self._buffer.append(bytes(msg))
                    else:
                        logger.warning("message buffer full, dropping oldest message")
                        self._buffer.pop(0)
                        self._buffer.append(bytes(msg))
        except Exception as exc:
            self._error = str(exc)
            logger.error("channel listener died: %s", exc)

    def drain_bytes(self) -> List[bytes]:
        """Return all buffered messages as raw bytes and clear the buffer."""
        with self._lock:
            msgs = list(self._buffer)
            self._buffer.clear()
        return msgs

    def drain(self) -> List[str]:
        """Return all buffered messages as UTF-8 strings and clear the buffer.

        Backward-compatible: decodes raw bytes to strings.
        """
        return [b.decode("utf-8", errors="replace") for b in self.drain_bytes()]

    def stop(self):
        self._running = False


class SkytaleChannelManager:
    """Manage encrypted channels with background message buffering.

    Designed for AI agent frameworks (LangGraph, CrewAI, MCP) where tool
    calls need non-blocking access to incoming messages.

    Args:
        identity: Agent identity bytes (or str, will be encoded).
        endpoint: Relay server URL. Defaults to ``SKYTALE_RELAY`` env var
            or ``https://relay.skytale.sh:5000``.
        data_dir: Local directory for MLS state. Defaults to
            ``SKYTALE_DATA_DIR`` env var or ``/tmp/skytale-<identity_hex>``.
        api_key: API key. Defaults to ``SKYTALE_API_KEY`` env var.
        api_url: API server URL. Defaults to ``SKYTALE_API_URL`` env var
            or ``https://api.skytale.sh``.
        mock: Enable local dev mode with in-memory transport (no relay
            needed). Also enabled by ``SKYTALE_MOCK`` env var.

    Example::

        mgr = SkytaleChannelManager(identity=b"agent-1")
        mgr.create("org/ns/chan")
        mgr.send("org/ns/chan", "hello")

        # Local dev mode (no relay required):
        mgr = SkytaleChannelManager(identity=b"agent-1", mock=True)
    """

    def __init__(
        self,
        identity: Union[bytes, str],
        endpoint: Optional[str] = None,
        data_dir: Optional[str] = None,
        api_key: Optional[str] = None,
        api_url: Optional[str] = None,
        mock: bool = False,
    ):
        if isinstance(identity, str):
            identity = identity.encode()

        self._identity = identity

        # Check SKYTALE_MOCK env var.
        env_mock = os.environ.get("SKYTALE_MOCK", "").lower() in ("1", "true", "yes")
        self._mock = mock or env_mock

        self._endpoint = endpoint or os.environ.get(
            "SKYTALE_RELAY", "https://relay.skytale.sh:5000"
        )
        id_hex = identity.hex()
        self._data_dir = data_dir or os.environ.get(
            "SKYTALE_DATA_DIR",
            os.path.join(os.path.expanduser("~"), ".skytale", id_hex),
        )
        from skytale_sdk._api import _RedactedStr

        _raw_key = api_key or os.environ.get("SKYTALE_API_KEY")
        self._api_key = _RedactedStr(_raw_key) if _raw_key else None
        self._api_url = api_url or os.environ.get(
            "SKYTALE_API_URL", "https://api.skytale.sh"
        )

        if self._mock:
            self._client = SkytaleClient(
                identity=self._identity, mock=True
            )
        else:
            kwargs = {}
            if self._api_key:
                kwargs["api_key"] = self._api_key
                kwargs["api_url"] = self._api_url

            self._client = SkytaleClient(
                self._endpoint, self._data_dir, self._identity, **kwargs
            )
        self._handles: Dict[str, ChannelHandle] = {}

        # API client for invite-token-based key exchange (skip in mock mode).
        self._api = None
        if not self._mock and self._api_key and self._api_url:
            from skytale_sdk._api import SkytaleAPI

            self._api = SkytaleAPI(self._api_url, self._api_key)
        self._join_poller_running = False
        self._join_poller_thread: Optional[threading.Thread] = None

    def create(self, channel_name: str) -> None:
        """Create a new encrypted channel and start listening for messages.

        If an API key is configured, the channel is registered with the API
        server and a background thread starts polling for incoming join
        requests (see :meth:`invite`).

        Args:
            channel_name: Channel name in ``org/namespace/service`` format.

        Raises:
            RuntimeError: If the channel name is invalid or creation fails.
        """
        channel = self._client.create_channel(channel_name)
        self._handles[channel_name] = ChannelHandle(channel)

        if self._api:
            try:
                self._api.register_channel(channel_name)
            except RuntimeError:
                logger.debug(
                    "channel registration with API skipped: %s", channel_name
                )
            self._start_join_poller()

    def join(self, channel_name: str, welcome: bytes) -> None:
        """Join an existing channel using an MLS Welcome message.

        Args:
            channel_name: Channel name to join.
            welcome: MLS Welcome message bytes from ``add_member()``.

        Raises:
            RuntimeError: If the Welcome is invalid or joining fails.
        """
        channel = self._client.join_channel(channel_name, welcome)
        self._handles[channel_name] = ChannelHandle(channel)

    def key_package(self) -> bytes:
        """Generate an MLS key package for joining a channel.

        Returns:
            Key package bytes to send to the channel creator.
        """
        return self._client.generate_key_package()

    def add_member(self, channel_name: str, key_package: bytes) -> bytes:
        """Add a member to a channel and return the MLS Welcome message.

        Args:
            channel_name: Channel to add the member to.
            key_package: MLS key package from the joining agent.

        Returns:
            MLS Welcome message bytes. Send these to the joining agent.

        Raises:
            KeyError: If the channel does not exist.
        """
        handle = self._handles[channel_name]
        return handle.channel.add_member(key_package)

    def send(self, channel_name: str, message: Union[str, bytes]) -> None:
        """Send a message on an encrypted channel.

        Args:
            channel_name: Target channel.
            message: Message content. Strings are UTF-8 encoded automatically.

        Raises:
            KeyError: If the channel does not exist.
        """
        if isinstance(message, str):
            message = message.encode("utf-8")
        handle = self._handles[channel_name]
        handle.channel.send(message)

    def receive(self, channel_name: str, timeout: float = 5.0) -> List[str]:
        """Drain all buffered messages from a channel.

        Waits up to *timeout* seconds for at least one message if the buffer
        is currently empty.

        Args:
            channel_name: Channel to read from.
            timeout: Max seconds to wait for messages. ``0`` returns
                immediately.

        Returns:
            List of message strings (may be empty if timeout expires).

        Raises:
            KeyError: If the channel does not exist.
            RuntimeError: If the background listener thread has died.
        """
        handle = self._handles[channel_name]
        if handle._error is not None:
            raise _listener_died_error(channel_name, handle._error)
        deadline = time.monotonic() + timeout
        while True:
            msgs = handle.drain()
            if msgs:
                return msgs
            if handle._error is not None:
                raise _listener_died_error(channel_name, handle._error)
            if time.monotonic() >= deadline:
                return []
            time.sleep(min(0.1, max(0, deadline - time.monotonic())))

    def receive_latest(
        self, channel_name: str, timeout: float = 5.0
    ) -> Optional[str]:
        """Return only the most recent buffered message, discarding older ones.

        Args:
            channel_name: Channel to read from.
            timeout: Max seconds to wait.

        Returns:
            The latest message string, or ``None`` if no messages arrive
            before timeout.

        Raises:
            RuntimeError: If the background listener thread has died.
        """
        handle = self._handles[channel_name]
        if handle._error is not None:
            raise _listener_died_error(channel_name, handle._error)
        msgs = self.receive(channel_name, timeout=timeout)
        return msgs[-1] if msgs else None

    def send_envelope(self, channel_name: str, envelope: Envelope) -> None:
        """Send a structured envelope on an encrypted channel.

        Serializes the envelope to wire format and sends through the channel.
        The receiver can reconstruct the envelope with :meth:`receive_envelopes`.

        Args:
            channel_name: Target channel.
            envelope: Envelope to send.

        Raises:
            KeyError: If the channel does not exist.

        Example::

            from skytale_sdk.envelope import Envelope, Protocol

            env = Envelope(Protocol.A2A, "application/json", b'{"parts":[]}')
            mgr.send_envelope("org/ns/chan", env)
        """
        handle = self._handles[channel_name]
        handle.channel.send(envelope.serialize())

    def receive_envelopes(
        self, channel_name: str, timeout: float = 5.0
    ) -> List[Envelope]:
        """Receive structured envelopes from a channel.

        Returns all buffered messages deserialized as :class:`Envelope` objects.
        Messages that fail deserialization (e.g. raw strings sent via
        :meth:`send`) are wrapped in a ``Protocol.RAW`` envelope.

        Args:
            channel_name: Channel to read from.
            timeout: Max seconds to wait for messages. ``0`` returns immediately.

        Returns:
            List of :class:`Envelope` objects (may be empty if timeout expires).

        Raises:
            KeyError: If the channel does not exist.
            RuntimeError: If the background listener thread has died.

        Example::

            envelopes = mgr.receive_envelopes("org/ns/chan")
            for env in envelopes:
                print(env.protocol, env.payload)
        """
        from skytale_sdk.envelope import Protocol

        handle = self._handles[channel_name]
        if handle._error is not None:
            raise _listener_died_error(channel_name, handle._error)
        deadline = time.monotonic() + timeout
        while True:
            raw_msgs = handle.drain_bytes()
            if raw_msgs:
                envelopes = []
                for raw in raw_msgs:
                    try:
                        envelopes.append(Envelope.deserialize(raw))
                    except (ValueError, KeyError):
                        # Not an envelope — wrap raw bytes as RAW protocol
                        envelopes.append(
                            Envelope(
                                protocol=Protocol.RAW,
                                content_type="application/octet-stream",
                                payload=raw,
                            )
                        )
                return envelopes
            if handle._error is not None:
                raise _listener_died_error(channel_name, handle._error)
            if time.monotonic() >= deadline:
                return []
            time.sleep(min(0.1, max(0, deadline - time.monotonic())))

    # ------------------------------------------------------------------
    # Invite-token-based key exchange
    # ------------------------------------------------------------------

    def invite(
        self, channel_name: str, max_uses: int = 1, ttl: int = 3600
    ) -> str:
        """Create an invite token for a channel.

        The returned token can be shared with another agent who calls
        :meth:`join_with_token` to join the channel.  The MLS key exchange
        happens automatically through the API server — neither agent needs
        to handle key packages or Welcome messages directly.

        Args:
            channel_name: Channel to create the invite for (must exist).
            max_uses: How many times the token can be used (default 1).
            ttl: Token lifetime in seconds (default 3600 = 1 hour).

        Returns:
            Invite token string (``skt_inv_...``).

        Raises:
            RuntimeError: If API is not configured or the request fails.

        Example::

            token = mgr.invite("acme/team/general")
            # Share 'token' with the agent that wants to join
        """
        if not self._api:
            raise AuthError(
                "unable to create invite: no API key configured. "
                "Pass api_key= to SkytaleChannelManager or set SKYTALE_API_KEY.",
                code="api_key_required",
            )
        resp = self._api.create_invite(channel_name, max_uses, ttl)
        return resp["token"]

    def join_with_token(
        self, channel_name: str, token: str, timeout: float = 60.0
    ) -> None:
        """Join a channel using an invite token.

        Generates an MLS key package, submits it to the API with the token,
        then polls until the channel owner processes it and uploads the
        MLS Welcome message.  Once received the agent joins the channel
        and starts listening for messages.

        Args:
            channel_name: Channel name to join.
            token: Invite token (``skt_inv_...``) from the channel owner.
            timeout: Max seconds to wait for the owner to process the join.

        Raises:
            RuntimeError: If API is not configured, the token is invalid,
                or the owner does not process the join within *timeout*.

        Example::

            mgr.join_with_token("acme/team/general", "skt_inv_abc123")
        """
        if not self._api:
            raise AuthError(
                "unable to join with token: no API key configured. "
                "Pass api_key= to SkytaleChannelManager or set SKYTALE_API_KEY.",
                code="api_key_required",
            )

        # Generate key package and submit join request.
        kp = self._client.generate_key_package()
        kp_b64 = base64.b64encode(kp).decode("ascii")
        identity_str = (
            self._identity.decode("utf-8", errors="replace")
            if isinstance(self._identity, bytes)
            else self._identity
        )

        resp = self._api.join(channel_name, token, kp_b64, identity_str)
        request_id = resp["request_id"]

        # Poll for the Welcome message.
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            resp = self._api.get_welcome(request_id)
            if resp.get("status") == "completed" and resp.get("welcome"):
                welcome = base64.b64decode(resp["welcome"])
                self.join(channel_name, welcome)
                logger.info("joined channel %s via invite token", channel_name)
                return
            time.sleep(1.0)

        raise TransportError(
            f"unable to join: timed out after {timeout}s waiting for channel "
            "owner to process the join request. Ensure the channel owner is "
            "running with an API key configured.",
            code="join_timeout",
        )

    # ------------------------------------------------------------------
    # Background join-request processing
    # ------------------------------------------------------------------

    def _start_join_poller(self) -> None:
        """Start the background thread that auto-processes pending joins."""
        if self._join_poller_running:
            return
        self._join_poller_running = True
        self._join_poller_thread = threading.Thread(
            target=self._join_poll_loop,
            daemon=True,
            name="skytale-join-poller",
        )
        self._join_poller_thread.start()

    def _join_poll_loop(self) -> None:
        """Poll the API for pending join requests and process them.

        Processes at most one request per channel per poll cycle to
        avoid rapid MLS epoch advancement under high join volume.
        """
        while self._join_poller_running:
            for channel_name in list(self._handles.keys()):
                try:
                    resp = self._api.get_pending(channel_name)
                    requests = resp.get("requests", [])
                    if not requests:
                        continue
                    # Process at most one request per poll cycle to rate-limit
                    # MLS epoch advancement.
                    req = requests[0]
                    kp = base64.b64decode(req["key_package"])
                    welcome = self.add_member(channel_name, kp)
                    welcome_b64 = base64.b64encode(welcome).decode("ascii")
                    self._api.submit_welcome(req["id"], welcome_b64)
                    logger.info(
                        "auto-processed join request %s on %s (identity=%s)",
                        req["id"],
                        channel_name,
                        req.get("identity", "?"),
                    )
                except Exception:
                    logger.warning(
                        "join poll error for %s", channel_name, exc_info=True
                    )
            time.sleep(2.0)

    def list_channels(self) -> List[str]:
        """Return the names of all active channels.

        Returns:
            Sorted list of channel name strings.
        """
        return sorted(self._handles.keys())

    def close(self) -> None:
        """Stop all background threads (listeners and join poller)."""
        self._join_poller_running = False
        for handle in self._handles.values():
            handle.stop()
        self._handles.clear()

    def __enter__(self) -> "SkytaleChannelManager":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    def __del__(self):
        self.close()
