# 工具选择准确率随工具数量变化（经验数据）
tool_count_vs_accuracy = {
    10: 0.96,    # 10个工具：96%准确率
    30: 0.88,    # 30个工具：88%准确率
    50: 0.78,    # 50个工具：78%准确率
    100: 0.65,   # 100个工具：65%准确率（随机水平为1%）
}

from dataclasses import dataclass, field
from enum import Enum, auto
from typing import Any, Callable, Dict, List, Optional, Union
import json
import time
import asyncio
from collections import defaultdict
import hashlib


class ToolCategory(Enum):
    FILE = auto()        # 文件操作
    DATABASE = auto()    # 数据库
    NETWORK = auto()     # 网络请求
    CODE = auto()        # 代码执行
    COMMUNICATION = auto()  # 通信/邮件
    SYSTEM = auto()      # 系统管理
    SEARCH = auto()      # 搜索查询
    MISC = auto()        # 杂项


@dataclass
class ToolMetadata:
    """工具元数据——描述工具的能力和约束"""
    name: str
    description: str
    category: ToolCategory
    tags: List[str] = field(default_factory=list)
    rate_limit: int = 60        # 每分钟最大调用次数
    timeout: float = 30.0       # 超时时间（秒）
    cost_per_call: float = 0.0  # 每次调用的估算成本（美元）
    requires_auth: bool = False
    is_sensitive: bool = False  # 是否需要额外确认
    parallel_safe: bool = True  # 是否可以并行调用


@dataclass
class ToolRegistration:
    """完整的工具注册信息"""
    metadata: ToolMetadata
    handler: Callable
    schema: Dict[str, Any]      # JSON Schema
    version: str = "1.0.0"

class LayeredToolRegistry:
    """分层工具注册中心——支持语义路由和动态注入"""

    def __init__(self):
        self._tools: Dict[str, ToolRegistration] = {}
        self._category_index: Dict[ToolCategory, List[str]] = defaultdict(list)
        self._tag_index: Dict[str, List[str]] = defaultdict(list)
        self._embedding_cache: Dict[str, List[float]] = {}

    def register(self, registration: ToolRegistration):
        """注册工具到所有索引"""
        name = registration.metadata.name
        self._tools[name] = registration
        self._category_index[registration.metadata.category].append(name)
        for tag in registration.metadata.tags:
            self._tag_index[tag].append(name)

    def route(self, query: str, top_k: int = 15) -> List[ToolRegistration]:
        """基于查询语义路由到最匹配的工具"""
        # 第一阶段：关键词预过滤
        candidates = self._keyword_filter(query)
        # 第二阶段：语义排序
        scored = self._semantic_score(query, candidates)
        # 第三阶段：取Top-K
        return [reg for reg, _ in scored[:top_k]]

    def _keyword_filter(self, query: str) -> List[str]:
        """基于关键词快速预过滤"""
        query_lower = query.lower()
        candidates = set()
        for name, reg in self._tools.items():
            if any(keyword in query_lower for keyword in
                   [name.lower(), reg.metadata.description.lower()]):
                candidates.add(name)
        if len(candidates) < 5:
            return list(self._tools.keys())
        return list(candidates)

    def _semantic_score(self, query: str, candidates: List[str]) -> List:
        """基于简单语义相关性打分"""
        query_embedding = self._get_embedding(query)
        results = []
        for name in candidates:
            reg = self._tools[name]
            tool_embedding = self._get_embedding(
                f"{reg.metadata.name}: {reg.metadata.description}"
            )
            similarity = self._cosine_similarity(query_embedding, tool_embedding)
            results.append((reg, similarity))
        return sorted(results, key=lambda x: x[1], reverse=True)

    def _get_embedding(self, text: str) -> List[float]:
        """简单的哈希嵌入——生产环境应使用text-embedding-3-small等模型"""
        if text in self._embedding_cache:
            return self._embedding_cache[text]
        hash_bytes = hashlib.md5(text.encode()).digest()
        embedding = [b / 255.0 for b in hash_bytes]
        self._embedding_cache[text] = embedding
        return embedding

    def _cosine_similarity(self, a: List[float], b: List[float]) -> float:
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = sum(x * x for x in a) ** 0.5
        norm_b = sum(x * x for x in b) ** 0.5
        return dot / (norm_a * norm_b) if norm_a and norm_b else 0.0

