"""Streaming API for the Meerkat Python SDK.

Provides :class:`_StdoutDispatcher` (internal) for multiplexing the rkat-rpc
stdout stream, and :class:`EventStream` (public) for consuming typed agent
events in real-time.
"""

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, AsyncIterator, Callable

from .errors import MeerkatError
from .events import Event, TextDelta, UnknownEvent, parse_event

if TYPE_CHECKING:
    from .session import Session
    from .types import RunResult

RPC_STDOUT_LIMIT_BYTES = 64 * 1024 * 1024


class _StdoutDispatcher:
    """Background reader that multiplexes stdout lines to response futures and event queues."""

    def __init__(self, stdout: asyncio.StreamReader):
        self._stdout = stdout
        self._pending_responses: dict[int, asyncio.Future[dict[str, Any]]] = {}
        self._event_queues: dict[str, asyncio.Queue[dict[str, Any] | None]] = {}
        self._stream_queues: dict[str, asyncio.Queue[dict[str, Any] | None]] = {}
        self._pending_stream_queue: asyncio.Queue[dict[str, Any] | None] | None = None
        self._pending_stream_request_id: int | None = None
        self._unmatched_buffer: dict[str, list[dict[str, Any]]] = {}
        self._unmatched_stream_buffer: dict[str, list[dict[str, Any]]] = {}
        self._stream_terminal: dict[str, dict[str, Any]] = {}
        self._unmatched_stream_end: set[str] = set()
        self._task: asyncio.Task[None] | None = None
        self._closed = False
        self._tool_handler: Any = None  # ToolRegistry, set lazily
        self._stdin_writer: asyncio.StreamWriter | None = None  # set on connect

    def start(self) -> None:
        self._task = asyncio.get_running_loop().create_task(self._read_loop())

    async def stop(self) -> None:
        self._closed = True
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        self._fail_all("CLIENT_CLOSED", "client stopped")

    def expect_response(self, request_id: int) -> asyncio.Future[dict[str, Any]]:
        future: asyncio.Future[dict[str, Any]] = asyncio.get_running_loop().create_future()
        self._pending_responses[request_id] = future
        return future

    def subscribe_events(self, session_id: str) -> asyncio.Queue[dict[str, Any] | None]:
        queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue()
        self._event_queues[session_id] = queue
        return queue

    def unsubscribe_events(self, session_id: str) -> None:
        self._event_queues.pop(session_id, None)

    def subscribe_stream(self, stream_id: str) -> asyncio.Queue[dict[str, Any] | None]:
        queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue()
        self._stream_queues[stream_id] = queue
        for event in self._unmatched_stream_buffer.pop(stream_id, []):
            queue.put_nowait(event)
        if stream_id in self._unmatched_stream_end:
            self._unmatched_stream_end.discard(stream_id)
            queue.put_nowait(None)
        return queue

    def unsubscribe_stream(self, stream_id: str) -> None:
        self._stream_queues.pop(stream_id, None)

    def stream_terminal(self, stream_id: str) -> dict[str, Any] | None:
        return self._stream_terminal.get(stream_id)

    def subscribe_pending_stream(self, request_id: int) -> asyncio.Queue[dict[str, Any] | None]:
        if self._pending_stream_queue is not None:
            raise RuntimeError("Only one pending stream at a time")
        self._pending_stream_queue = asyncio.Queue()
        self._pending_stream_request_id = request_id
        return self._pending_stream_queue

    def unsubscribe_pending_stream(self) -> None:
        self._pending_stream_queue = None
        self._pending_stream_request_id = None
        self._unmatched_buffer.clear()

    def set_tool_handler(self, handler: Any) -> None:
        """Set the tool registry for handling callback tool requests."""
        self._tool_handler = handler

    def set_stdin_writer(self, process: Any) -> None:
        """Set the process for writing callback responses."""
        self._stdin_writer = process

    async def _read_loop(self) -> None:
        while not self._closed:
            line = await self._stdout.readline()
            if not line:
                self._fail_all("CONNECTION_CLOSED", "rkat-rpc process closed")
                return
            try:
                data = json.loads(line)
            except json.JSONDecodeError:
                continue
            # Server→client callback request (has both id and method).
            # Spawn as a separate task to avoid blocking the read loop — if the
            # tool handler makes a nested SDK call, the response needs to flow
            # through this same _read_loop.
            if "id" in data and "method" in data:
                asyncio.get_running_loop().create_task(
                    self._handle_callback_request(data)
                )
                continue
            if "id" in data:
                request_id = data["id"]
                future = self._pending_responses.pop(request_id, None)
                if future and not future.done():
                    if data.get("error"):
                        err = data["error"]
                        raw_message = err.get("message", "Unknown error")
                        nested_code = None
                        nested_message = None
                        nested_details = err.get("data")
                        if isinstance(nested_details, dict):
                            nested_code = nested_details.get("code")
                            nested_message = nested_details.get("message")
                            nested_details = nested_details.get(
                                "details",
                                nested_details.get("reason", nested_details),
                            )
                        if isinstance(raw_message, str):
                            try:
                                parsed = json.loads(raw_message)
                            except json.JSONDecodeError:
                                parsed = None
                            if isinstance(parsed, dict):
                                nested_code = parsed.get("code")
                                nested_message = parsed.get("message")
                                nested_details = parsed.get("details", parsed.get("reason", nested_details))
                        future.set_exception(
                            MeerkatError(
                                str(nested_code or err.get("code", "UNKNOWN")),
                                str(nested_message or raw_message),
                                nested_details,
                            )
                        )
                    else:
                        result = data.get("result", {})
                        future.set_result(result)
                        if (
                            request_id == self._pending_stream_request_id
                            and self._pending_stream_queue is not None
                        ):
                            sid = result.get("session_id", "")
                            if sid:
                                for evt in self._unmatched_buffer.pop(sid, []):
                                    await self._pending_stream_queue.put(evt)
                                self._event_queues[sid] = self._pending_stream_queue
                            self._pending_stream_queue = None
                            self._pending_stream_request_id = None
                            self._unmatched_buffer.clear()
            elif "method" in data:
                method = str(data.get("method", ""))
                params = data.get("params", {})
                if method in {"session/stream_event", "mob/stream_event"}:
                    stream_id = str(params.get("stream_id", ""))
                    raw_event = params.get("event") or params
                    # Preserve scope fields when present (delegated-branch / mob-member scoped events).
                    scope_id = params.get("scope_id")
                    scope_path = params.get("scope_path")
                    if scope_id is not None or scope_path is not None:
                        event = {
                            "event": raw_event,
                            "scope_id": scope_id,
                            "scope_path": scope_path,
                        }
                    else:
                        event = raw_event
                    queue = self._stream_queues.get(stream_id)
                    if queue is not None:
                        await queue.put(event)
                    elif stream_id:
                        self._unmatched_stream_buffer.setdefault(stream_id, []).append(event)
                    continue
                if method in {"session/stream_end", "mob/stream_end"}:
                    stream_id = str(params.get("stream_id", ""))
                    if stream_id:
                        self._stream_terminal[stream_id] = dict(params)
                    queue = self._stream_queues.get(stream_id)
                    if queue is not None:
                        await queue.put(None)
                    elif stream_id:
                        self._unmatched_stream_end.add(stream_id)
                    continue
                session_id = params.get("session_id", "")
                event = params.get("event")
                queue = self._event_queues.get(session_id)
                if queue is not None:
                    await queue.put(event)
                elif self._pending_stream_queue is not None:
                    self._unmatched_buffer.setdefault(session_id, []).append(event)

    async def _handle_callback_request(self, data: dict[str, Any]) -> None:
        """Handle an incoming server→client tool/execute request."""
        request_id = data.get("id")
        method = data.get("method", "")
        params = data.get("params", {})

        if method == "tool/execute" and self._tool_handler is not None:
            tool_name = params.get("name", "")
            arguments = params.get("arguments", {})
            content, is_error = await self._tool_handler.handle(tool_name, arguments)
            response = {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {"content": content, "is_error": is_error},
            }
        else:
            response = {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32601,
                    "message": f"Method not supported: {method}",
                },
            }

        self._write_response(response)

    def _write_response(self, response: dict[str, Any]) -> None:
        """Write a JSON-RPC response back to the server via stdin."""
        if self._stdin_writer is not None:
            try:
                data = (json.dumps(response) + "\n").encode()
                self._stdin_writer.write(data)
            except Exception:
                pass  # Best-effort — process may have died.

    def _fail_all(self, code: str, message: str) -> None:
        for future in self._pending_responses.values():
            if not future.done():
                future.set_exception(MeerkatError(code, message))
        self._pending_responses.clear()
        for queue in self._event_queues.values():
            queue.put_nowait(None)
        self._event_queues.clear()
        for queue in self._stream_queues.values():
            queue.put_nowait(None)
        self._stream_queues.clear()
        self._unmatched_stream_buffer.clear()
        self._unmatched_stream_end.clear()
        self._stream_terminal.clear()
        if self._pending_stream_queue is not None:
            self._pending_stream_queue.put_nowait(None)
            self._pending_stream_queue = None
            self._pending_stream_request_id = None
            self._unmatched_buffer.clear()


