"""函数注册表 - 管理字符串函数名到 Python 函数的映射。

支持自动发现和手动注册，与 FormulaParser 配合使用。

示例:
    >>> from src.factors.registry import FunctionRegistry
    >>> registry = FunctionRegistry(auto_scan=True)  # 自动加载 api.py 函数
    >>> registry.available_functions()[:5]
    ['abs', 'clip', 'cs_demean', 'cs_neutralize', 'cs_rank']
"""

import inspect
import typing
from typing import Any, Callable, Dict, List, Optional, Set

from src.factors.dsl import Node, FunctionNode
from src.factors.exceptions import DuplicateFunctionError


class FunctionRegistry:
    """函数注册表。

    管理字符串函数名到可调用对象的映射。
    自动从 api.py 加载标准函数，支持用户自定义函数注册。

    Attributes:
        _functions: 函数字典，name -> callable
    """

    def __init__(self, auto_scan: bool = True) -> None:
        """初始化注册表。

        Args:
            auto_scan: 是否自动扫描 api.py 模块，默认 True
        """
        self._functions: Dict[str, Callable] = {}

        if auto_scan:
            self._scan_api_module()

    def register(
        self, name: str, func: Callable, force: bool = False
    ) -> "FunctionRegistry":
        """注册自定义函数。

        Args:
            name: 函数名称（字符串形式）
            func: 可调用对象
            force: 是否强制覆盖已存在的函数，默认 False

        Returns:
            self（支持链式调用）

        Raises:
            DuplicateFunctionError: 当函数名已存在且 force=False 时

        Example:
            >>> registry = FunctionRegistry(auto_scan=False)
            >>> registry.register("my_func", lambda x: x * 2)
            >>> registry.get("my_func")(5)
            10
        """
        if name in self._functions and not force:
            raise DuplicateFunctionError(name)

        self._functions[name] = func
        return self

    def unregister(self, name: str) -> "FunctionRegistry":
        """注销函数。

        Args:
            name: 要注销的函数名

        Returns:
            self（支持链式调用）

        Raises:
            KeyError: 函数不存在时
        """
        if name not in self._functions:
            raise KeyError(f"函数 '{name}' 不存在")
        del self._functions[name]
        return self

    def get(self, name: str) -> Optional[Callable]:
        """获取函数。

        Args:
            name: 函数名称

        Returns:
            函数对象，不存在返回 None
        """
        return self._functions.get(name)

    def has(self, name: str) -> bool:
        """检查函数是否存在。

        Args:
            name: 函数名称

        Returns:
            是否存在
        """
        return name in self._functions

    def available_functions(self) -> List[str]:
        """返回所有可用函数名列表（按字母序）。

        Returns:
            排序后的函数名列表
        """
        return sorted(self._functions.keys())

    def clear(self) -> "FunctionRegistry":
        """清空所有注册的函数。

        Returns:
            self（支持链式调用）
        """
        self._functions.clear()
        return self

    def scan_module(
        self, module: Any, prefix: str = "", force: bool = False
    ) -> "FunctionRegistry":
        """扫描指定模块，自动注册符合条件的函数。

        扫描规则:
        1. 模块级别的函数（排除私有函数 _*）
        2. 返回类型注解为 Node 或 FunctionNode

        Args:
            module: 要扫描的模块对象
            prefix: 函数名前缀，用于避免命名冲突
            force: 是否强制覆盖已存在的函数

        Returns:
            self（支持链式调用）

        Example:
            >>> import my_custom_module
            >>> registry.scan_module(my_custom_module, prefix="custom_")
        """
        for name, obj in inspect.getmembers(module):
            # 只处理非私有函数
            if not inspect.isfunction(obj) or name.startswith("_"):
                continue

            # 检查是否应该注册
            if self._should_register(obj):
                full_name = prefix + name
                self.register(full_name, obj, force=force)

        return self

    def _scan_api_module(self) -> None:
        """自动扫描 api.py 模块，注册所有符合条件的函数。

        这是默认的自动扫描行为，在 __init__ 中调用。
        """
        try:
            from src.factors import api

            self.scan_module(api)
        except ImportError:
            # api 模块可能不存在，静默跳过
            pass

    def _should_register(self, func: Callable) -> bool:
        """检查函数是否应该被注册。

        基于类型提示检查函数返回类型，只注册返回 Node 或 FunctionNode 的函数。

        Args:
            func: 要检查的函数

        Returns:
            是否应该注册该函数
        """
        try:
            hints = typing.get_type_hints(func)
            return_type = hints.get("return")

            if return_type is None:
                return False

            # 处理 Union 类型（如 Union[Node, FunctionNode]）
            origin = typing.get_origin(return_type)
            args = typing.get_args(return_type)

            if origin is typing.Union:
                # Union 类型，检查任一参数
                return any(self._is_node_type(arg) for arg in args)
            else:
                # 单一类型
                return self._is_node_type(return_type)

        except Exception:
            return False

    def _is_node_type(self, typ: Any) -> bool:
        """检查类型是否是 Node 或 FunctionNode 的子类。

        Args:
            typ: 要检查的类型

        Returns:
            是否是 Node 相关类型
        """
        if not isinstance(typ, type):
            return False

        return issubclass(typ, (Node, FunctionNode))

    def __len__(self) -> int:
        """返回已注册函数数量。"""
        return len(self._functions)

    def __contains__(self, name: str) -> bool:
        """检查是否包含某个函数名。"""
        return name in self._functions

    def __repr__(self) -> str:
        """返回注册表字符串表示。"""
        return f"FunctionRegistry({len(self._functions)} functions: {self.available_functions()[:5]}...)"
