"""Capture what a node prints, so the editor can show it.

Node code is written by the user and ``print`` is the obvious way to look at a
value, but a node runs on a pool thread and its output would otherwise land
unattributed in the server log. ``sys.stdout`` is replaced once by a tee that
also hands what it is given to whichever capture is active on *this* thread —
so only node executions are captured, and everything else passes through
untouched.
"""

from __future__ import annotations

import io
import sys
import traceback
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from typing import TextIO

# Set for the duration of one node execution, on the thread running it.
_sink: ContextVar[Callable[[str], None] | None] = ContextVar(
    "fluksio_node_log_sink", default=None
)


class _Tee(io.TextIOBase):
    """Writes through to the real stream, and to the active capture."""

    def __init__(self, real: TextIO) -> None:
        self._real = real

    def write(self, text: str) -> int:
        sink = _sink.get()
        if sink is not None and text:
            sink(text)
        try:
            return self._real.write(text)
        except OSError:
            # A dead stdout — `fluksio serve` runs the engine as a child of the
            # dashboard, which holds the far end of that pipe — must not fail
            # the node whose output was being teed. The capture above has it,
            # and there is nowhere left to report the loss to anyway.
            return len(text)

    def flush(self) -> None:
        try:
            self._real.flush()
        except OSError:
            pass

    def isatty(self) -> bool:
        return self._real.isatty()


def install() -> None:
    """Put the tee in place. Safe to call more than once."""
    if not isinstance(sys.stdout, _Tee):
        sys.stdout = _Tee(sys.stdout)
    if not isinstance(sys.stderr, _Tee):
        sys.stderr = _Tee(sys.stderr)


@contextmanager
def capture(sink: Callable[[str], None]) -> Iterator[None]:
    """Send everything printed on this thread to ``sink`` for the duration."""
    token = _sink.set(sink)
    try:
        yield
    finally:
        _sink.reset(token)


def node_traceback() -> str:
    """The exception being handled, from the node's own code onward.

    The frames above it are the engine calling the node, which is noise to the
    person who wrote it — the same trimming ``_short_error`` does for the one
    line shown on the node itself.
    """
    exc_type, exc, tb = sys.exc_info()
    if exc is None:
        return ""
    # A node running out of process already trimmed its own; the frames on this
    # side are the RPC that carried it.
    remote = getattr(exc, "remote_traceback", "")
    if remote:
        return str(remote)
    frames = traceback.extract_tb(tb)
    start = next(
        (i for i, frame in enumerate(frames) if frame.filename.startswith("<node ")),
        0,
    )
    return "".join(
        ["Traceback (most recent call last):\n"]
        + traceback.format_list(frames[start:])
        + traceback.format_exception_only(exc_type, exc)
    )


class Collector:
    """Gathers one execution's output, with a ceiling on how much it keeps."""

    __slots__ = ("chunks", "truncated", "_limit", "_size", "_max_bytes")

    def __init__(self, limit: int = 200, max_bytes: int = 8192) -> None:
        self.chunks: list[str] = []
        self.truncated = False
        self._limit = limit
        self._max_bytes = max_bytes
        self._size = 0

    def __call__(self, text: str) -> None:
        if len(self.chunks) >= self._limit or self._size >= self._max_bytes:
            self.truncated = True
            return
        self.chunks.append(text)
        self._size += len(text)

    @property
    def text(self) -> str:
        out = "".join(self.chunks)
        return out[: self._max_bytes] if len(out) > self._max_bytes else out