"""
基础类的定义
"""
import os
import time
from dataclasses import dataclass, field, asdict
from collections.abc import Callable
from typing import (
    List, Dict, Literal, Union, Optional, Any, TypeAlias
    
    )
import json
import inspect
from inspect import signature, Parameter, ismethod, iscoroutinefunction


AnyFunction: TypeAlias = Callable[..., Any]

@dataclass
class permission:
    groups: List[str] = field(default_factory=list, metadata={"help": "Model's groups"})
    users: List[str] = field(default_factory=list, metadata={"help": "Model's users"})
    owner: Union[str, None] = field(default=None, metadata={"help": "Model's owner"})

    def to_dict(self):
        return asdict(self)

@dataclass
class HModelConfig:
    name: str = field(default=None, metadata={"help": "Model's name"})
    permission: Union[str, Dict] = field(default=None, metadata={"help": "Model's permission, separated by ;, e.g., 'groups: all; users: a, b; owner: c', will inherit from worker permissions if not setted"})
    version: str = field(default="2.0", metadata={"help": "Model's version"})
    enable_mcp: bool = field(default=False, metadata={"help": "Enable MCP (Model Context Protocol) for LLM worker"})
    mcp_transport: Literal["sse", "streamable-http"] = field(default="streamable-http", metadata={"help": "MCP transport type, could be 'sse' or 'streamable-http'"})

# from ..worker.singletons import mcp_manager

class BaseWorkerModel:
 
    # _mcp: Any = None  # 占位，避免mypy报错
    
    # @property
    # def mcp(self):
    #     if self._mcp is None:
    #         from mcp.server.fastmcp import FastMCP
    #         self._mcp = FastMCP(name=self.name)
    #     return self._mcp
    
    # @mcp.setter
    # def mcp(self, value):
    #     self._mcp = value
    
    @classmethod
    def remote_callable(
        cls,
        name: str | None = None,
        title: str | None = None,
        description: str | None = None,
        annotations:  None = None,
        structured_output: bool | None = None,
        **decorator_kwargs,
    ) -> Callable[[AnyFunction], AnyFunction]:
        """
        用来修饰一个函数，使得其可以被远程调用，未修饰的函数不能被远程调用
        Decorator to mark a method as remotely callable.
        """
        
        enable_mcp = decorator_kwargs.pop("enable_mcp", False)
        
        # def decorator(func):
        #     func.is_remote_callable = True
        #     # 可以在这里添加更多基于decorator_args和decorator_kwargs的处理
        def decorator(fn: AnyFunction) -> AnyFunction:
            mcp_manager.register_mcp(name, fn)
            cls().mcp.add_tool(
                fn,
                name=name,
                title=title,
                description=description,
                annotations=annotations,
                structured_output=structured_output,
            )
            return fn

        # 如果没有参数直接调用，则返回装饰器函数
        # if len(decorator_args) == 1 and callable(decorator_args[0]):
        #     # 这种情况是直接 @HRModel.remote_callable 使用，没有括号
        #     # return decorator(decorator_args[0])
        #     func = decorator_args[0]
        #     func.is_remote_callable = True
        #     return func
        if callable(name):
            # 说明是@BaseWorkerModel.remote_callable直接修饰函数的形式
            func = name
            func.is_remote_callable = True
            return func
        else:
            return decorator
       
    
    @classmethod
    def remote_callable_backup(
        cls, 
        func
        ):
        """
        用来修饰一个函数，使得其可以被远程调用，未修饰的函数不能被远程调用
        Decorator to mark a method as remotely callable.
        """
        func.is_remote_callable = True
        return func
    
    @property
    def all_remote_callables(self):
        """
        获取所有被修饰的远程调用函数
        """
        methods = {}
        # Traverse the method resolution order (MRO), which includes the class and its bases
        for subclass in inspect.getmro(self.__class__):
            for name, method in inspect.getmembers(subclass, inspect.isfunction):
                methods[name] = method
        callable_funcs = list(methods.keys())
        # 再判断这些函数是不是可以远程调用的
        callable_funcs = [f for f in callable_funcs if hasattr(methods[f], "is_remote_callable")]
        return callable_funcs
    
    @classmethod
    def run(cls, **worker_overrides):
        """
        Run the worker model as a standalone service.

        """

        import uvicorn
        import hepai as hai
        from ..worker.worker_app import HWorkerAPP, HWorkerConfig

        model_config, worker_config = hai.parse_args((HModelConfig, HWorkerConfig))
        app = HWorkerAPP(
            cls(config=model_config),
            worker_config=worker_config,
            **worker_overrides,
            )  # Instantiate the APP, which is a FastAPI application.
        print(app.worker.get_worker_info(), flush=True)
        # 启动服务
        uvicorn.run(app, host=app.host, port=app.port)

    @classmethod
    def connect(
        cls,
        name: str,  # 远程模型的名称
        base_url: str = "https://aiapi.ihep.ac.cn/apiv2",  # 远程模型的地址
        api_key: str = None,  # 远程模型的API Key
        **kwargs,
        ):

        import logging
        import httpx
        # 将 httpx 的日志记录级别调整为 WARNING
        logging.getLogger("httpx").setLevel(logging.WARNING)
        
        api_key = api_key if api_key else os.environ.get("HEPAI_API_KEY")
        
        from hepai import HepAI
        client = HepAI(
            base_url=base_url,
            api_key=api_key,
            **kwargs,
            )

        from ..hclient._remote_model import LRModel

        model: LRModel = client.get_remote_model(model_name=name)
        return model
    
    @classmethod
    def list_models(
        cls,
        base_url: str = None,  # 远程模型的地址
        **kwargs,
        ):
        import logging
        import httpx
        # 将 httpx 的日志记录级别调整为 WARNING
        logging.getLogger("httpx").setLevel(logging.WARNING)

        base_url = base_url if base_url else "http://localhost:4260/apiv2"
        
        from hepai import HepAI
        client = HepAI(base_url=base_url,**kwargs)

        from ..hclient._remote_model import LRModel
        models = client.models.list()
        return models

    @classmethod
    async def async_connect(
            cls,
            name: str,  # 远程模型的名称
            base_url: str,  # 远程模型的地址
            # api_key: str = None,  # 远程模型的API Key
            **kwargs,
    ):

        import logging
        import httpx
        # 将 httpx 的日志记录级别调整为 WARNING
        logging.getLogger("httpx").setLevel(logging.WARNING)

        from hepai import AsyncHepAI
        client = AsyncHepAI(
            base_url=base_url,
            # api_key=api_key,
            **kwargs,
        )

        from ..hclient._remote_model import LRModel

        model: LRModel = await client.get_remote_model(model_name=name)
        return model

