"""One user-code process: compiles a node's source once, then runs it on demand.

This runs under the *user* venv's interpreter, so a node's imports resolve
against what the Modules page installed rather than against the engine's own
packages. Nothing of the engine is importable here — the file is handed to the
interpreter by path and is deliberately standard library only.

The protocol is one JSON object per line: a request arrives on stdin, the reply
goes out on a private duplicate of fd 1 taken before anything else can write to
it. fd 1 itself is pointed at stderr, so a stray ``write(1, ...)`` — from a
native library, or a node printing during an import — lands in the server log
instead of corrupting the reply stream.

A node may also send lines back *while* it is still running: anything carrying
an ``event`` key is a report rather than the answer, and the engine keeps
reading. That is what makes a training curve visible during the hours it takes
to draw, and what tells the engine a long node is alive rather than hung —
each event resets its deadline, so the timeout measures silence rather than
duration.

What travels that way is not a log. A node that produces values over time is a
generator: every ``yield`` is a dict keyed by output port, published the moment
it happens, and whatever the generator returns at the end is the node's result.
Nothing leaves a node except through a port it declared, which is the whole
point — a number worth keeping is an output, not a side effect. Where a yield
cannot reach, because the value comes from inside somebody else's callback,
``fluksio.emit(loss=0.3)`` writes the same ports the same way.
"""

from __future__ import annotations

import os
import sys

# Handed to the interpreter by path, so CPython puts this file's own directory
# — ``fluksio_worker`` — on ``sys.path[0]``. A node doing ``import
# queue`` would then get the engine's queue module rather than the standard
# library's. Drop it before anything else can import. (``-P`` and
# ``PYTHONSAFEPATH`` do this at startup, but both are 3.11+.)
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path[:] = [p for p in sys.path if os.path.abspath(p or ".") != _HERE]

import contextlib
import hashlib
import inspect
import io
import json
import logging
import tempfile
import time
import traceback
import urllib.parse
import urllib.request
from collections.abc import Callable
from types import ModuleType
from typing import Any, cast

#: Enough of a print to debug with. A node printing in a loop must not fill the
#: pipe or the reply.
MAX_LOG = 16 * 1024

#: Where the artifact store is, from this worker's point of view: a directory
#: when it shares the engine's filesystem, a URL when it does not.
ARTIFACT_DIR_ENV = "FLUKSIO_ARTIFACT_DIR"
#: The ring, for frames a flow only shows live. Same layout as the store above;
#: what differs is that the oldest fall out of it.
ARTIFACT_VOLATILE_DIR_ENV = "FLUKSIO_ARTIFACT_VOLATILE_DIR"
ARTIFACT_URL_ENV = "FLUKSIO_ARTIFACT_URL"
ARTIFACT_TOKEN_ENV = "FLUKSIO_ARTIFACT_TOKEN"
ARTIFACT_CACHE_ENV = "FLUKSIO_ARTIFACT_CACHE"
#: How much a worker that fetches over HTTP keeps of what it downloaded. The
#: cache is keyed by digest and nothing ever invalidates an entry, so without
#: a bound a node reading a media stream fills the disk with chunks nothing
#: will ask for twice.
ARTIFACT_CACHE_BYTES_ENV = "FLUKSIO_ARTIFACT_CACHE_BYTES"
DEFAULT_CACHE_BYTES = 1024 * 1024 * 1024
ARTIFACT_TIMEOUT_S = 300
#: How much of an artifact is held in memory at a time, matching the engine's
#: own store. Repeated rather than imported: nothing of the engine is here.
CHUNK = 1024 * 1024

#: The reply channel, opened by ``main``. Also what an event line goes down.
_RPC: Any = None
#: The call being served, so an event can say which one it belongs to.
_CALL_ID = ""


def _emit(event: dict[str, Any]) -> None:
    """Send one line back without ending the call."""
    if _RPC is None:
        return
    event["call_id"] = _CALL_ID
    event["ts"] = time.time()
    _RPC.write(json.dumps(event) + "\n")
    _RPC.flush()


