# 文件: conductor/scheduler/cache_manager.py (推测)
from typing import Optional, List
from dataclasses import dataclass

@dataclass
class PrefixCacheInfo:
    """前缀缓存信息"""
    cache_key: str          # 缓存键
    prefix_length: int      # 前缀长度
    locations: List[str]    # 存储位置列表 ["node1:seg1", "node2:seg2"]
    hit_count: int          # 命中次数
    last_access_time: float # 最后访问时间
    frequency_score: float  # 频率得分 (用于热度评估)

class CacheManager:
    def __init__(self):
        self.prefix_tree = PrefixTree()  # Trie 树索引
        self.cache_metadata = {}          # 缓存元数据

    def find_longest_prefix(
        self,
        prompt_tokens: List[int]
    ) -> Optional[PrefixCacheInfo]:
        """
        查找最长匹配前缀

        算法:
        1. 在 Trie 树中查找最长公共前缀
        2. 返回匹配的缓存信息

        时间复杂度: O(L), L 为 prompt 长度
        """
        # 1. 在 Trie 树中查找
        matched_node = self.prefix_tree.find_longest_match(prompt_tokens)

        if not matched_node:
            return None

        # 2. 获取缓存元数据
        cache_key = matched_node.cache_key
        cache_info = self.cache_metadata.get(cache_key)

        if not cache_info:
            return None

        # 3. 检查缓存是否过期
        if self._is_expired(cache_info):
            self._evict_cache(cache_key)
            return None

        # 4. 更新访问统计
        cache_info.hit_count += 1
        cache_info.last_access_time = time.time()
        cache_info.frequency_score = self._calculate_frequency_score(cache_info)

        logger.debug(
            f"Found prefix cache: key={cache_key}, "
            f"length={cache_info.prefix_length}, "
            f"hit_count={cache_info.hit_count}"
        )

        return cache_info

    def _calculate_frequency_score(self, cache_info: PrefixCacheInfo) -> float:
        """
        计算频率得分 (用于 LFU 驱逐策略)

        公式: score = hit_count / (time_since_creation)^0.5
        """
        time_since_creation = time.time() - cache_info.creation_time
        return cache_info.hit_count / (time_since_creation ** 0.5)

class PrefixTree:
    """
    Trie 树,用于快速前缀匹配

    示例:
        prompt1: [1, 2, 3, 4, 5]
        prompt2: [1, 2, 3, 6, 7]

        Trie:
            1 → 2 → 3 → 4 → 5 (cache_key: "abc123")
                      └─ 6 → 7 (cache_key: "def456")
    """
    class Node:
        def __init__(self):
            self.children = {}       # token → Node
            self.cache_key = None    # 叶子节点存储 cache_key
            self.is_leaf = False

    def __init__(self):
        self.root = self.Node()

    def insert(self, tokens: List[int], cache_key: str):
        """插入新的前缀"""
        node = self.root
        for token in tokens:
            if token not in node.children:
                node.children[token] = self.Node()
            node = node.children[token]

        node.is_leaf = True
        node.cache_key = cache_key

    def find_longest_match(self, tokens: List[int]) -> Optional[Node]:
        """查找最长匹配前缀"""
        node = self.root
        last_leaf = None

        for token in tokens:
            if token not in node.children:
                break

            node = node.children[token]
            if node.is_leaf:
                last_leaf = node  # 记录最后一个叶子节点

        return last_leaf

