#  Copyright 2026 zhaoxi826
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""进程内 actor 运行时（位置透明的组件寻址层）。

设计目标：**一处编写，处处运行**。业务代码永远通过 ``get_actor(name)`` 拿到一个
``ActorHandle``，再 ``await handle.method(args)`` 调用，**不关心对端在哪**。

今天所有核心组件（PostgresDatabase / GlobalStateMachine / ...）都是本进程内的
普通异步对象，``get_actor`` 返回 ``LocalActorHandle``，直接转发到本地实例。

未来若要把某一层（如跑 vLLM 的重型 worker）拆到独立进程/机器上，只需新增一个
``RemoteActorHandle``（gRPC / 子进程 / 消息队列后端）并在注册时选择它——业务侧的
调用代码一行都不用改。这层接缝就是当初 Ray ``.remote()`` 提供的价值，这里用零依赖
的方式把它保留下来。
"""

from __future__ import annotations

import asyncio
from typing import Any, Awaitable, Dict, Optional, Set

_registry: Dict[str, Any] = {}
_handle_cache: Dict[str, "ActorHandle"] = {}

# 已 spawn 的后台任务：持强引用，防止事件循环只持弱引用导致任务被 GC 掉
_background_tasks: Set["asyncio.Task[Any]"] = set()


class _BoundMethod:
    """包装目标对象的一个方法，使调用统一返回 awaitable。

    - 目标方法是 ``async def`` → 直接返回其协程
    - 目标方法是普通函数 → 把返回值包成一个已完成的 awaitable
    这样调用方永远可以 ``await handle.method(...)``，无需关心同步/异步、本地/远程。
    """

    __slots__ = ("_fn",)

    def __init__(self, fn: Any) -> None:
        self._fn = fn

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        result = self._fn(*args, **kwargs)
        if asyncio.iscoroutine(result):
            return result

        async def _wrap() -> Any:
            return result

        return _wrap()


class ActorHandle:
    """actor 句柄基类：定义位置透明的调用接口。

    未来的 ``RemoteActorHandle`` 继承本类，改写方法解析为跨进程 RPC 即可。
    """


class LocalActorHandle(ActorHandle):
    """本进程句柄：方法调用直接转发到已注册的实例。"""

    __slots__ = ("_impl",)

    def __init__(self, impl: Any) -> None:
        object.__setattr__(self, "_impl", impl)

    def __getattr__(self, name: str) -> Any:
        attr = getattr(object.__getattribute__(self, "_impl"), name)
        if callable(attr):
            return _BoundMethod(attr)
        return attr


def register_actor(name: str, instance: Any) -> None:
    """注册一个 actor 实例。重复注册会覆盖（支持插件热重载）。"""
    _registry[name] = instance
    _handle_cache.pop(name, None)


def unregister_actor(name: str) -> None:
    """注销一个 actor（插件卸载用）；不存在时静默。"""
    _registry.pop(name, None)
    _handle_cache.pop(name, None)


def get_actor(name: str) -> ActorHandle:
    """按名字取 actor 句柄；未注册时抛 KeyError。

    动态名（如 ``f"org_{plugin}"``）也走这个统一入口。
    """
    if name not in _registry:
        raise KeyError(f"actor {name!r} 未注册")
    handle = _handle_cache.get(name)
    if handle is None:
        handle = LocalActorHandle(_registry[name])
        _handle_cache[name] = handle
    return handle


def try_get_actor(name: str) -> Optional[ActorHandle]:
    """按名字取 actor 句柄；未注册时返回 None（可选依赖用）。"""
    try:
        return get_actor(name)
    except KeyError:
        return None


def clear_actors() -> None:
    """清空注册表（测试用）。"""
    _registry.clear()
    _handle_cache.clear()
    _background_tasks.clear()


def spawn_background(coro: "Awaitable[Any]") -> "asyncio.Task[Any]":
    """在当前事件循环中后台调度一个协程，并维持强引用防止被 GC。

    等价于旧 Ray `.remote()` 的 fire-and-forget 语义：
    - 协程被调度执行；
    - 异常不会静默消失，而是记录到日志；
    - 调用方无需 await，也无需保存返回值。
    """
    from kilostar.utils.logger import get_logger

    task: asyncio.Task[Any] = asyncio.ensure_future(coro)
    _background_tasks.add(task)

    def _on_done(t: "asyncio.Task[Any]") -> None:
        _background_tasks.discard(t)
        if not t.cancelled() and t.exception() is not None:
            get_logger("actor").exception(
                f"后台任务 {t.get_name()!r} 异常退出", exc_info=t.exception()
            )

    task.add_done_callback(_on_done)
    return task


# ─── 类型化 getter（核心单例，语义糖）──────────────────────────


def get_postgres() -> ActorHandle:
    return get_actor("postgres_database")


def get_gsm() -> ActorHandle:
    return get_actor("global_state_machine")


def get_gwm() -> ActorHandle:
    return get_actor("global_workflow_manager")


def get_regulatory() -> ActorHandle:
    return get_actor("regulatory_node")


def get_consciousness() -> ActorHandle:
    return get_actor("consciousness_node")


def get_worker_cluster() -> ActorHandle:
    return get_actor("worker_cluster")


def get_plugin_manager() -> ActorHandle:
    return get_actor("global_plugin_manager")