class _Reporter(ModuleType):
    """``import fluksio`` — the parts of a node's job that need the engine."""

    #: The same name the SDK exports at top level, so a node written against
    #: one surface finds it on the other. Set by ``_install_reporter``.
    logger: logging.Logger

    def emit(self, **ports: Any) -> None:
        """Publish on this node's output ports without returning yet.

        Yielding is the better way to say this and should be preferred; use
        this where a yield cannot reach — inside a training framework's
        callback, say, which calls you rather than the other way round. It is
        the same publication either way: the values go to the ports the node
        declared, and are checked against them.
        """
        _emit({"event": "emit", "outputs": dict(ports)})

    def save_artifact(
        self,
        source: Any,
        name: str = "",
        media_type: str = "application/octet-stream",
        volatile: bool = False,
    ) -> dict[str, Any]:
        """Store bytes or a file and return the reference to return onward.

        Return the reference from an ``artifact`` port: a checkpoint is far too
        big to be a message, and the reference is what the next node opens.

        ``volatile`` is for a frame rather than a result: the bytes go to a
        ring held in memory and are pushed to whatever screen is watching,
        instead of being written to the data volume. They last seconds. A
        volatile reference a node *returns* from a run is kept anyway, so what
        this costs is only what nobody asked to keep.
        """
        if isinstance(source, (bytes, bytearray)):
            return _store_bytes(
                bytes(source), name or "artifact.bin", media_type, volatile
            )
        path = str(source)
        return _store_file(
            path, name or os.path.basename(path), media_type, volatile
        )

    def load_artifact(self, ref: dict[str, Any]) -> str:
        """Fetch an artifact and hand back a local path to read it from."""
        digest = str((ref or {}).get("digest") or "")
        if not digest.startswith("sha256:"):
            raise ValueError("not an artifact reference")
        return _fetch(digest)

    # ---------------------------------------------------------------------
    # The authoring API, inert.
    #
    # A node generated by `fluksio sync` imports the caller's own module, and
    # that module says `from fluksio import Port, node, Flow` at the top —
    # which, in here, is this. The declarations were read at sync time and are
    # already in the flow document, so what they have to do now is import
    # without doing anything: the decorators hand the function back, and `flow`
    # builds nothing.
    # ---------------------------------------------------------------------

    def Port(self, *args: Any, **kwargs: Any) -> Any:  # noqa: N802
        """A port declaration, already read by `fluksio sync`."""
        return _Declared()

    def node(self, *args: Any, **kwargs: Any) -> Any:
        """The decorator, which here gives the function straight back."""
        return lambda fn: fn

    def use(self, *args: Any, **kwargs: Any) -> Any:
        """One use of a node in a flow, already read by `fluksio sync`."""
        return _Declared()

    def Flow(self, *args: Any, **kwargs: Any) -> Any:  # noqa: N802
        """A flow declaration, already read by `fluksio sync`."""
        return _Declared()


class _LogHandler(logging.Handler):
    """Writes to whatever ``sys.stderr`` is at the time of the record.

    A call runs under ``redirect_stderr`` into the buffer that becomes the
    node's logs, so a handler holding the stream it was built with would write
    past every one of them.
    """

    def emit(self, record: logging.LogRecord) -> None:
        sys.stderr.write(self.format(record) + "\n")


class _Declared:
    """Stands in for a declaration whose work was done before the run.

    Tolerant on purpose: a module may keep one at module level and touch it in
    ways a node never exercises, and none of that should fail an import.
    """

    def __getattr__(self, name: str) -> Any:
        if name.startswith("__"):
            raise AttributeError(name)
        return _Declared()

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        return _Declared()

    def __repr__(self) -> str:
        return "<fluksio declaration>"


def _artifact_dir(volatile: bool) -> str | None:
    """Which store this worker writes to, if it can see one at all.

    A volatile write falls back to the durable directory rather than failing:
    an engine with its ring turned off still has to take the frame.
    """
    if volatile:
        ring = os.environ.get(ARTIFACT_VOLATILE_DIR_ENV)
        if ring:
            return ring
    return os.environ.get(ARTIFACT_DIR_ENV)