class ParallelToolScheduler:
    """智能并行工具调度器——自动识别可并行的工具调用"""

    def __init__(self, max_concurrent: int = 10):
        self.max_concurrent = max_concurrent
        self._dependency_graph: Dict[str, List[str]] = {}

    async def execute_batch(
        self,
        calls: List[Dict[str, Any]],
        registry: LayeredToolRegistry
    ) -> List[Dict[str, Any]]:
        """批量执行工具调用，自动并行化独立调用"""
        self._build_dependency_graph(calls)
        layers = self._topological_sort(calls)
        results = []
        for layer in layers:
            batch_tasks = []
            for call in layer:
                tool_name = call["tool"]
                args = call.get("arguments", {})
                if tool_name in registry._tools:
                    reg = registry._tools[tool_name]
                    batch_tasks.append(self._execute_safe(reg, args, call.get("id")))
            layer_results = await asyncio.gather(*batch_tasks, return_exceptions=True)
            results.extend(layer_results)
        return results

    def _build_dependency_graph(self, calls: List[Dict]):
        """构建调用依赖图"""
        self._dependency_graph = {}
        call_map = {c.get("id", i): c for i, c in enumerate(calls)}
        for call in calls:
            call_id = call.get("id", str(id(call)))
            deps = []
            args_str = json.dumps(call.get("arguments", {}))
            for other_id in call_map:
                if other_id != call_id and f"${{{other_id}}}" in args_str:
                    deps.append(other_id)
            self._dependency_graph[call_id] = deps

    def _topological_sort(self, calls: List[Dict]) -> List[List]:
        """拓扑排序——将调用组织成可并行执行的层级"""
        call_ids = [c.get("id", str(id(c))) for c in calls]
        remaining = set(call_ids)
        layers = []
        while remaining:
            current_layer = []
            for cid in remaining:
                deps = self._dependency_graph.get(cid, [])
                if all(d not in remaining for d in deps):
                    current_layer.append(cid)
            if not current_layer:
                current_layer = list(remaining)
            layers.append([
                c for c in calls if c.get("id", str(id(c))) in current_layer
            ])
            remaining -= set(current_layer)
        return layers

    async def _execute_safe(self, reg, args, call_id=None):
        """安全执行单个工具调用"""
        start = time.time()
        try:
            result = await asyncio.wait_for(
                reg.handler(**args) if asyncio.iscoroutinefunction(reg.handler)
                else asyncio.to_thread(reg.handler, **args),
                timeout=reg.metadata.timeout
            )
            return {"id": call_id, "tool": reg.metadata.name, "status": "success",
                    "result": result, "latency": time.time() - start}
        except asyncio.TimeoutError:
            return {"id": call_id, "tool": reg.metadata.name, "status": "timeout",
                    "error": f"Timed out after {reg.metadata.timeout}s",
                    "latency": reg.metadata.timeout}
        except Exception as e:
            return {"id": call_id, "tool": reg.metadata.name, "status": "error",
                    "error": str(e), "latency": time.time() - start}

class RateLimiter:
    """多维度速率限制器——支持全局/工具/用户三级限流"""

    def __init__(self):
        self._limits: Dict[str, List[float]] = defaultdict(list)

    def set_limit(self, key: str, max_calls: int, window: int = 60):
        self._limits[key] = {"max_calls": max_calls, "window": window, "timestamps": []}

    def check(self, tool_name: str, user_id: str = "default") -> bool:
        keys = [f"global", f"tool:{tool_name}", f"user:{user_id}",
                f"user:{user_id}:tool:{tool_name}"]
        now = time.time()
        for key in keys:
            if key not in self._limits:
                continue
            limit = self._limits[key]
            limit["timestamps"] = [t for t in limit["timestamps"] if now - t < limit["window"]]
            if len(limit["timestamps"]) >= limit["max_calls"]:
                return False
        return True

    def wait_time(self, tool_name: str, user_id: str = "default") -> float:
        keys = [f"global", f"tool:{tool_name}", f"user:{user_id}",
                f"user:{user_id}:tool:{tool_name}"]
        now = time.time()
        max_wait = 0.0
        for key in keys:
            if key not in self._limits:
                continue
            limit = self._limits[key]
            if len(limit["timestamps"]) >= limit["max_calls"]:
                oldest = min(limit["timestamps"])
                wait = limit["window"] - (now - oldest)
                max_wait = max(max_wait, wait)
        return max_wait