# 文件: conductor/scheduler/scheduler.py (推测)
class Scheduler:
    def select_node_near_cache(
        self,
        cache_locations: List[str],    # ["node1:seg1", "node2:seg2"]
        required_memory: int,           # 所需显存
        task_type: str                  # "prefill" or "decode"
    ) -> str:
        """
        选择靠近缓存的节点

        策略:
        1. 优先选择缓存所在节点 (如果有资源)
        2. 否则选择网络拓扑最近的节点
        3. 考虑负载均衡

        Returns:
            目标节点 ID
        """
        # 1. 提取缓存所在节点
        cache_nodes = [loc.split(":")[0] for loc in cache_locations]

        # 2. 过滤可用节点 (有足够资源)
        available_nodes = []
        for node_id in cache_nodes:
            node_state = self.get_node_state(node_id)

            if node_state.task_type != task_type:
                continue  # 跳过类型不匹配的节点

            if node_state.free_memory < required_memory:
                continue  # 资源不足

            if node_state.queue_length > self.config.max_queue_length:
                continue  # 队列过长

            available_nodes.append(node_id)

        # 3. 如果缓存节点有资源,直接选择
        if available_nodes:
            # 选择负载最低的缓存节点
            selected_node = min(
                available_nodes,
                key=lambda nid: self.get_node_state(nid).queue_length
            )

            logger.info(f"Selected cache node: {selected_node}")
            return selected_node

        # 4. 否则,选择网络拓扑最近的节点
        all_nodes = self.get_all_nodes(task_type)

        # 计算网络距离 (基于拓扑)
        min_distance = float('inf')
        best_node = None

        for node_id in all_nodes:
            node_state = self.get_node_state(node_id)

            # 检查资源
            if node_state.free_memory < required_memory:
                continue

            # 计算到所有缓存节点的平均距离
            avg_distance = sum(
                self.topology.get_distance(node_id, cache_node)
                for cache_node in cache_nodes
            ) / len(cache_nodes)

            # 加权:距离 + 负载
            score = avg_distance * 0.7 + node_state.queue_length * 0.3

            if score < min_distance:
                min_distance = score
                best_node = node_id

        logger.info(
            f"Selected nearby node: {best_node}, "
            f"distance={min_distance:.2f}"
        )

        return best_node

class NetworkTopology:
    """
    网络拓扑管理

    拓扑结构示例:
        ToR1 ─── [node1, node2]
        ToR2 ─── [node3, node4]
        Spine ── [ToR1, ToR2]
    """

    def get_distance(self, node1: str, node2: str) -> int:
        """
        计算两节点间的网络距离

        距离定义:
        - 同节点: 0
        - 同机架 (同 ToR): 1
        - 同集群 (经过 Spine): 2
        - 跨集群: 3
        """
        if node1 == node2:
            return 0

        tor1 = self.node_to_tor[node1]
        tor2 = self.node_to_tor[node2]

        if tor1 == tor2:
            return 1  # 同机架

        spine1 = self.tor_to_spine[tor1]
        spine2 = self.tor_to_spine[tor2]

        if spine1 == spine2:
            return 2  # 同集群

        return 3  # 跨集群

# 文件: conductor/scheduler/migration_manager.py (推测)
class MigrationManager:
    def should_migrate_cache(
        self,
        cache_info: PrefixCacheInfo,
        target_node: str
    ) -> bool:
        """
        判断是否应该迁移缓存

        决策因素:
        1. 缓存热度 (频率得分)
        2. 迁移成本 (网络带宽)
        3. 负载不均衡程度

        Returns:
            是否应该迁移
        """
        # 1. 检查缓存热度
        if cache_info.frequency_score < self.config.hot_cache_threshold:
            return False  # 不是热缓存,不值得迁移

        # 2. 计算迁移成本
        source_node = cache_info.locations[0].split(":")[0]
        migration_cost = self._estimate_migration_cost(
            source_node,
            target_node,
            cache_info.data_size
        )

        # 3. 估算迁移收益
        expected_benefit = self._estimate_migration_benefit(
            cache_info,
            target_node
        )

        # 4. 决策: 收益 > 成本?
        should_migrate = expected_benefit > migration_cost * 1.5  # 1.5x 安全系数

        logger.debug(
            f"Migration decision: {should_migrate}, "
            f"benefit={expected_benefit:.2f}, "
            f"cost={migration_cost:.2f}"
        )

        return should_migrate

    def _estimate_migration_cost(
        self,
        source_node: str,
        target_node: str,
        data_size: int
    ) -> float:
        """
        估算迁移成本 (单位: ms)

        成本 = 传输时间 + CPU 开销
        """
        # 1. 计算传输时间
        network_bandwidth = self.topology.get_bandwidth(source_node, target_node)
        transfer_time_ms = (data_size / network_bandwidth) * 1000

        # 2. CPU 开销 (序列化/反序列化)
        cpu_overhead_ms = 5.0

        return transfer_time_ms + cpu_overhead_ms

    def _estimate_migration_benefit(
        self,
        cache_info: PrefixCacheInfo,
        target_node: str
    ) -> float:
        """
        估算迁移收益 (单位: ms)

        收益 = 未来请求节省的传输时间
        """
        # 1. 预测未来请求数 (基于历史频率)
        predicted_requests = cache_info.frequency_score * 60  # 未来 1 分钟

        # 2. 计算当前位置的传输时间
        source_node = cache_info.locations[0].split(":")[0]
        current_transfer_time = self._estimate_migration_cost(
            source_node,
            target_node,
            cache_info.data_size
        )

        # 3. 迁移后,传输时间变为 0 (本地访问)
        benefit = predicted_requests * current_transfer_time

        return benefit

    def trigger_migration(
        self,
        cache_info: PrefixCacheInfo,
        target_node: str
    ):
        """
        触发异步缓存迁移

        流程:
        1. 在目标节点分配空间
        2. 通过 Transfer Engine 传输数据
        3. 更新元数据
        4. (可选) 删除原副本
        """
        migration_task = MigrationTask(
            cache_key=cache_info.cache_key,
            source_location=cache_info.locations[0],
            target_node=target_node,
            data_size=cache_info.data_size
        )

        # 异步执行迁移
        self.migration_queue.put(migration_task)

        logger.info(
            f"Triggered migration: {cache_info.cache_key} → {target_node}"
        )