def _store_bytes(
    data: bytes, name: str, media_type: str, volatile: bool = False
) -> dict[str, Any]:
    """Write to the artifact store, whichever end of it this worker can see.

    A worker in the engine's own container writes the file; one on another host
    puts it over HTTP. Node code cannot tell the difference, which is the point
    — the same flow runs either place. The hashing is repeated from
    ``fluksio.flow.artifacts`` rather than imported, because nothing of the engine
    is importable here.
    """
    digest = "sha256:" + hashlib.sha256(data).hexdigest()
    directory = _artifact_dir(volatile)
    if directory:
        target = os.path.join(directory, digest[7:9], digest[7:])
        if not os.path.exists(target):
            os.makedirs(os.path.dirname(target), exist_ok=True)
            # Written beside the target and moved, so a reader never opens a
            # half-written artifact.
            handle, temporary = tempfile.mkstemp(dir=os.path.dirname(target))
            with os.fdopen(handle, "wb") as out:
                out.write(data)
            os.replace(temporary, target)
    else:
        _put_over_http(data, name, media_type, volatile=volatile)
    return {
        "digest": digest,
        "size": len(data),
        "media_type": media_type,
        "name": name,
    }


def _store_file(
    path: str, name: str, media_type: str, volatile: bool = False
) -> dict[str, Any]:
    """The same, for a file already on disk — read a chunk at a time.

    A video segment or a checkpoint is as likely to be handed over as a path as
    as bytes, and reading one into memory to hash it defeats the point of it
    being a file.
    """
    digester = hashlib.sha256()
    size = 0
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(CHUNK), b""):
            digester.update(chunk)
            size += len(chunk)
    digest = "sha256:" + digester.hexdigest()

    directory = _artifact_dir(volatile)
    if directory:
        target = os.path.join(directory, digest[7:9], digest[7:])
        if not os.path.exists(target):
            os.makedirs(os.path.dirname(target), exist_ok=True)
            handle_fd, temporary = tempfile.mkstemp(dir=os.path.dirname(target))
            with os.fdopen(handle_fd, "wb") as out, open(path, "rb") as source:
                for chunk in iter(lambda: source.read(CHUNK), b""):
                    out.write(chunk)
            os.replace(temporary, target)
    else:
        with open(path, "rb") as handle:
            _put_over_http(handle, name, media_type, size, volatile=volatile)
    return {"digest": digest, "size": size, "media_type": media_type, "name": name}


def _put_over_http(
    data: Any,
    name: str,
    media_type: str,
    length: int | None = None,
    volatile: bool = False,
) -> None:
    base = os.environ.get(ARTIFACT_URL_ENV)
    if not base:
        raise RuntimeError(
            "this worker has no artifact store — neither "
            f"{ARTIFACT_DIR_ENV} nor {ARTIFACT_URL_ENV} is set"
        )
    fields = {"name": name, "media_type": media_type}
    if volatile:
        fields["volatile"] = "1"
    query = urllib.parse.urlencode(fields)
    request = urllib.request.Request(
        f"{base.rstrip('/')}?{query}", data=data, method="PUT"
    )
    if length is not None:
        # urllib streams a file object only when it does not have to measure it.
        request.add_header("Content-Length", str(length))
    _authorize(request)
    with urllib.request.urlopen(request, timeout=ARTIFACT_TIMEOUT_S):
        pass


def _fetch(digest: str) -> str:
    """A local path holding this artifact's bytes, downloading it if needed."""
    for variable in (ARTIFACT_DIR_ENV, ARTIFACT_VOLATILE_DIR_ENV):
        directory = os.environ.get(variable)
        if not directory:
            continue
        target = os.path.join(directory, digest[7:9], digest[7:])
        if os.path.exists(target):
            return target
    if os.environ.get(ARTIFACT_DIR_ENV):
        # A worker sharing the engine's filesystem has nowhere else to look:
        # either the bytes are in one of the two stores or they have gone.
        raise FileNotFoundError(digest)

    # Cached by digest: content-addressing means a file once fetched is never
    # stale, so a sweep pulls a shared input across the network once.
    cache = os.environ.get(ARTIFACT_CACHE_ENV) or os.path.join(
        tempfile.gettempdir(), "fluksio-artifacts"
    )
    target = os.path.join(cache, digest[7:])
    if os.path.exists(target):
        return target
    base = os.environ.get(ARTIFACT_URL_ENV)
    if not base:
        raise RuntimeError(f"this worker cannot reach the artifact store: {digest}")
    os.makedirs(cache, exist_ok=True)
    request = urllib.request.Request(f"{base.rstrip('/')}/{digest}")
    _authorize(request)
    handle, temporary = tempfile.mkstemp(dir=cache)
    with urllib.request.urlopen(request, timeout=ARTIFACT_TIMEOUT_S) as response:
        with os.fdopen(handle, "wb") as out:
            while chunk := response.read(CHUNK):
                out.write(chunk)
    os.replace(temporary, target)
    _trim_cache(cache)
    return target


