"""Declare flows and nodes in the repository you already have.

A research repository keeps its functions where they are; the decorators here
say which of them are nodes and which nodes make up a flow. ``fluksio sync``
uploads the flow document and, per node, a generated import shim — so the flow
store still holds a complete, runnable, git-versioned definition, and the code
it imports stays yours.

Nothing here imports the engine. A ``from fluksio import node`` in a training
script must not drag FastAPI and SQLAlchemy in behind it, and the same names
have to resolve inside a worker, where ``fluksio`` is the reporter module.
"""

from __future__ import annotations

import inspect
import json
import re
import warnings
from collections.abc import Callable, Sequence
from pathlib import Path
from typing import Any, TypeVar

__all__ = [
    "FLOWS",
    "Flow",
    "Port",
    "SyncError",
    "build",
    "node",
    "shims",
    "use",
]

F = TypeVar("F", bound=Callable[..., Any])

#: Mirrors :class:`fluksio.flow.messages.DType`. Mirrored rather than imported:
#: importing it would pull the engine into a research process.
DTYPES = frozenset(
    {"float", "int", "str", "bool", "json", "series", "record", "list", "artifact"}
)
#: What a list may hold, mirroring ``messages._ITEM_TYPES``.
ITEM_TYPES = frozenset({"float", "int", "str", "bool", "json", "record"})
#: Settings the engine reads itself, mirroring ``nodes.base.RESERVED_SETTINGS``.
RESERVED_SETTINGS = frozenset({"synchronous"})
#: Mirrors ``fluksio.flow.schemas.NAME_PATTERN``.
NAME_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$")
#: First line of every generated node body. Its absence is how ``sync`` knows a
#: node's source was edited on the canvas.
MARKER = "# generated by fluksio sync"


class SyncError(Exception):
    """A declaration the engine would refuse, caught where it was written."""


#: Every flow declared in the modules imported so far. ``fluksio sync`` reads it.
FLOWS: dict[str, Flow] = {}


def _name(value: str, what: str) -> str:
    if not NAME_PATTERN.match(value):
        raise SyncError(
            f"{what} '{value}': use lowercase letters, digits and underscores, "
            "starting with a letter"
        )
    return value


class Port:
    """One port of a node, and the message it binds to.

    A thin mirror of :class:`fluksio.flow.messages.MessageSpec`. ``name`` is
    the wiring — two nodes are connected because one provides the name another
    requires — and ``port`` is what the function calls its parameter, which
    defaults to the name.
    """

    def __init__(
        self,
        name: str,
        dtype: str | None = None,
        *,
        port: str = "",
        stream: bool = False,
        initial: Any = None,
        item: str | None = None,
        interval: float = 0.0,
        trigger: bool = True,
    ) -> None:
        #: Whether the caller said the type or left it to be resolved from
        #: whoever provides the message. A message has one type, declared where
        #: it is produced; requiring it does not mean declaring it again.
        self.explicit = dtype is not None
        dtype = dtype or "float"
        if dtype not in DTYPES:
            raise SyncError(f"port '{name}': unknown type '{dtype}'")
        if item is not None and item not in ITEM_TYPES:
            raise SyncError(f"port '{name}': a list cannot hold '{item}' items")
        if interval < 0:
            raise SyncError(f"port '{name}': interval cannot be negative")
        self.name = name
        self.dtype = dtype
        self.port = port or name.rsplit(".", 1)[-1]
        self.stream = stream
        self.initial = initial
        self.item = item
        self.interval = interval
        self.trigger = trigger
        if not self.port.isidentifier():
            raise SyncError(
                f"port '{name}': '{self.port}' is not a Python identifier, so no "
                "parameter can carry it — pass port='...'"
            )

    def spec(self) -> dict[str, Any]:
        """The ``MessageSpec`` this port serialises to."""
        out: dict[str, Any] = {
            "name": self.name,
            "port": self.port,
            "dtype": self.dtype,
            "interval": self.interval,
            "trigger": self.trigger,
            "stream": self.stream,
        }
        if self.item is not None:
            out["item"] = self.item
        return out

    def replace(self, name: str = "") -> Port:
        """A copy of this port, optionally bound to a different message."""
        copy = Port(
            name or self.name,
            self.dtype,
            port=self.port,
            stream=self.stream,
            initial=self.initial,
            item=self.item,
            interval=self.interval,
            trigger=self.trigger,
        )
        copy.explicit = self.explicit
        return copy

    def __repr__(self) -> str:
        return f"Port({self.name!r}, {self.dtype!r})"