# 文件: conductor/scheduler/prefill_scheduler.py (推测)
@dataclass
class ScheduleDecision:
    target_node: str                    # 目标节点
    use_prefix_cache: bool              # 是否使用前缀缓存
    cache_location: Optional[str]       # 缓存位置
    cache_length: int                   # 前缀长度
    need_migration: bool                # 是否需要迁移
    estimated_latency: float            # 预估延迟 (ms)

class PrefillScheduler:
    def schedule(self, request: PrefillRequest) -> ScheduleDecision:
        """
        为 Prefill 请求做调度决策

        流程:
        1. 查询前缀缓存
        2. 评估是否使用缓存
        3. 选择目标节点
        4. 判断是否需要迁移
        5. 生成调度决策
        """
        # 1. 查询前缀缓存
        cache_info = self.cache_manager.find_longest_prefix(
            request.prompt_tokens
        )

        if not cache_info:
            # 无缓存,选择负载最低节点
            target_node = self.select_least_loaded_node("prefill")

            return ScheduleDecision(
                target_node=target_node,
                use_prefix_cache=False,
                cache_location=None,
                cache_length=0,
                need_migration=False,
                estimated_latency=self._estimate_latency(request, None)
            )

        # 2. 评估缓存命中率
        hit_ratio = cache_info.prefix_length / len(request.prompt_tokens)

        if hit_ratio < self.config.min_hit_ratio:
            # 命中率过低,不使用缓存
            target_node = self.select_least_loaded_node("prefill")

            return ScheduleDecision(
                target_node=target_node,
                use_prefix_cache=False,
                cache_location=None,
                cache_length=0,
                need_migration=False,
                estimated_latency=self._estimate_latency(request, None)
            )

        # 3. 选择靠近缓存的节点
        target_node = self.select_node_near_cache(
            cache_locations=cache_info.locations,
            required_memory=request.estimate_memory(),
            task_type="prefill"
        )

        # 4. 判断是否需要迁移
        need_migration = False
        if target_node not in [loc.split(":")[0] for loc in cache_info.locations]:
            need_migration = self.migration_manager.should_migrate_cache(
                cache_info,
                target_node
            )

            if need_migration:
                self.migration_manager.trigger_migration(cache_info, target_node)

        # 5. 生成调度决策
        decision = ScheduleDecision(
            target_node=target_node,
            use_prefix_cache=True,
            cache_location=cache_info.locations[0],
            cache_length=cache_info.prefix_length,
            need_migration=need_migration,
            estimated_latency=self._estimate_latency(request, cache_info)
        )

        # 6. 记录调度决策
        self._log_decision(request, decision)

        return decision

    def _estimate_latency(
        self,
        request: PrefillRequest,
        cache_info: Optional[PrefixCacheInfo]
    ) -> float:
        """
        估算请求延迟

        延迟 = 队列等待 + KVCache 加载 + 模型计算
        """
        # 队列等待时间
        queue_latency = self.get_queue_latency(target_node)

        # KVCache 加载时间
        if cache_info:
            cache_load_latency = self._estimate_cache_load_time(cache_info)
        else:
            cache_load_latency = 0

        # 模型计算时间
        compute_latency = self._estimate_compute_time(
            request,
            cache_info.prefix_length if cache_info else 0
        )

        total_latency = queue_latency + cache_load_latency + compute_latency

        return total_latency