def _trim_cache(cache: str) -> None:
    """Hold the fetch cache inside its bound, oldest first.

    Content addressing means an entry is never stale, so nothing here ever
    expires on its own — and a node reading a media stream asks for a digest it
    will never ask for again. Cheap: one listing per download.
    """
    try:
        limit = int(os.environ.get(ARTIFACT_CACHE_BYTES_ENV) or DEFAULT_CACHE_BYTES)
    except ValueError:
        limit = DEFAULT_CACHE_BYTES
    if limit <= 0:
        return
    entries = []
    total = 0
    with os.scandir(cache) as listing:
        for entry in listing:
            try:
                stat = entry.stat()
            except OSError:
                continue
            if not entry.is_file():
                continue
            entries.append((stat.st_mtime, stat.st_size, entry.path))
            total += stat.st_size
    if total <= limit:
        return
    for _mtime, size, path in sorted(entries):
        if total <= limit:
            break
        try:
            os.unlink(path)
        except OSError:
            continue
        total -= size


def _authorize(request: Any) -> None:
    token = os.environ.get(ARTIFACT_TOKEN_ENV)
    if token:
        request.add_header("Authorization", f"Bearer {token}")


def _install_reporter() -> None:
    """Put ``fluksio`` on the import path of every node this worker runs."""
    module = _Reporter("fluksio")
    module.__doc__ = "Report metrics and progress from inside a node."
    # The SDK exports a logger at top level, so a node written against it says
    # `fluksio.logger.info(...)` — which used to die with AttributeError, after
    # whatever it was reporting on had already succeeded. Its records land in
    # the same capture as a print.
    logger = logging.getLogger("fluksio")
    logger.setLevel(logging.INFO)
    logger.addHandler(_LogHandler())
    # Nothing above it should see these twice: a node calling `basicConfig` for
    # its own logging installs a root handler on the same stream.
    logger.propagate = False
    module.logger = logger
    sys.modules["fluksio"] = module


def load_function(flow: str, node_id: str, code: str) -> Callable[..., Any]:
    """Compile a node's source and return the function to run.

    A node file defines ``process(...)``; if it defines exactly one public
    function under another name, that one is used.
    """
    digest = hashlib.md5(code.encode()).hexdigest()[:8]
    module_name = f"_fluksio_node_{flow}_{node_id}_{digest}"

    module = ModuleType(module_name)
    module.__dict__["__name__"] = module_name
    sys.modules[module_name] = module
    exec(compile(code, f"<node {flow}.{node_id}>", "exec"), module.__dict__)

    if callable(getattr(module, "process", None)):
        return cast(Callable[..., Any], module.process)

    functions = [
        value
        for name, value in vars(module).items()
        if callable(value)
        and not name.startswith("_")
        and getattr(value, "__module__", None) == module_name
    ]
    if len(functions) == 1:
        return cast(Callable[..., Any], functions[0])
    raise ValueError(
        "Define a function named 'process' — this file has "
        f"{len(functions)} functions to choose from."
    )


class _Capped(io.StringIO):
    """Keeps the first ``MAX_LOG`` characters and quietly drops the rest."""

    def write(self, text: str) -> int:
        room = MAX_LOG - self.tell()
        if room > 0:
            super().write(text[:room])
        return len(text)


def _short_error(exc: BaseException) -> str:
    """One line a node author can act on — the engine's rule, applied here."""
    if isinstance(exc, SyntaxError):
        # Its own message already names the compiled file, which is noise here.
        return f"{type(exc).__name__}: {exc.msg} (line {exc.lineno})"

    frames = [
        frame
        for frame in traceback.extract_tb(exc.__traceback__)
        if frame.filename.startswith("<node ")
    ]
    where = f" (line {frames[-1].lineno})" if frames else ""
    return f"{type(exc).__name__}: {exc}{where}"