def _ports(value: Sequence[str | Port] | str | Port) -> tuple[list[Port], bool]:
    """Normalise a ports declaration; the flag says it was a single port.

    A single ``Port`` rather than a list of one is the opt-in to a bare return:
    the shim wraps the value in the dict the engine expects. It has to be
    explicit, because an artifact reference is itself a dict and a dict whose
    keys match no port is dropped in silence.
    """
    if isinstance(value, (str, Port)):
        return [value if isinstance(value, Port) else Port(value)], True
    return [item if isinstance(item, Port) else Port(item) for item in value], False


class NodeSpec:
    """What ``@node`` recorded about a function."""

    def __init__(
        self,
        fn: Callable[..., Any],
        *,
        id: str,
        requires: list[Port],
        provides: list[Port],
        single: bool,
        settings: dict[str, Any],
        title: str,
        timeout: float | None,
        device: str | None,
        device_policy: str,
        cache: bool = True,
    ) -> None:
        self.fn = fn
        self.id = id
        self.requires = requires
        self.provides = provides
        self.single = single
        self.settings = settings
        self.title = title
        self.timeout = timeout
        self.device = device
        self.device_policy = device_policy
        self.cache = cache

    def rebind(
        self, *, id: str = "", wire: dict[str, str] | None = None, **settings: Any
    ) -> NodeSpec:
        """A copy of this node wired and configured for one flow."""
        wire = wire or {}
        known = {port.port for port in self.requires} | {p.port for p in self.provides}
        unknown = sorted(set(wire) - known)
        if unknown:
            raise SyncError(
                f"node '{self.id}': wire={{'{unknown[0]}': ...}} names no port of it "
                f"(it has {', '.join(sorted(known)) or 'none'})"
            )

        # Copied rather than shared: the same decorated function appears in
        # several flows, and each resolves its own port types.
        def rewire(ports: list[Port]) -> list[Port]:
            return [port.replace(wire.get(port.port, "")) for port in ports]

        return NodeSpec(
            self.fn,
            id=_name(id, "node id") if id else self.id,
            requires=rewire(self.requires),
            provides=rewire(self.provides),
            single=self.single,
            settings={**self.settings, **settings},
            title=self.title,
            timeout=self.timeout,
            device=self.device,
            device_policy=self.device_policy,
            cache=self.cache,
        )


class NodeUse:
    """One use of a decorated function in one flow — see :func:`use`."""

    def __init__(
        self,
        fn: Callable[..., Any],
        *,
        id: str = "",
        wire: dict[str, str] | None = None,
        settings: dict[str, Any] | None = None,
    ) -> None:
        self.fn = fn
        self.id = id
        self.wire = wire
        self.settings = settings or {}


def node(
    *,
    requires: Sequence[str | Port] = (),
    provides: Sequence[str | Port] | str | Port = (),
    id: str = "",
    settings: dict[str, Any] | None = None,
    title: str = "",
    timeout: float | None = None,
    device: str | None = None,
    device_policy: str = "require",
    cache: bool = True,
) -> Callable[[F], F]:
    """Mark a function as a node, declaring its ports.

    The function is returned untouched: it is still callable, testable and
    importable as what it was. A bare string in ``requires`` or ``provides`` is
    shorthand for a ``float`` port of that name.

    Parameters with a default that are not ports become the node's settings, so
    the canvas can tune them without editing code. ``device`` picks the worker
    the code runs on — ``"gpu"``, say — and ``device_policy="prefer"`` runs it
    locally when no such worker is attached rather than waiting for one.

    A batch run skips this node when an earlier one already ran the same source
    with the same settings and the same input values, and restores what it
    returned. ``cache=False`` says not to: the answer can change on its own.
    """
    if device_policy not in ("require", "prefer"):
        raise SyncError("device_policy is 'require' or 'prefer'")

    def decorate(fn: F) -> F:
        required, _ = _ports(requires)
        provided, single = _ports(provides)
        spec = NodeSpec(
            fn,
            id=_name(id or fn.__name__, "node id"),
            requires=required,
            provides=provided,
            single=single,
            settings=dict(settings or {}),
            title=title,
            timeout=timeout,
            device=device,
            device_policy=device_policy,
            cache=cache,
        )
        _check_signature(spec)
        fn.__fluksio__ = spec  # type: ignore[attr-defined]
        return fn

    return decorate