DEFAULT_STREAM_DATA = [
    1, 2, 3,
    "x", "y", "z",
    [[1, 2], [3, 4], [5, 6]],
    {"a": "b", "c": "d"},
]

from mcp.server.fastmcp import FastMCP
    
class HRemoteModel(BaseWorkerModel):
    """
    The Remote Model of HAI Framework
    """
    def __init__(
            self,
            name: str = None,
            config: HModelConfig = None,
            mcp: FastMCP = None,
            ):
        self.config = config if config is not None else HModelConfig()
        assert isinstance(self.config, HModelConfig), "config must be an instance of HModelConfig"
        self.config.name = name if name is not None else self.config.name
        self.config.name = self.config.name if self.config.name else self.__class__.__name__
        self.name = self.config.name
        
        from ..worker import utils
        self.model_id = utils.gen_one_id(lenth=15, prefix="md-", extra_indicators=["model", self.name])
      
        self.permission = self.config.permission
        
        self.created = int(time.time())
        
        self.register_functions: List[Dict[str, str]] = self.load_functions()
        
        enable_mcp = self.config.enable_mcp if hasattr(self.config, 'enable_mcp') else False
        if enable_mcp:
            if mcp is not None:
                self.mcp = mcp
            else:
                # 这是把用RemoteCallable修饰的函数注册到MCP服务器上
                from mcp.server.fastmcp import FastMCP
                self.mcp = FastMCP(name=self.name)
                remote_callables = self.all_remote_callables
                for func_name in remote_callables:
                    func = getattr(self, func_name)
                    self.mcp.add_tool(func)
        else:
            self.mcp = None
        

    @BaseWorkerModel.remote_callable
    def hello_world(self, *args, **kwargs):
        """An example of a function that returns a string"""
        return f"Hello world! You are using the HepAI worker model with args: `{args}`, kwargs: `{kwargs}`"

    # @BaseWorkerModel.remote_callable
    # def get_int(self, a: int = 1, b: int = 2) -> int:
    #     """An example of a function that returns an int type"""
    #     return a + b
    
    # @BaseWorkerModel.remote_callable
    # def get_float(self, a: float = 1.1, b: float = 2.2) -> float:
    #     """An example of a function that returns a float type"""
    #     return a + b
    
    # @BaseWorkerModel.remote_callable
    # def get_list(self, a: List[int] = [1, 2], b: List[int] = [3, 4]) -> List[int]:
    #     """An example of a function that returns a list type"""
    #     return a + b
    
    # @BaseWorkerModel.remote_callable
    # def get_dict(self, a: Dict[str, int] = {"a1": 1}, b: Dict[str, int] = {"a2": 2}) -> Dict[str, int]:
    #     """An example of a function that returns a dict type"""
    #     return {**a, **b}
    
    # @BaseWorkerModel.remote_callable
    # def get_stream(self, data: Any = None, interval: float = 0.2):
    #     """An example of a function that returns a stream type"""

    #     data = data if data is not None else DEFAULT_STREAM_DATA
    #     for i, x in enumerate(data):
    #         time.sleep(interval)  # 注：此处为了演示，故意加了延迟，实际使用时应该去掉
    #         yield f"data: {json.dumps(x)}\n\n"

    @BaseWorkerModel.remote_callable
    def __call__(self, *args, **kwargs):
        return f"Hello world! You are calling function `__call__` of the HepAI remote model with args: `{args}`, kwargs: `{kwargs}`"

    
    def _clean_type_str(self, annotation):
        """类型注解清洗方法"""
        if annotation == Parameter.empty:
            return "Any"
        
        # 处理标准类型
        if isinstance(annotation, type):
            return annotation.__name__
        
        # 处理typing模块类型
        type_str = str(annotation)
        
        # 去除模块前缀
        type_str = type_str.replace('typing.', '')
        
        # 处理特殊类型
        if '[' not in type_str:  # 简单类型
            return type_str.split('.')[-1]  # 获取最终类型名
        
        # 处理泛型类型 (如List[int])
        try:
            main_type, args = type_str.split('[', 1)
            main_type = main_type.split('.')[-1]  # 提取主类型名
            # args = args.rstrip(']')
            last_bracket = args.rfind(']')
            if last_bracket != -1:
                args = args[:last_bracket]  # 保留内部所有括号
            return f"{main_type}[{args}]"
        except:
            return type_str
    
    def load_functions(self) -> List[Dict[str, str]]:
        register_functions = []
        # 遍历当前类及其所有父类
        for cls in self.__class__.__mro__:
            # 仅处理CustomWorkerModel及其子类（假设CustomWorkerModel是基类）
            if cls is HRModel or issubclass(cls, HRModel):
                for attr_name in dir(cls):
                    attr = getattr(cls, attr_name)
                    # 检查类属性是否被标记为remote_callable
                    if getattr(attr, "is_remote_callable", False):
                        if attr.is_remote_callable:
                            # 获取实例对应的属性（处理可能的实例覆盖）
                            instance_attr = getattr(self, attr_name)
                            if callable(instance_attr):
                                # 获取原始函数（处理装饰器包装）
                                raw_func = getattr(instance_attr, '__wrapped__', instance_attr)
                                # 处理异步函数
                                if iscoroutinefunction(raw_func):
                                    sig = signature(raw_func)
                                else:
                                    sig = signature(instance_attr)
                                
                                # 构建参数结构
                                params = []
                                for param in sig.parameters.values():
                                    if param.name == 'self' and ismethod(instance_attr):
                                        continue  # 跳过实例方法的self参数
                                    params.append({
                                        "name": param.name,
                                        # "type": str(param.annotation) if param.annotation != Parameter.empty else "Any",
                                        "type": self._clean_type_str(param.annotation) if param.annotation != Parameter.empty else "Any",
                                        "default": param.default if param.default != Parameter.empty else None
                                    })
                                
                                register_functions.append({
                                    "__name__": attr_name,
                                    "__doc__": instance_attr.__doc__,
                                    "__signature__": params,
                                    # "__return__": str(sig.return_annotation) if sig.return_annotation != Parameter.empty else "Any"
                                    "__return__": self._clean_type_str(sig.return_annotation) if sig.return_annotation != Parameter.empty else "Any",
                                })
        other_functions = ['__call__', 'get_dict', 'get_float', 'get_int', 'get_list', 'get_stream', 'hello_world', 'get_register_functions']
        # 过滤掉不需要的函数
        register_functions = [func for func in register_functions if func['__name__'] not in other_functions]
        return register_functions

    @BaseWorkerModel.remote_callable
    async def get_register_functions(self) -> list[dict[str, str]]:
        """
        获取当前模型中可调用的函数列表

        Returns:
            List[Dict[str, str]]: 注册函数列表，包含函数名、函数描述、函数签名
        """
        return self.register_functions
    
    