def _node_traceback(exc: BaseException) -> str:
    """The traceback from the node's own code onward; the rest is this loop."""
    frames = traceback.extract_tb(exc.__traceback__)
    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(type(exc), exc)
    )


def _handle(request: dict[str, Any], cache: dict[tuple[str, str], Any]) -> Any:
    """Compile if needed, then run — the part that may raise the user's error."""
    flow, node = request["flow"], request["node"]
    source = request.get("source") or ""
    digest = hashlib.md5(source.encode()).hexdigest()

    if request["op"] == "compile" and not request.get("keep", True):
        # A question about a draft, not a load: does this source compile? The
        # answer is all the editor wants, and keeping it would throw away the
        # published source this worker is still serving calls from.
        load_function(flow, node, source)
        return None

    cached = cache.get((flow, node))
    if cached is None or cached[0] != digest:
        cache[(flow, node)] = (digest, load_function(flow, node, source))
    function = cache[(flow, node)][1]

    if request["op"] == "compile":
        return None

    result = function(**(request.get("kwargs") or {}))
    if inspect.isgenerator(result):
        result = _drain(result)
    return result


def _drain(generator: Any) -> Any:
    """Run a generator node, sending each yield the moment it happens.

    Two shapes work, and they mean the same thing. Yield throughout and
    ``return`` the result at the end, which is the explicit one; or just yield,
    and the last one is the result. Which of the two a node is cannot be known
    until it ends, so the last yield has to be held back somewhere — and that
    somewhere is the engine, which is the only side that knows what ports the
    node declared. Holding it here instead cost a pass: a yield naming a port
    that does not exist was only checked once the *next* one arrived.
    """
    try:
        while True:
            _emit({"event": "yield", "outputs": next(generator)})
    except StopIteration as stop:
        return stop.value


def main() -> None:
    global _RPC, _CALL_ID

    _RPC = os.fdopen(os.dup(1), "w")
    # Everything the node writes to the real stdout now goes to the server log.
    os.dup2(2, 1)
    _install_reporter()

    cache: dict[tuple[str, str], Any] = {}
    for line in sys.stdin:
        if not line.strip():
            continue
        request = json.loads(line)
        _CALL_ID = str(request.get("call_id") or "")
        captured = _Capped()
        # ``id`` back untouched: the engine matches it against what it sent, so
        # a pipe that has slipped a call is caught rather than handing one node
        # another node's answer.
        response: dict[str, Any] = {"id": request.get("id"), "call_id": _CALL_ID}
        try:
            with (
                contextlib.redirect_stdout(captured),
                contextlib.redirect_stderr(captured),
            ):
                response["result"] = _handle(request, cache)
            response["ok"] = True
        except Exception as exc:
            response["ok"] = False
            response["error"] = {
                "type": type(exc).__name__,
                "message": str(exc),
                "short": _short_error(exc),
                "traceback": _node_traceback(exc),
            }
        response["logs"] = captured.getvalue()
        try:
            # allow_nan=False: a bare NaN is what json.dumps would write, and
            # nothing downstream can read it back — the ports refuse one too,
            # but this is the crossing, so refuse it where it is still the
            # node's own reply rather than a row somebody queries later.
            reply = json.dumps(response, allow_nan=False)
        except (TypeError, ValueError, RecursionError):
            # Encoding the reply is also the check that the node returned
            # something the typed-message contract can carry — the result is
            # the only part of this dict a node controls, so it is the only
            # part that can fail. Doing it once is why there is no separate
            # dumps of the result above.
            result = response.pop("result", None)
            try:
                json.dumps(result)
            except (TypeError, ValueError, RecursionError):
                trouble = (
                    f"returned {type(result).__name__}, which cannot be sent "
                    "back as JSON — return numbers, strings, booleans, lists "
                    "or dicts."
                )
            else:
                # It encodes with NaN allowed, so that is what is wrong with it.
                trouble = (
                    "returned a NaN or an infinity, which JSON cannot carry — "
                    "publish None, or a number that says the measurement had "
                    "nothing in it."
                )
            response["ok"] = False
            response["error"] = {
                "type": "ValueError",
                "message": trouble,
                "short": f"ValueError: {trouble}",
                "traceback": "",
            }
            reply = json.dumps(response)
        _RPC.write(reply + "\n")
        _RPC.flush()


if __name__ == "__main__":
    main()