"""KiloStar Ray 兼容层：单机/分布式模式无感切换 + 序列化工具。

单机模式下，所有 Actor 退化为普通 Python 异步单例，通过 StandaloneProxy
包装后暴露与 Ray Actor Handle 相同的 `.method.remote(args)` 调用接口，
使上层代码在两种模式间无感切换。
"""

from __future__ import annotations

import asyncio
import os
from typing import Any, Type, TypeVar

from pydantic import BaseModel

_STANDALONE = os.environ.get("KILOSTAR_MODE", "distributed") == "standalone"

T = TypeVar("T", bound=Type[BaseModel])


class _MethodProxy:
    """包装单个方法，使 .remote(*args, **kwargs) 返回一个可 await 的 Task。"""

    __slots__ = ("_method",)

    def __init__(self, method: Any):
        self._method = method

    def remote(self, *args: Any, **kwargs: Any) -> asyncio.Task:
        async def _invoke():
            result = self._method(*args, **kwargs)
            if asyncio.iscoroutine(result):
                return await result
            return result

        return asyncio.ensure_future(_invoke())


class StandaloneProxy:
    """包装一个普通 Python 实例，模拟 Ray Actor Handle 的属性访问接口。

    用法：proxy.some_method.remote(x, y) → 等效于 await instance.some_method(x, y)
    """

    __slots__ = ("_instance",)

    def __init__(self, instance: Any):
        object.__setattr__(self, "_instance", instance)

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


# ─── 条件装饰器 ───


def actor_class(cls):
    """条件装饰器：分布式模式 → @ray.remote，单机模式 → 原样返回类。"""
    if _STANDALONE:
        return cls
    import ray
    return ray.remote(cls)


def remote_task(func):
    """条件装饰器：分布式 → @ray.remote(func)，单机 → .remote() 转为 asyncio task。

    单机模式下返回一个 stub 对象，其 .remote() 方法把函数以协程方式调度到
    当前事件循环（workflow task 需要用 await 版本的 _entry，由调用方处理）。
    """
    if _STANDALONE:

        class _TaskProxy:
            @staticmethod
            def remote(*args, **kwargs):
                async def _run():
                    result = func(*args, **kwargs)
                    if asyncio.iscoroutine(result):
                        return await result
                    return result

                return asyncio.ensure_future(_run())

        return _TaskProxy()

    import ray
    return ray.remote(func)


# ─── Pickle (Ray 序列化优化) ───


def pickle(cls: T) -> T:
    """类装饰器：用 Pydantic 的高效 JSON 序列化替代 Python 原生 __reduce__，
    使 Ray 跨进程通信时对 BaseModel 子类走 Rust 级序列化。
    """

    def __reduce__(self):
        data = self.model_dump_json()
        return cls.model_validate_json, (data,)

    cls.__reduce__ = __reduce__
    return cls