class HRModel(HRemoteModel):
    """
    Alias of HepAI Remote Model
    """
    ...
    
class HCloudModel(HRemoteModel):
    """
    The Cloud Model of HAI Framework
    """
    ...

        
    

@dataclass
class ModelResourceInfo:
    """
    Model Resource Info for Worker, such as llm, nn, preceptor, actuator, etc.
    """
    model_name: str = field(default="<default_modelname>", metadata={"help": "Model's name"})
    model_type: str = field(default="common", metadata={"help": "Model's type"})
    model_version: str = field(default="1.0", metadata={"help": "Model's version"})
    model_description: str =field(default="<This is model description.>", metadata={"help": "Model's description"})
    model_author: Union[str, List[str], None] = field(default="", metadata={"help": "Model's author"})
    model_owner: Union[str, None] = field(default="", metadata={"help": "Model's owner"})
    model_groups: List[str] = field(default_factory=list, metadata={"help": "Model's groups"})
    model_users: List[str] = field(default_factory=list, metadata={"help": "Model's users"})
    model_functions: List[str] = field(default_factory=list, metadata={"help": "Model's functions that can be called by remote"})
    id: Optional[str] = field(default=None, metadata={"help": "Model's id, usually set by the system"})
    created: Optional[int] = field(default=None, metadata={"help": "Model's created timestamp, usually set by the system"})
    object: Optional[str] = field(default="model", metadata={"help": "Model's object type, usually set by the system"})
    owned_by: Optional[str] = field(default=None, metadata={"help": "Model's owned by, usually set by the system"})

    def __post_init__(self):
        if self.id is None and self.model_name:
            self.id = self.model_name
        if self.created is None:
            self.created = int(time.time())
        if self.owned_by is None and self.model_owner:
            self.owned_by = self.model_owner

    def to_dict(self):
        return asdict(self)
    
    def __repr__(self):
        return f'ModelResourceInfo(model_id={self.id!r}, model_name={self.model_name!r}, model_type={self.model_type!r})'

