"""Editor-side stdio client. Rejects non-JSON-RPC stdout on every exchange.

A paused reader forces pipe backpressure and adapter queue overflow.
"""

from __future__ import annotations

import itertools
import json
import os
import subprocess
import threading
import time
from typing import Any, NamedTuple

from harness import ROOT

CLI = ROOT / "cli/bin/albedo"


class Exchange(NamedTuple):
    """A response, and the notifications the adapter wrote before it."""

    response: dict[str, Any]
    notifications: list[dict[str, Any]]


def _key(value: Any) -> str:
    return json.dumps(value, sort_keys=True)


def _problem(message: Any) -> str | None:
    """Why a parsed stdout line is not a JSON-RPC 2.0 message, or None."""
    if not isinstance(message, dict):
        return "not an object"
    if message.get("jsonrpc") != "2.0":
        return "missing jsonrpc 2.0"
    if "method" in message:
        return None if isinstance(message["method"], str) else "method is not a string"
    if "id" not in message:
        return "neither a request, a notification, nor a response"
    if ("result" in message) == ("error" in message):
        return "a response needs exactly one of result or error"
    return None


class AcpClient:
    def __init__(self, env: dict[str, str]):
        self.env = env
        # One condition guards every list below. It is an RLock, so helpers
        # that hold it may call each other.
        self._cond = threading.Condition()
        self._messages: list[dict[str, Any]] = []
        self._violations: list[str] = []
        self._stderr = bytearray()
        self._eof = False
        self._started: dict[str, int] = {}
        self._write_lock = threading.Lock()
        self._ids = itertools.count(1)
        self._reading = threading.Event()
        self._reading.set()
        self.process = subprocess.Popen(
            [env.get("ALBEDO_TEST_CLI", str(CLI)), "acp"],
            cwd=ROOT,
            env=env,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        self._threads = [
            threading.Thread(target=self._read_stdout, daemon=True),
            threading.Thread(target=self._read_stderr, daemon=True),
        ]
        for thread in self._threads:
            thread.start()

    def write(self, data: bytes) -> None:
        """Write raw bytes to stdin, flushed as one write."""
        stream = self.process.stdin
        assert stream is not None
        with self._write_lock:
            try:
                stream.write(data)
                stream.flush()
            except (OSError, ValueError) as error:
                raise AssertionError(
                    f"adapter no longer reads stdin: {error}\n{self.stderr()}"
                ) from error

    def write_fragmented(self, data: bytes, size: int, pause: float = 0.02) -> None:
        """Write raw bytes in pieces, so the adapter reads a partial line."""
        for offset in range(0, len(data), size):
            self.write(data[offset : offset + size])
            time.sleep(pause)

    def expect(self, request_id: Any) -> None:
        """Mark where a response to ``request_id`` may start, before it is sent."""
        with self._cond:
            self._started[_key(request_id)] = len(self._messages)

    def send(
        self,
        method: str,
        params: Any = None,
        *,
        id: Any = None,
        fragment: int = 0,
    ) -> Any:
        """Send a request and return its id. ``fragment`` splits the write."""
        if id is None:
            id = next(self._ids)
        message: dict[str, Any] = {"jsonrpc": "2.0", "id": id, "method": method}
        if params is not None:
            message["params"] = params
        data = json.dumps(message).encode() + b"\n"
        self.expect(id)
        if fragment:
            self.write_fragmented(data, fragment)
        else:
            self.write(data)
        return id

    def notify(self, method: str, params: Any = None) -> None:
        message: dict[str, Any] = {"jsonrpc": "2.0", "method": method}
        if params is not None:
            message["params"] = params
        self.write(json.dumps(message).encode() + b"\n")

    def request(
        self, method: str, params: Any = None, *, id: Any = None, timeout: float = 30
    ) -> Exchange:
        return self.response(self.send(method, params, id=id), timeout=timeout)

    def response(self, request_id: Any, timeout: float = 30) -> Exchange:
        """Wait for the response to a request sent with ``send`` or ``expect``."""
        key = _key(request_id)
        with self._cond:
            start = self._started.pop(key)
        deadline = time.monotonic() + timeout
        with self._cond:
            while True:
                self._raise_violations()
                notifications = []
                for message in self._messages[start:]:
                    if "method" in message:
                        if "id" not in message:
                            notifications.append(message)
                    elif _key(message.get("id")) == key:
                        return Exchange(message, notifications)
                if self._eof:
                    raise AssertionError(
                        f"adapter exited before answering {request_id}\n{self.stderr()}"
                    )
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise AssertionError(
                        f"no response to {request_id} within {timeout}s\n{self.stderr()}"
                    )
                self._cond.wait(remaining)

    def messages(self) -> list[dict[str, Any]]:
        """Every stdout message so far, in arrival order (the same objects)."""
        with self._cond:
            return list(self._messages)

    def notifications(self, method: str | None = None) -> list[dict[str, Any]]:
        return [
            message
            for message in self.messages()
            if "method" in message
            and "id" not in message
            and (method is None or message["method"] == method)
        ]

    def updates(self, session_id: str) -> list[dict[str, Any]]:
        """The ``update`` payloads of ``session/update`` for one session."""
        return [
            message["params"]["update"]
            for message in self.notifications("session/update")
            if message["params"].get("sessionId") == session_id
        ]

    def requests_from_adapter(self) -> list[dict[str, Any]]:
        """Requests the adapter sent to the editor, such as a permission request."""
        return [
            message
            for message in self.messages()
            if "method" in message and "id" in message
        ]

    def wait_for(self, predicate, timeout: float = 30, message: str = "condition"):
        """Poll ``predicate`` until it returns a truthy value, and return that."""
        deadline = time.monotonic() + timeout
        while True:
            value = predicate()
            if value:
                return value
            with self._cond:
                self._raise_violations()
            if time.monotonic() >= deadline:
                raise AssertionError(
                    f"timed out after {timeout}s waiting for {message}\n{self.stderr()}"
                )
            time.sleep(0.02)

    def stderr(self) -> str:
        with self._cond:
            return bytes(self._stderr).decode(errors="replace")

    def check_stdout(self) -> None:
        """Fail if stdout ever carried anything but ACP messages."""
        with self._cond:
            self._raise_violations()

    def _accept(self, raw: bytes) -> None:
        try:
            message = json.loads(raw)
        except ValueError:
            problem = "not JSON"
        else:
            problem = _problem(message)
        with self._cond:
            if problem is None:
                self._messages.append(message)
            else:
                self._violations.append(f"{problem}: {raw[:300]!r}")
            self._cond.notify_all()

    def _read_stdout(self) -> None:
        stream = self.process.stdout
        assert stream is not None
        buffered = b""
        while True:
            self._reading.wait()
            chunk = os.read(stream.fileno(), 1 << 16)
            if not chunk:
                break
            buffered += chunk
            start = 0
            while (end := buffered.find(b"\n", start)) >= 0:
                self._accept(buffered[start:end])
                start = end + 1
            buffered = buffered[start:]
        with self._cond:
            if buffered:
                self._violations.append(f"stdout ended mid-line: {buffered[:300]!r}")
            self._eof = True
            self._cond.notify_all()

    def _read_stderr(self) -> None:
        stream = self.process.stderr
        assert stream is not None
        while chunk := os.read(stream.fileno(), 1 << 16):
            with self._cond:
                self._stderr.extend(chunk)

    def _raise_violations(self) -> None:
        if self._violations:
            raise AssertionError(
                "adapter wrote to stdout outside ACP:\n"
                + "\n".join(self._violations[:5])
                + "\nstderr:\n"
                + self.stderr()
            )

    def pause_reading(self) -> None:
        self._reading.clear()

    def signal(self, number: int) -> None:
        self.process.send_signal(number)

    def close_stdin(self) -> None:
        """Close the editor's side so the adapter sees EOF on stdin."""
        stream = self.process.stdin
        if stream is not None and not stream.closed:
            try:
                stream.close()
            except BrokenPipeError:
                pass

    def wait_exit(self, timeout: float = 30) -> int:
        try:
            return self.process.wait(timeout=timeout)
        except subprocess.TimeoutExpired as error:
            raise AssertionError(
                f"adapter still running after {timeout}s\n{self.stderr()}"
            ) from error

    def close(self, timeout: float = 30) -> int:
        """Close stdin, wait for the adapter to exit, and check its stdout."""
        self.close_stdin()
        code = self.wait_exit(timeout)
        self._join()
        self.check_stdout()
        return code

    def kill(self) -> None:
        """Stop the adapter during cleanup."""
        if self.process.poll() is None:
            self.process.kill()
        self.process.wait(timeout=10)
        # A paused reader must wake to see EOF and finish.
        self._reading.set()
        self._join()

    def _join(self) -> None:
        for thread in self._threads:
            thread.join(timeout=2)