"""
Function Call 功能模块

提供AI大模型的function call功能，支持动态注册和调用ketacli的各种功能。
"""

import json
import inspect
import asyncio
import os
from concurrent.futures import ThreadPoolExecutor
from typing import Dict, List, Any, Callable, Optional
from dataclasses import dataclass
from rich.console import Console
import logging

console = Console(markup=False)
logger = logging.getLogger("ketacli.textual")
tool_logger = logging.getLogger("ketacli.toolcalls")
if not tool_logger.handlers:
    tool_logger.setLevel(logging.DEBUG)
    try:
        base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
        log_dir = os.path.join(base_dir, "log")
        os.makedirs(log_dir, exist_ok=True)
        log_path = os.path.join(log_dir, "tool_calls.log")
        fh = logging.FileHandler(log_path)
        fh.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s"))
        tool_logger.addHandler(fh)
    except Exception:
        pass

def _preview_text(value: Any, limit: int = 1000) -> str:
    try:
        if value is None:
            return ""
        if isinstance(value, bytes):
            try:
                value = value.decode("utf-8", errors="replace")
            except Exception:
                value = str(value)
        if isinstance(value, (dict, list)):
            try:
                text = json.dumps(value, ensure_ascii=False)[:limit]
                return text
            except Exception:
                value = str(value)
        text = str(value)
        if len(text) > limit:
            return text[:limit] + "..."
        return text
    except Exception:
        return ""


@dataclass
class FunctionSchema:
    """函数模式定义"""
    name: str
    description: str
    parameters: Dict[str, Any]
    function: Callable


class FunctionRegistry:
    """函数注册表"""
    
    def __init__(self):
        self._functions: Dict[str, FunctionSchema] = {}
    
    def register(self, name: str, description: str, parameters: Dict[str, Any] = None):
        """注册函数装饰器"""
        def decorator(func: Callable):
            if parameters is None:
                # 自动从函数签名生成参数模式
                sig = inspect.signature(func)
                auto_params = self._generate_parameters_from_signature(sig)
            else:
                auto_params = parameters
            
            schema = FunctionSchema(
                name=name,
                description=description,
                parameters=auto_params,
                function=func
            )
            self._functions[name] = schema
            return func
        return decorator
    
    def _generate_parameters_from_signature(self, sig: inspect.Signature) -> Dict[str, Any]:
        """从函数签名自动生成参数模式"""
        properties = {}
        required = []
        
        for param_name, param in sig.parameters.items():
            if param_name.startswith('_'):  # 跳过私有参数
                continue
                
            param_info = {
                "type": "string"  # 默认类型
            }
            
            # 根据类型注解推断参数类型
            if param.annotation != inspect.Parameter.empty:
                if param.annotation == int:
                    param_info["type"] = "integer"
                elif param.annotation == float:
                    param_info["type"] = "number"
                elif param.annotation == bool:
                    param_info["type"] = "boolean"
                elif param.annotation == list:
                    param_info["type"] = "array"
                elif param.annotation == dict:
                    param_info["type"] = "object"
            
            # 如果没有默认值，则为必需参数
            if param.default == inspect.Parameter.empty:
                required.append(param_name)
            else:
                param_info["default"] = param.default
            
            properties[param_name] = param_info
        
        return {
            "type": "object",
            "properties": properties,
            "required": required
        }
    
    def get_function(self, name: str) -> Optional[FunctionSchema]:
        """获取注册的函数"""
        return self._functions.get(name)
    
    def list_functions(self) -> List[FunctionSchema]:
        """列出所有注册的函数"""
        return list(self._functions.values())
    
    def get_openai_tools_format(self) -> List[Dict[str, Any]]:
        """获取OpenAI tools格式的函数定义"""
        tools = []
        for func_schema in self._functions.values():
            tool = {
                "type": "function",
                "function": {
                    "name": func_schema.name,
                    "description": func_schema.description,
                    "parameters": func_schema.parameters
                }
            }
            tools.append(tool)
        return tools