# 文件: conductor/scheduler/decode_scheduler.py (推测)
class DecodeScheduler:
    def schedule(self, request: DecodeRequest) -> ScheduleDecision:
        """
        为 Decode 请求做调度决策

        Decode 调度更简单:
        1. 检查 KVCache 位置
        2. 选择靠近 KVCache 的 Decode 节点
        3. 考虑负载均衡
        """
        # 1. 获取 KVCache 位置 (Prefill 阶段已创建)
        cache_location = self.cache_manager.get_cache_location(
            request.session_id
        )

        if not cache_location:
            raise ValueError(f"KVCache not found for session {request.session_id}")

        # 2. 选择靠近 KVCache 的 Decode 节点
        target_node = self.select_node_near_cache(
            cache_locations=[cache_location],
            required_memory=request.estimate_memory(),
            task_type="decode"
        )

        # 3. 生成调度决策
        decision = ScheduleDecision(
            target_node=target_node,
            use_prefix_cache=True,
            cache_location=cache_location,
            cache_length=request.current_length,
            need_migration=False,
            estimated_latency=self._estimate_decode_latency(request)
        )

        return decision

class SchedulerMetrics:
    """调度器指标"""

    def __init__(self):
        # 缓存命中率
        self.cache_hit_rate = Gauge("cache_hit_rate", "Cache hit rate")

        # 调度延迟
        self.schedule_latency = Histogram(
            "schedule_latency_ms",
            "Scheduling latency",
            buckets=[0.1, 0.5, 1.0, 2.0, 5.0]
        )

        # 节点负载
        self.node_load = Gauge(
            "node_load",
            "Node load (queue length)",
            ["node_id", "task_type"]
        )

        # 迁移次数
        self.migration_count = Counter(
            "cache_migration_total",
            "Total cache migrations"
        )

    def record_schedule_decision(self, decision: ScheduleDecision):
        """记录调度决策"""
        # 记录缓存使用
        if decision.use_prefix_cache:
            self.cache_hit_rate.inc()

        # 记录迁移
        if decision.need_migration:
            self.migration_count.inc()

# 在 config.yaml 中配置
logging:
  level: DEBUG
  modules:
    - conductor.scheduler
    - conductor.cache_manager
    - conductor.migration_manager

# 添加调度决策分析工具
def analyze_schedule_decision(decision: ScheduleDecision, request: PrefillRequest):
    """
    分析调度决策的合理性
    """
    print(f"=== Schedule Decision Analysis ===")
    print(f"Request ID: {request.id}")
    print(f"Prompt length: {len(request.prompt_tokens)}")
    print(f"Target node: {decision.target_node}")
    print(f"Use cache: {decision.use_prefix_cache}")

    if decision.use_prefix_cache:
        hit_ratio = decision.cache_length / len(request.prompt_tokens)
        print(f"Cache hit ratio: {hit_ratio:.2%}")
        print(f"Cache location: {decision.cache_location}")
        print(f"Need migration: {decision.need_migration}")

    print(f"Estimated latency: {decision.estimated_latency:.2f}ms")

    # 检查是否有更优决策
    alternative_nodes = get_all_nodes("prefill")
    for node in alternative_nodes[:3]:
        alt_latency = estimate_latency_for_node(request, node)
        print(f"  Alternative node {node}: {alt_latency:.2f}ms")

# 1. 检查 Trie 树状态
cache_stats = cache_manager.get_statistics()
print(f"Total cache entries: {cache_stats.total_entries}")
print(f"Trie tree depth: {cache_stats.tree_depth}")

# 2. 分析未命中原因
for request in recent_requests:
    prefix = cache_manager.find_longest_prefix(request.prompt_tokens)
    if not prefix:
        print(f"No match: {request.prompt_tokens[:10]}...")
    else:
        hit_ratio = prefix.prefix_length / len(request.prompt_tokens)
        print(f"Partial match: hit_ratio={hit_ratio:.2%}")

# 调整配置
config.yaml:
  scheduler:
    max_queue_length: 10       # 降低队列上限,强制分散
    migration_threshold: 0.3   # 降低迁移阈值,增加迁移频率