@dataclass
class WorkerStatusInfo:
    """
    Worker Status Info, will be dynamic updated
    """
    speed: int = field(default=1, metadata={"help": "Worker's speed, the number of requests that can be processed per second"})
    queue_length: int = field(default=0, metadata={"help": "Worker's queue length"})
    status: Literal["idle", "ready", "busy", "error"] = "idle"
    start_time: Optional[float] = field(default=None, metadata={"help": "Worker's start time, timestamp"})

    def is_valid(self):
        """是信息是否可用，即相关信息是否已被填入，而不是None"""
        return self.speed > 0 and self.queue_length >= 0
    
    def to_dict(self):
        return asdict(self)


@dataclass
class WorkerNetworkInfo:
    """
    Network Info for Worker
    """
    host: str = field(default="127.0.0.1", metadata={"help": "Worker's host"})
    port: int = field(default=42602, metadata={"help": "Worker's port"})
    route_prefix: Union[str, None] = field(default="", metadata={"help": "Worker's route prefix, default is '/', the worker's address will be `http://host:port/route_prefix/other_router` if setted"})
    host_name: str = field(default="localhost", metadata={"help": "Worker's host name"})
    worker_address: Union[None, str] = field(default="", metadata={"help": "Worker's address, will be auto generated if not setted"})

    def check_and_autoset_worker_address(self):
        """自动检查并设置worker_address"""
        if self.worker_address in ["", None]:
            if self.route_prefix in ["", None]:
                self.worker_address = f"http://{self.host}:{self.port}"
            else:
                self.worker_address = f"http://{self.host}:{self.port}/{self.route_prefix}"
        return self.worker_address

    def to_dict(self):
        return asdict(self)