def use(
    fn: Callable[..., Any],
    *,
    id: str = "",
    wire: dict[str, str] | None = None,
    **settings: Any,
) -> NodeUse:
    """The same node, wired or configured differently in this flow.

    ``wire={"dataset": "augmented"}`` binds the ``dataset`` port to another
    message; keyword arguments override settings. Give ``id`` when one function
    appears twice in one flow.
    """
    return NodeUse(fn, id=id, wire=wire, settings=settings)


def _spec_of(entry: Callable[..., Any] | NodeUse) -> NodeSpec:
    fn = entry.fn if isinstance(entry, NodeUse) else entry
    spec = getattr(fn, "__fluksio__", None)
    if not isinstance(spec, NodeSpec):
        raise SyncError(
            f"{getattr(fn, '__name__', fn)!r} is not a node — decorate it with @node()"
        )
    if isinstance(entry, NodeUse):
        return spec.rebind(id=entry.id, wire=entry.wire, **entry.settings)
    return spec.rebind()


def _check_signature(spec: NodeSpec) -> None:
    """Whether the declaration matches the function it is on.

    Raised where the decorator is rather than where the flow is: this is a
    disagreement between a function and its own ports, and the engine would
    otherwise only find it when the node is first called.
    """
    fn, where = spec.fn, f"node '{spec.id}'"
    if inspect.iscoroutinefunction(fn):
        raise SyncError(f"{where}: async functions cannot be nodes")
    if fn.__module__ == "__main__":
        raise SyncError(
            f"{where}: {fn.__name__}() is defined in a script run directly, so the "
            "generated node could not import it — put it in an importable module"
        )
    generator = inspect.isgeneratorfunction(fn)
    if generator and spec.single:
        raise SyncError(
            f"{where}: {fn.__name__}() yields, so it says which port each value is "
            f"on — pass provides=[Port(...)] rather than a single port"
        )

    signature = inspect.signature(fn)
    parameters = signature.parameters
    var_keyword = any(
        p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()
    )
    ports = {port.port for port in spec.requires}

    for port in ports:
        if port not in parameters and not var_keyword:
            raise SyncError(
                f"{where}: port '{port}' has no parameter to arrive in — "
                f"{fn.__name__}({', '.join(parameters)}) takes no such argument"
            )
    for name, parameter in parameters.items():
        if parameter.kind is inspect.Parameter.VAR_POSITIONAL:
            raise SyncError(f"{where}: *{name} cannot be filled from ports")
        if parameter.kind is inspect.Parameter.VAR_KEYWORD or name in ports:
            continue
        if parameter.default is inspect.Parameter.empty:
            raise SyncError(
                f"{where}: parameter '{name}' is neither a port nor a setting — "
                f"add Port('{name}', ...) to requires=, or give it a default"
            )


def _check(spec: NodeSpec, mode: str) -> dict[str, Any]:
    """The node as stored, with the settings and wiring of one flow.

    Mirrors ``FlowController._build_node`` and ``flow.runs.batch_issues``. Kept
    apart from the signature check because ``use(fn, epochs=3)`` adds settings
    per flow, so this is the first point at which they are all known.
    """
    fn, where = spec.fn, f"node '{spec.id}'"
    parameters = inspect.signature(fn).parameters
    var_keyword = any(
        p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values()
    )
    ports = {port.port for port in spec.requires}

    params: dict[str, Any] = {}
    for name, parameter in parameters.items():
        if name in ports or parameter.default is inspect.Parameter.empty:
            continue
        if name in RESERVED_SETTINGS:
            continue
        if _serialisable(parameter.default):
            params[name] = parameter.default
    for name, value in spec.settings.items():
        if name in ports:
            raise SyncError(
                f"{where}: '{name}' is both a port and a setting — a node reads its "
                "settings beside its ports, so the names have to differ"
            )
        if name in RESERVED_SETTINGS:
            raise SyncError(f"{where}: '{name}' is a setting the engine reads itself")
        if name not in parameters and not var_keyword:
            raise SyncError(
                f"{where}: setting '{name}' is not a parameter of {fn.__name__}()"
            )
        if not _serialisable(value):
            raise SyncError(
                f"{where}: setting '{name}' is not JSON — settings are stored with "
                "the flow, so they have to survive a round trip"
            )
        params[name] = value

    if mode == "batch":
        for port in spec.requires + spec.provides:
            if port.interval > 0 and not port.stream:
                raise SyncError(
                    f"{where}: port '{port.port}' has interval={port.interval} but a "
                    "batch run delivers once — drop it, or declare stream=True"
                )

    return {
        "id": spec.id,
        "type": "python",
        "title": spec.title,
        "params": params,
        "requires": [port.spec() for port in spec.requires],
        "provides": [port.spec() for port in spec.provides],
        "timeout": spec.timeout,
        "device": spec.device,
        "device_policy": spec.device_policy,
        "cache": spec.cache,
    }


