"""In-process transport: bridges AICC clients to a Bridge within one event loop.

Useful for tests, embedded use, and single-process agents that drive the bridge directly.
"""

from __future__ import annotations

import asyncio
from typing import Any, Self

from aicc.bridge import Bridge
from aicc.errors import ConnectionError_
from aicc.protocol import SessionInit


class InProcessTransport:
    """Two-sided transport: one side is the bridge, the other is the client.

    Both sides share the same asyncio event loop. The bridge side pushes its
    manifest into the client's receive queue when start() is awaited.

    Example:
        async with InProcessTransport.start(bridge) as t:
            client = AICCClient(t)
            async with client:
                manifest = await client.handshake()
    """

    def __init__(self) -> None:
        self._send_q: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        self._recv_q: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        self._bridge: Bridge | None = None
        self._session_id: str | None = None
        self._connected = False

    @classmethod
    def start(cls, bridge: Bridge) -> InProcessTransport:
        t = cls()
        t._bridge = bridge
        return t

    async def connect(self) -> None:
        if self._bridge is None:
            raise ConnectionError_("InProcessTransport.start(bridge) must be used for server-side.")
        if self._connected:
            return
        self._connected = True
        session = self._bridge.open_session()
        self._session_id = session.session_id
        manifest: SessionInit = session.manifest  # type: ignore[assignment]
        await self._recv_q.put(manifest.model_dump(by_alias=True))

    async def close(self) -> None:
        if self._bridge is not None and self._session_id is not None:
            self._bridge.close_session(self._session_id)
            self._session_id = None
        self._connected = False

    async def send(self, message: dict[str, Any]) -> None:
        if not self._connected or self._bridge is None or self._session_id is None:
            raise ConnectionError_("Transport is not connected.")
        sid = message.get("session_id") or self._session_id
        msg_type = message.get("type")
        # session_resume carries the session_id; bind it
        if msg_type == "session_resume" and not self._bridge._sessions.get(sid):
            # Auto-open if bridge forgot
            pass
        response = await self._bridge.handle_message(sid, message)
        if response is not None:
            await self._recv_q.put(response)
        # Drain any events the bridge queued during handling (e.g. from
        # tool handlers calling bridge.emit_event). Deliver them after the
        # response so the client sees a coherent request -> response order.
        sess = self._bridge._sessions.get(sid)
        if sess is not None:
            while not sess.send_queue.empty():
                event = sess.send_queue.get_nowait()
                await self._recv_q.put(event)

    async def receive(self) -> dict[str, Any]:
        if not self._connected:
            raise ConnectionError_("Transport is not connected.")
        return await self._recv_q.get()

    async def push_event(self, event: dict[str, Any]) -> None:
        """Bridge-side helper to push an event to the connected client."""
        await self._recv_q.put(event)

    async def __aenter__(self) -> Self:
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.close()