@dataclass
class WorkerInfo:
    """
    Worker Info
    v2.1.2 support multi model
    """
    id: str
    type: Literal["llm", "actuator", "preceptor", "memory", "common"] = "common"
    network_info: WorkerNetworkInfo = field(default_factory=WorkerNetworkInfo, metadata={"help": "Worker's network info"})
    resource_info: List[ModelResourceInfo] = field(default_factory=list, metadata={"help": "Model's resource info"})
    status_info: WorkerStatusInfo = field(default_factory=WorkerStatusInfo, metadata={"help": "Worker's status info"})
    check_heartbeat: bool = True
    last_heartbeat: Union[int, None] = None
    version: str = "2.0"
    metadata: Dict = field(default_factory=dict, metadata={"help": "Worker's metadata"})

    def __post_init__(self):
        """在实例化本类时，自动检查network_info等是否Dict, 并转换为相应的对象"""
        if isinstance(self.network_info, dict):
            self.network_info = WorkerNetworkInfo(**self.network_info)
        if isinstance(self.resource_info, list):
            # 后处理自动把json处理成对象
            for i, mr in enumerate(self.resource_info):
                if isinstance(mr, dict):
                    self.resource_info[i] = ModelResourceInfo(**mr)
        if isinstance(self.status_info, dict):
            self.status_info = WorkerStatusInfo(**self.status_info)

    def to_dict(self):
        return asdict(self)
    
    @property
    def model_names(self):
        return [x.model_name for x in self.resource_info]
    
    def get_model_resource(self, model_name: str) -> ModelResourceInfo:
        """v2.1.2多模型模式下，获取模型信息"""
        mr = [x for x in self.resource_info if x.model_name == model_name]
        if len(mr) != 1:
            # print(f'ERROR: {self.resource_info}')
            raise ValueError(f'[WorkerInfo] Failed to get_model_info by name `{model_name}`, got {len(mr)} models')
        return mr[0]
    
    def get_model_index(self, model_name: str) -> int:
        """根据模型名，获取其在resource_info中的索引"""
        return self.model_names.index(model_name)
    
    
    def to_openai_list_models(self) -> Dict:
        """转换为OpenAI格式的list_models的返回列表"""
        data = []
        for i, resc in enumerate(self.resource_info):
            tmp = {}
            owned_by = resc.model_owner if resc.model_owner else resc.model_author
            tmp["id"] = resc.model_name
            tmp["created"] = resc.created
            tmp["object"] = "model"
            tmp['owned_by'] = owned_by
            data.append(tmp)
        return data
    
    
    def __repr__(self):
        """自定义repr，resource_info很长时只显示前5和后5个"""
        base = f"WorkerInfo(\n  id={self.id!r},\n  type={self.type!r}, "
        n = len(self.resource_info)
        if n > 10:
            shown = (
                [repr(x) for x in self.resource_info[:5]] +
                ["..."] +
                [repr(x) for x in self.resource_info[-5:]]
            )
        else:
            shown = [repr(x) for x in self.resource_info]
        return (
            f"{base}\n"
            f"  resource_info=[\n    " +
            ",\n    ".join(shown) +
            f"\n  ] ({n}), " +
            f"\n  network_info={self.network_info!r}," +
            f"\n  status_info={self.status_info!r}," +
            f"\n  check_heartbeat={self.check_heartbeat!r}," +
            f"\n  last_heartbeat={self.last_heartbeat!r}," +
            f"\n  version={self.version!r}" +
            # f"\n  metadata={self.metadata!r}," +
            "\n)"
        )


from pydantic import BaseModel

class WorkerInfoItem(BaseModel):
    """用于FastAPI的WorkerInfo"""
    id: str = field(default="wk-default_id", metadata={"help": "Worker's id"})     
    type: Union[Literal["llm", "actuator", "preceptor", "memory", "common"], str] = "common"
    network_info: WorkerNetworkInfo = field(default_factory=WorkerNetworkInfo, metadata={"help": "Worker's network info"})
    resource_info: List[ModelResourceInfo] = field(default_factory=list, metadata={"help": "Model's resource info"})
    status_info: WorkerStatusInfo = field(default_factory=WorkerStatusInfo, metadata={"help": "Worker's status info"})
    check_heartbeat: bool = field(default=True, metadata={"help": "Check worker's heartbeat"})
    last_heartbeat: Union[int, None] = field(default=None, metadata={"help": "Worker's last heartbeat"})
    version: str = field(default="2.0", metadata={"help": "Worker's version"})
    metadata: Dict = field(default_factory=dict, metadata={"help": "Worker's metadata"})


@dataclass
class WorkerStoppedInfo:
    id: str = field(metadata={"help": "Worker's id"}) 
    stopped: bool = field(metadata={"help": "Worker's stopped flag"})
    message: str = field(default=None, metadata={"help": "Worker's stopped message"})
    shutdown: bool = field(default=False, metadata={"help": "Worker's shutdown"})

    def to_dict(self):
        return {
            "id": self.id,
            "stopped": self.stopped,
            "message": self.message,
            "shutdown": self.shutdown,
        }