def _serialisable(value: Any) -> bool:
    try:
        json.dumps(value)
    except (TypeError, ValueError):
        return False
    return True


class Flow:
    """A flow: which nodes are in it, what it takes and what it reports."""

    def __init__(
        self,
        name: str,
        *,
        nodes: Sequence[Callable[..., Any] | NodeUse],
        mode: str = "batch",
        inputs: Sequence[Port] = (),
        outputs: Sequence[str] = (),
        title: str = "",
    ) -> None:
        """Declare a flow over decorated functions, and register it for sync.

        ``nodes`` is the membership: any decorated function, from anywhere in
        the repository. Wiring is not membership — nodes connect because one
        provides a message another requires, never because one imported the
        other. Wrap an entry in :func:`use` to rebind or reconfigure it here
        without touching the other flows using it.
        """
        if mode not in ("batch", "live"):
            raise SyncError(f"flow '{name}': mode is 'batch' or 'live'")
        self.name = _name(name, "flow name")
        self.title = title
        self.mode = mode
        self.inputs = list(inputs)
        self.outputs = list(outputs)
        self.nodes = [_spec_of(entry) for entry in nodes]
        self._resolve_types()
        self._validate()
        self._nodes = [_check(spec, mode) for spec in self.nodes]
        self._register()

    def _register(self) -> None:
        """Put this flow where `fluksio sync` will find it."""
        existing = FLOWS.get(self.name)
        if existing is not None and existing is not self:
            warnings.warn(
                f"two flows are called '{self.name}'; the later one wins, and only "
                "it will be synced",
                stacklevel=3,
            )
        FLOWS[self.name] = self

    def _resolve_types(self) -> None:
        """Give every required port the type of the message it reads.

        A message has one type, declared once by whoever produces it. Requiring
        it is not a second declaration, so `requires=["dataset"]` takes the
        artifact type from the node providing `dataset` rather than defaulting
        to a float the engine would then refuse at run time. Spelling the type
        out on both sides is still allowed — and a disagreement is an error,
        because one of the two is wrong about what it is handling.
        """
        declared: dict[str, Port] = {}
        for port in self.inputs:
            declared.setdefault(port.name, port)
        for spec in self.nodes:
            for port in spec.provides:
                declared.setdefault(port.name, port)
        for spec in self.nodes:
            for port in spec.requires:
                source = declared.get(port.name)
                if source is None:
                    continue
                if not port.explicit:
                    port.dtype = source.dtype
                    port.item = source.item
                elif port.dtype != source.dtype:
                    raise SyncError(
                        f"flow '{self.name}': node '{spec.id}' reads '{port.name}' as "
                        f"{port.dtype}, but it is provided as {source.dtype}"
                    )

    def _validate(self) -> None:
        seen: set[str] = set()
        for spec in self.nodes:
            if spec.id in seen:
                raise SyncError(
                    f"flow '{self.name}': two nodes are called '{spec.id}' — "
                    "pass use(fn, id='...') to tell them apart"
                )
            seen.add(spec.id)

        produced: dict[str, list[str]] = {}
        for spec in self.nodes:
            for port in spec.provides:
                produced.setdefault(port.name, []).append(spec.id)
        for name, writers in produced.items():
            if len(writers) > 1:
                warnings.warn(
                    f"flow '{self.name}': {', '.join(writers)} all publish "
                    f"'{name}', so a reader gets whichever ran last",
                    stacklevel=4,
                )

        available = set(produced) | {port.name for port in self.inputs}
        for spec in self.nodes:
            for port in spec.requires:
                if "." in port.name or port.name in available:
                    continue
                raise SyncError(
                    f"flow '{self.name}': node '{spec.id}' requires '{port.name}', "
                    "which no node in it provides and no input declares"
                )
        for output in self.outputs:
            if output not in available:
                raise SyncError(
                    f"flow '{self.name}': output '{output}' is not produced by any "
                    "of its nodes"
                )

    def document(self, origin: dict[str, Any] | None = None) -> dict[str, Any]:
        """The flow document to store, without a version — sync sets that."""
        doc: dict[str, Any] = {
            "name": self.name,
            "title": self.title,
            "mode": self.mode,
            "nodes": self._nodes,
            "inputs": [
                {"spec": port.spec(), "initial": port.initial} for port in self.inputs
            ],
            "outputs": self.outputs,
        }
        if origin is not None:
            doc["origin"] = origin
        return doc

    def shims(self) -> dict[str, str]:
        """One generated node body per node, keyed by node id."""
        return {spec.id: _shim(spec) for spec in self.nodes}

    def submit(
        self, *, seed: int | None = None, client: Any = None, **params: Any
    ) -> Any:
        """Start a run of this flow and hand back a handle to it."""
        from fluksio.sdk.client import Client

        return (client or Client()).submit(self.name, params, seed=seed)

    def runs(self, limit: int = 20, client: Any = None) -> list[dict[str, Any]]:
        """This flow's runs, newest first."""
        from fluksio.sdk.client import Client

        rows: list[dict[str, Any]] = (client or Client()).runs(
            flow=self.name, limit=limit
        )
        return rows

    def __repr__(self) -> str:
        return f"Flow({self.name!r}, nodes={[n.id for n in self.nodes]})"