class EventStream:
    """Async iterable of **typed** events from a running turn.

    Use as an async context manager::

        async with session.stream("prompt") as events:
            async for event in events:
                match event:
                    case TextDelta(delta=chunk):
                        print(chunk, end="")
            result = events.result

    Or consume in one shot::

        result = await session.stream("prompt").collect()
    """

    def __init__(
        self,
        *,
        session_id: str,
        event_queue: asyncio.Queue[dict[str, Any] | None],
        response_future: asyncio.Future[dict[str, Any]],
        dispatcher: _StdoutDispatcher,
        parse_result: Callable[[dict[str, Any]], RunResult],
        pending_send: tuple[Any, bytes] | None = None,
        session: Session | None = None,
    ):
        self._session_id = session_id
        self._event_queue = event_queue
        self._response_future = response_future
        self._dispatcher = dispatcher
        self._parse_result = parse_result
        self._result: RunResult | None = None
        self._pending_send = pending_send
        self._session = session

    @property
    def session_id(self) -> str:
        return self._session_id

    @property
    def result(self) -> RunResult:
        """The final :class:`~meerkat.RunResult`.  Available after iteration completes."""
        if self._result is None:
            raise MeerkatError(
                "STREAM_NOT_CONSUMED",
                "Iterate the stream before accessing result",
            )
        return self._result

    async def __aenter__(self) -> EventStream:
        if self._pending_send is not None:
            stdin, data = self._pending_send
            self._pending_send = None
            stdin.write(data)
            await stdin.drain()
        return self

    async def __aexit__(self, *_exc: object) -> None:
        if self._session_id:
            self._dispatcher.unsubscribe_events(self._session_id)
        else:
            self._dispatcher.unsubscribe_pending_stream()

    def __aiter__(self) -> AsyncIterator[Event]:
        return self._iter_events()

    async def _iter_events(self) -> AsyncIterator[Event]:
        """Yield typed events until the response future resolves."""
        queue_get: asyncio.Task[dict[str, Any] | None] | None = None
        response_task = asyncio.ensure_future(self._response_future)
        try:
            while True:
                if response_task.done():
                    while not self._event_queue.empty():
                        raw = self._event_queue.get_nowait()
                        if raw is None:
                            break
                        yield parse_event(raw)
                    self._finalise(await response_task)
                    return

                if queue_get is None or queue_get.done():
                    queue_get = asyncio.ensure_future(self._event_queue.get())

                done, _ = await asyncio.wait(
                    [queue_get, response_task],
                    return_when=asyncio.FIRST_COMPLETED,
                )

                if queue_get in done:
                    raw = queue_get.result()
                    queue_get = None
                    if raw is None:
                        self._finalise(await response_task)
                        return
                    yield parse_event(raw)
        finally:
            if queue_get is not None and not queue_get.done():
                queue_get.cancel()

    def _finalise(self, raw_result: dict[str, Any]) -> None:
        self._result = self._parse_result(raw_result)
        if not self._session_id:
            self._session_id = self._result.session_id
        if self._session is not None:
            self._session._last_result = self._result  # noqa: SLF001

    async def collect(self) -> RunResult:
        """Consume all events silently and return the final result."""
        async with self:
            async for _ in self:
                pass
        return self.result

    async def collect_text(self) -> tuple[str, RunResult]:
        """Consume events, accumulate text deltas, return ``(full_text, result)``."""
        parts: list[str] = []
        async with self:
            async for event in self:
                if isinstance(event, TextDelta):
                    parts.append(event.delta)
        return "".join(parts), self.result


