# DOCS_REF:
# - docs/user/07-stateless-op-scopes.md
# DOCS_SYNC:
# - If stateless scope resolution semantics change (global/serialized/local),
#   update stateless-op-scopes docs and examples/docs/stateless_op_scopes_example.py.

from __future__ import annotations

import importlib
import uuid
from dataclasses import dataclass
from typing import Any, Callable

import cloudpickle

from flownet.api.actor_method_ref import ActorMethodRef
from flownet.core.stateless import StatelessOpConfig, StatelessOpRef
from flownet.core.request import ActorCall, FlownetPacket
from flownet.core.types import PayloadData
from flownet.runtime.kernel import get_runtime_kernel


@dataclass
class _OpEntry:
    fn: Callable[..., Any]
    config: StatelessOpConfig
    name: str | None = None
    module: str | None = None
    qualname: str | None = None
    blob: bytes | None = None


class StatelessOpRegistry:
    def __init__(self) -> None:
        self._ops: dict[str, _OpEntry] = {}

    def register(
        self,
        fn: Callable[..., Any],
        *,
        op_id: str,
        config: StatelessOpConfig,
        name: str | None = None,
    ) -> StatelessOpRef:
        module = getattr(fn, "__module__", None)
        qualname = getattr(fn, "__qualname__", None)
        blob: bytes | None = None
        if config.scope == "serialized":
            blob = cloudpickle.dumps(fn)
        entry = _OpEntry(
            fn=fn,
            config=config,
            name=name,
            module=module,
            qualname=qualname,
            blob=blob,
        )
        self._ops[op_id] = entry
        origin = _local_address()
        return StatelessOpRef(
            op_id=op_id,
            scope=config.scope,
            name=name,
            module=module,
            qualname=qualname,
            origin=origin,
            config=config,
        )

    def get_blob(self, op_id: str) -> bytes | None:
        entry = self._ops.get(op_id)
        if entry is None:
            return None
        return entry.blob

    async def resolve(self, op_ref: StatelessOpRef) -> Callable[..., Any]:
        if op_ref.scope == "global":
            cached = self._ops.get(op_ref.op_id)
            if cached is not None:
                return cached.fn
            fn = _import_global(op_ref.module, op_ref.qualname)
            self._ops[op_ref.op_id] = _OpEntry(
                fn=fn,
                config=op_ref.config or StatelessOpConfig(scope="global"),
                name=op_ref.name,
                module=op_ref.module,
                qualname=op_ref.qualname,
            )
            return fn

        if op_ref.scope == "local":
            entry = self._ops.get(op_ref.op_id)
            if entry is None:
                raise RuntimeError(f"Local stateless op not found: {op_ref.op_id}")
            return entry.fn

        entry = self._ops.get(op_ref.op_id)
        if entry is not None:
            return entry.fn

        blob = None
        entry = self._ops.get(op_ref.op_id)
        if entry is not None:
            blob = entry.blob
        if blob is None:
            blob = await _pull_serialized(op_ref)
        if blob is None:
            raise RuntimeError(f"Serialized stateless op not found: {op_ref.op_id}")
        fn = cloudpickle.loads(blob)
        self._ops[op_ref.op_id] = _OpEntry(
            fn=fn,
            config=op_ref.config or StatelessOpConfig(scope=op_ref.scope),
            name=op_ref.name,
            module=op_ref.module,
            qualname=op_ref.qualname,
            blob=blob,
        )
        return fn


_registry = StatelessOpRegistry()


def register_stateless_op(
    fn: Callable[..., Any],
    *,
    op_id: str | None = None,
    config: StatelessOpConfig,
    name: str | None = None,
) -> StatelessOpRef:
    if op_id is None:
        op_id = uuid.uuid4().hex
    return _registry.register(fn, op_id=op_id, config=config, name=name)


def get_serialized_op(op_id: str) -> bytes | None:
    return _registry.get_blob(op_id)


async def resolve_stateless_op(op_ref: StatelessOpRef) -> Callable[..., Any]:
    return await _registry.resolve(op_ref)


def _local_address() -> str:
    return get_runtime_kernel().local_address()


def _import_global(module: str | None, qualname: str | None) -> Callable[..., Any]:
    if not module or not qualname:
        raise RuntimeError("Global stateless op missing module/qualname")
    mod = importlib.import_module(module)
    obj: Any = mod
    for attr in qualname.split("."):
        obj = getattr(obj, attr)
    if not callable(obj):
        raise RuntimeError("Global stateless op target is not callable")
    return obj


async def _pull_serialized(op_ref: StatelessOpRef) -> bytes | None:
    if not op_ref.origin:
        raise RuntimeError("Serialized stateless op missing origin")
    hub = get_runtime_kernel()
    packet = FlownetPacket(
        type="actorcall",
        uuid=uuid.uuid4().hex,
        reply=True,
        packet_content=ActorCall(
            source=hub.local_address(),
            destination=ActorMethodRef(
                address=op_ref.origin,
                actor_id="__container__",
                method="get_stateless_op_blob",
            ),
            payload=PayloadData(args=[op_ref.op_id], kwargs={}),
        ),
    )
    result = await hub.require_io().send_packet(op_ref.origin, packet)
    return result


__all__ = [
    "StatelessOpRegistry",
    "register_stateless_op",
    "resolve_stateless_op",
    "get_serialized_op",
]