def build(target: Flow, origin: dict[str, Any] | None = None) -> dict[str, Any]:
    """The flow document for a flow — checked, without a version."""
    return target.document(origin)


def shims(target: Flow) -> dict[str, str]:
    """The generated node bodies for a flow."""
    return target.shims()


def import_root(fn: Callable[..., Any]) -> str:
    """The directory that has to be on the path for ``fn`` to be importable.

    Its module's dotted name says how deep in a package it sits, so counting
    the segments back up from the file lands on the directory the import
    resolves against — which is not always the repository root, and is what a
    generated shim actually needs.
    """
    source = inspect.getsourcefile(fn)
    if not source:
        return ""
    path = Path(source).resolve()
    for _ in fn.__module__.split("."):
        path = path.parent
    return str(path)


def _shim(spec: NodeSpec) -> str:
    """The node body the store keeps: an import of the real function.

    The store therefore still holds a complete, runnable definition — the body
    simply happens to be generated, which is why it says so and says where the
    real thing is.
    """
    fn = spec.fn
    where = inspect.getsourcefile(fn) or fn.__module__
    repo = import_root(fn)
    ports = [port.port for port in spec.requires]
    signature = ", ".join(ports + ["**settings"])
    arguments = ", ".join([f"{port}={port}" for port in ports] + ["**settings"])
    call = f"{fn.__name__}({arguments})"

    lines = [
        f"{MARKER} from {where} — edit that file instead",
        "import sys",
        "",
    ]
    if repo:
        # ponytail: the repository is on the path because the shim puts it
        # there. An editable install (`-e /repo` in the Modules manifest, or a
        # VCS requirement on a remote worker) makes these three lines a no-op.
        lines += [
            f"_REPO = {repo!r}",
            "if _REPO not in sys.path:",
            "    sys.path.insert(0, _REPO)",
            "",
        ]
    lines += [
        f"from {fn.__module__} import {fn.__name__}",
        "",
        "",
        f"def process({signature}):",
    ]
    if inspect.isgeneratorfunction(fn):
        # `return (yield from ...)`, not a bare `yield from`: the delegating
        # form is what carries the generator's own return value out, and that
        # return value is everything the node produces at the end.
        lines.append(f"    return (yield from {call})")
    elif spec.single:
        lines.append(f"    return {{{spec.provides[0].port!r}: {call}}}")
    else:
        lines.append(f"    return {call}")
    return "\n".join(lines) + "\n"