import re


class ToolArgumentSanitizer:
    """工具参数清洗——防止参数注入攻击"""

    def __init__(self, policies: Dict[str, SecurityPolicy]):
        self.policies = policies

    def sanitize(self, tool_name: str, args: Dict) -> Dict:
        policy = self.policies.get(tool_name, SecurityPolicy(
            allowed_categories=[], sensitive_actions=[]
        ))
        cleaned = {}
        for key, value in args.items():
            str_value = str(value)
            for pattern in policy.deny_patterns:
                if re.search(pattern, str_value, re.IGNORECASE):
                    raise ValueError(f"Argument '{key}' contains denied pattern: {pattern}")
            if len(str_value) > policy.max_args_size:
                raise ValueError(f"Argument '{key}' exceeds max size of {policy.max_args_size}")
            cleaned[key] = value
        return cleaned

@dataclass
class ToolCallRecord:
    """工具调用记录——完整的可观测性数据"""
    call_id: str
    tool_name: str
    arguments: Dict
    result: Any
    status: str          # success / error / timeout / rate_limited
    latency: float       # 毫秒
    user_id: str
    session_id: str
    timestamp: float
    token_cost: int      # LLM Token消耗
    api_cost: float      # 外部API调用成本


class ToolCallTracker:
    """工具调用追踪器——记录所有调用用于审计和分析"""

    def __init__(self):
        self._records: List[ToolCallRecord] = []
        self._metrics: Dict[str, Any] = defaultdict(lambda: {
            "total_calls": 0, "success_calls": 0, "error_calls": 0,
            "total_latency_ms": 0, "total_cost": 0.0
        })

    def record(self, record: ToolCallRecord):
        self._records.append(record)
        metrics = self._metrics[record.tool_name]
        metrics["total_calls"] += 1
        if record.status == "success":
            metrics["success_calls"] += 1
        else:
            metrics["error_calls"] += 1
        metrics["total_latency_ms"] += record.latency
        metrics["total_cost"] += record.api_cost

    def get_metrics(self, tool_name=None) -> Dict:
        if tool_name:
            m = self._metrics.get(tool_name, {})
            return {
                "tool": tool_name,
                "total_calls": m["total_calls"],
                "success_rate": m["success_calls"] / m["total_calls"] * 100 if m["total_calls"] else 0,
                "avg_latency_ms": m["total_latency_ms"] / m["total_calls"] if m["total_calls"] else 0,
                "total_cost": m["total_cost"]
            }
        return {tool: {"total_calls": m["total_calls"],
                       "success_rate": m["success_calls"] / m["total_calls"] * 100 if m["total_calls"] else 0}
                for tool, m in self._metrics.items()}

# tool-engine.yaml - 生产部署配置
tool_engine:
  registry:
    embedding_model: "text-embedding-3-small"
    route_top_k: 15
    enable_auto_discovery: true

  scheduler:
    max_concurrent: 10
    default_timeout: 30
    enable_dependency_analysis: true

  rate_limiter:
    global:
      max_calls: 1000
      window: 60
    per_tool:
      max_calls: 60
      window: 60
    per_user:
      max_calls: 100
      window: 60

  cache:
    ttl: 300
    max_size_mb: 512
    auto_invalidate: true

  security:
    max_args_size: 10240
    deny_patterns:
      - "rm\\s+-rf"
      - "DROP\\s+TABLE"
      - "sudo"
    require_confirmation:
      - "delete_file"
      - "send_email"
      - "execute_sql"

  monitoring:
    enable_tracing: true
    metrics_export_interval: 60
    failure_alert_threshold: 0.1
    cost_tracking: true