class FunctionCallExecutor:
    """函数调用执行器"""
    
    def __init__(self, registry: FunctionRegistry):
        self.registry = registry
        self.executor = ThreadPoolExecutor(max_workers=4)  # 线程池用于异步执行
        # 设置线程池为守护模式，这样主线程退出时不会阻塞
        import atexit
        atexit.register(self.shutdown)
    
    def shutdown(self):
        """关闭线程池，确保资源释放"""
        if hasattr(self, 'executor') and self.executor:
            self.executor.shutdown(wait=False)
    
    def execute(self, function_name: str, arguments: Dict[str, Any]) -> Any:
        """执行函数调用"""
        func_schema = self.registry.get_function(function_name)
        if not func_schema:
            raise ValueError(f"未找到函数: {function_name}")
        
        try:
            logger.debug(f"[fc.sync] 调用函数: {function_name} args={arguments}")
            tool_logger.info(f"[start] fn={function_name} args={_preview_text(arguments)}")
            # 执行函数
            result = func_schema.function(**arguments)
            tool_logger.info(
                f"[success] fn={function_name} type={type(result).__name__} preview={_preview_text(result)}"
            )
            return {
                "success": True,
                "result": result,
                "function_name": function_name
            }
        except Exception as e:
            console.print(f"[red]执行函数 {function_name} 时出错: {e}[/red]")
            tool_logger.error(f"[error] fn={function_name} err={str(e)}")
            return {
                "success": False,
                "error": str(e),
                "function_name": function_name
            }
    
    def execute_from_tool_calls(self, tool_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """从AI模型的tool_calls执行函数"""
        results = []
        for tool_call in tool_calls:
            if tool_call.get("type") == "function":
                function_info = tool_call.get("function", {})
                function_name = function_info.get("name")
                raw_args = function_info.get("arguments", "{}")

                try:
                    if isinstance(raw_args, dict):
                        arguments = raw_args
                    elif isinstance(raw_args, str):
                        arguments = json.loads(raw_args or "{}")
                    else:
                        # 兜底：非常规类型一律包裹为 value
                        arguments = {"value": raw_args}

                    tool_logger.info(
                        f"[start] fn={function_name} tool_call_id={tool_call.get('id')} args={_preview_text(arguments)}"
                    )
                    result = self.execute(function_name, arguments)
                    result["tool_call_id"] = tool_call.get("id")
                    results.append(result)
                except Exception as e:
                    tool_logger.error(
                        f"[error] fn={function_name} tool_call_id={tool_call.get('id')} err={str(e)}"
                    )
                    results.append({
                        "success": False,
                        "error": f"参数解析错误: {e}",
                        "function_name": function_name,
                        "tool_call_id": tool_call.get("id")
                    })

        return results
    
    # ==================== 异步方法 ====================
    
    async def execute_async(self, function_name: str, arguments: Dict[str, Any]) -> Any:
        """异步执行函数调用"""
        func_schema = self.registry.get_function(function_name)
        if not func_schema:
            raise ValueError(f"未找到函数: {function_name}")
        
        try:
            # 在线程池中执行函数，避免阻塞主线程
            loop = asyncio.get_event_loop()
            
            # 创建一个包装函数来正确传递参数
            def execute_with_args():
                return func_schema.function(**arguments)
            
            logger.debug(f"[fc.async] 调用函数: {function_name} args={arguments}")
            tool_logger.info(f"[start.async] fn={function_name} args={_preview_text(arguments)}")
            result = await loop.run_in_executor(
                self.executor, 
                execute_with_args
            )
            tool_logger.info(
                f"[success.async] fn={function_name} type={type(result).__name__} preview={_preview_text(result)}"
            )
            return {
                "success": True,
                "result": result,
                "function_name": function_name
            }
        except Exception as e:
            console.print(f"[red]异步执行函数 {function_name} 时出错: {e}[/red]")
            tool_logger.error(f"[error.async] fn={function_name} err={str(e)}")
            return {
                "success": False,
                "error": str(e),
                "function_name": function_name
            }
    
    async def execute_from_tool_calls_async(self, tool_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """异步从AI模型的tool_calls执行函数"""
        tasks = []
        timeouts: List[int] = []

        for tool_call in tool_calls:
            if tool_call.get("type") == "function":
                function_info = tool_call.get("function", {})
                function_name = function_info.get("name")
                raw_args = function_info.get("arguments", "{}")

                try:
                    if isinstance(raw_args, dict):
                        arguments = raw_args
                    elif isinstance(raw_args, str):
                        arguments = json.loads(raw_args or "{}")
                    else:
                        arguments = {"value": raw_args}

                    logger.debug(f"[fc.tool] 解析参数成功: fn={function_name} args={arguments}")
                    tool_logger.info(
                        f"[start.tool] fn={function_name} tool_call_id={tool_call.get('id')} args={_preview_text(arguments)}"
                    )
                    # 创建异步任务
                    task = self._execute_single_tool_call_async(
                        function_name, arguments, tool_call.get("id")
                    )
                    tasks.append(task)
                    try:
                        timeouts.append(int(arguments.get("timeout", 30)))
                    except Exception:
                        timeouts.append(30)
                except Exception as e:
                    # 对于解析错误，创建一个返回错误的协程
                    async def create_error_result():
                        return {
                            "success": False,
                            "error": f"参数解析错误: {e}",
                            "function_name": function_name,
                            "tool_call_id": tool_call.get("id")
                        }
                    tasks.append(create_error_result())
        
        # 并发执行所有工具调用，添加超时处理
        if tasks:
            try:
                # 计算整体超时，取各工具声明的最大timeout
                overall_timeout = max(timeouts) if timeouts else 30
                logger.debug(f"[fc.tool] 并发执行工具调用，overall_timeout={overall_timeout}s, tasks={len(tasks)}")
                # 添加超时处理，避免工具执行时间过长导致卡死
                results = await asyncio.wait_for(
                    asyncio.gather(*tasks, return_exceptions=True),
                    timeout=overall_timeout
                )
                
                # 处理异常结果
                processed_results = []
                for result in results:
                    if isinstance(result, Exception):
                        processed_results.append({
                            "success": False,
                            "error": f"工具执行异常: {str(result)}",
                            "function_name": "unknown"
                        })
                        logger.error(f"[fc.tool] 工具执行异常: {str(result)} fn={function_name}")
                        tool_logger.error(f"[error.tool] fn={function_name} err={str(result)}")
                    else:
                        tool_logger.info(
                            f"[success.tool] fn={result.get('function_name')} tool_call_id={result.get('tool_call_id')} preview={_preview_text(result.get('result'))}"
                        )
                        processed_results.append(result)
                return processed_results
            except asyncio.TimeoutError:
                # 超时处理
                print("工具执行超时，强制返回")
                return [{
                    "success": False,
                    "error": f"工具执行超时，请稍后重试",
                    "function_name": "timeout"
                }]
            except Exception as e:
                # 捕获其他异常
                print(f"工具执行过程中发生异常: {e}")
                return [{
                    "success": False,
                    "error": f"工具执行过程中发生异常: {str(e)}",
                    "function_name": "error"
                }]
        
        return []
    
    async def _execute_single_tool_call_async(self, function_name: str, arguments: Dict[str, Any], tool_call_id: str) -> Dict[str, Any]:
        """异步执行单个工具调用"""
        result = await self.execute_async(function_name, arguments)
        result["tool_call_id"] = tool_call_id
        return result


# 全局函数注册表
function_registry = FunctionRegistry()
function_executor = FunctionCallExecutor(function_registry)