class EventSubscription:
    """Async iterable standalone event subscription (session/member/mob)."""

    def __init__(
        self,
        *,
        stream_id: str,
        event_queue: asyncio.Queue[dict[str, Any] | None],
        dispatcher: _StdoutDispatcher,
        close_remote: Callable[[str], Any],
        parse_event_fn: Callable[[dict[str, Any]], Any],
    ):
        self._stream_id = stream_id
        self._event_queue = event_queue
        self._dispatcher = dispatcher
        self._close_remote = close_remote
        self._parse_event_fn = parse_event_fn
        self._closed = False

    @property
    def stream_id(self) -> str:
        return self._stream_id

    @property
    def terminal_outcome(self) -> dict[str, Any] | None:
        return self._dispatcher.stream_terminal(self._stream_id)

    async def __aenter__(self) -> EventSubscription:
        return self

    async def __aexit__(self, *_exc: object) -> None:
        await self.close()

    async def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        self._dispatcher.unsubscribe_stream(self._stream_id)
        self._event_queue.put_nowait(None)
        await self._close_remote(self._stream_id)

    def __aiter__(self) -> AsyncIterator[Any]:
        return self._iter_events()

    async def _iter_events(self) -> AsyncIterator[Any]:
        while True:
            raw = await self._event_queue.get()
            if raw is None:
                return
            yield self._parse_event_fn(raw)
