鸽子的回忆录

# launch/slam.launch.py
from launch import LaunchDescription
from launch_ros.actions import Node

def generate_launch_description():
    return LaunchDescription([
        Node(
            package='slam_toolbox',
            executable='async_slam_toolbox_node',
            name='slam_toolbox',
            output='screen',
            parameters=[{
                'use_sim_time': False,
                'slam_toolbox.map_file_name': '',
                'slam_toolbox.mode': 'mapping',
                'slam_toolbox.scan_topic': '/scan',
                'slam_toolbox.publish_map': True,
            }],
        ),
    ])

import numpy as np
from typing import List, Tuple

def generate_coverage_route(
    map_data: np.ndarray,
    resolution: float,
    origin: Tuple[float, float],
    line_spacing: float = 1.0,
) -> List[Tuple[float, float]]:
    traversable = map_data == 0
    rows, cols = np.where(traversable)
    min_row, max_row = rows.min(), rows.max()
    min_col, max_col = cols.min(), cols.max()

    def to_world(row, col):
        x = origin[0] + col * resolution
        y = origin[1] + row * resolution
        return (x, y)

    route = []
    direction = 1
    current_row = min_row

    while current_row <= max_row:
        if direction == 1:
            cols_range = range(min_col, max_col + 1)
        else:
            cols_range = range(max_col, min_col - 1, -1)

        for col in cols_range:
            if traversable[current_row, col]:
                route.append(to_world(current_row, col))

        current_row += int(line_spacing / resolution)
        direction *= -1

    return route

import asyncio
import json
import websockets

class RoverEdgeClient:
    def __init__(self, server_url: str):
        self.server_url = server_url
        self._ws = None

    async def connect(self):
        while True:
            try:
                async with websockets.connect(self.server_url) as ws:
                    self._ws = ws
                    logger.success("已连接到服务端")
                    await self._message_handler()
            except Exception as e:
                logger.error(f"连接失败: {e}")
                await asyncio.sleep(5)

    async def _message_handler(self):
        async for message in self._ws:
            data = json.loads(message)
            action = data.get("action")
            if action == "navigate":
                await self._handle_navigate(data["data"])
            elif action == "start_mapping":
                await self._handle_start_mapping()
            elif action == "stop_mapping":
                await self._handle_stop_mapping()
            elif action == "update_config":
                await self._handle_update_config(data["data"])

    async def _handle_navigate(self, payload):
        target_x = payload["x"]
        target_y = payload["y"]
        goal_pose = self._create_goal_pose(target_x, target_y)
        self._nav_client.send_goal(goal_pose)
        await self._report_navigation_status("navigating")

    async def report_position(self, x, y, theta):
        await self._ws.send(json.dumps({
            "action": "position_update",
            "data": {"x": x, "y": y, "theta": theta}
        }))

    async def report_anomaly(self, anomaly_type, details):
        await self._ws.send(json.dumps({
            "action": "anomaly",
            "data": {"type": anomaly_type, "details": details}
        }))

class AnomalyDetector:
    def __init__(self):
        self._anomaly_types = {
            "low_battery": self._check_low_battery,
            "stuck": self._check_stuck,
            "motor_stall": self._check_motor_stall,
            "lidar_timeout": self._check_lidar_timeout,
            "depth_stream_error": self._check_depth_stream,
            "nav2_timeout": self._check_nav2_timeout,
        }

    def check(self, rover_state: dict) -> List[dict]:
        anomalies = []
        for anomaly_type, checker in self._anomaly_types.items():
            if checker(rover_state):
                anomalies.append({
                    "type": anomaly_type,
                    "position": rover_state["position"],
                    "timestamp": time.time(),
                    "details": rover_state,
                })
        return anomalies

    def _check_low_battery(self, state: dict) -> bool:
        return state.get("battery_voltage", 0) < 3.3

    def _check_stuck(self, state: dict) -> bool:
        if state.get("speed", 0) > 0.01:
            self._last_moving_time = time.time()
        return time.time() - self._last_moving_time > 30

import aiohttp
import asyncio
import json

class NodeWebSocketClient:
    def __init__(self):
        self._session: aiohttp.ClientSession = None
        self._ws: aiohttp.ClientWebSocketResponse = None

    async def run_forever(self) -> None:
        while True:
            try:
                self._session, auth_status, access_token = await self._authenticate()
                if not auth_status:
                    await asyncio.sleep(5)
                    continue

                async with self._session.ws_connect(
                    self._ws_url(access_token),
                    autoping=True,
                ) as ws:
                    logger.success("WebSocket 已连接")
                    self._ws = ws
                    await self._message_handler()

            except aiohttp.ClientError as err:
                logger.error(f"WebSocket 连接失败: {err}")
            finally:
                self._ws = None
                await asyncio.sleep(5)

async def _authenticate(self):
    async with aiohttp.ClientSession() as session:
        async with session.post(self._auth_url, json=self._auth_data) as resp:
            data = await resp.json()
            if resp.status == 200:
                return session, True, data["access_token"]
            return session, False, None

import hashlib
import hmac

def generate_signature(node_token: str, action: str, payload: dict) -> str:
    message = f"{action}:{json.dumps(payload, sort_keys=True)}"
    signature = hmac.new(
        node_token.encode(),
        message.encode(),
        hashlib.sha256,
    ).hexdigest()
    return signature

async def websocket_send_json(self, data: dict) -> None:
    action = data.get("action", "")
    payload = data.get("data", {})

    if self._needs_signature(action):
        data["_sign"] = generate_signature(self.node_token(), action, payload)

    await self._ws.send_str(json.dumps(data))

def verify_signature(node_token: str, action: str, payload: dict, signature: str) -> bool:
    expected = generate_signature(node_token, action, payload)
    return hmac.compare_digest(expected, signature)

class NodeActionDispatcher:
    def __init__(self, client: NodeWebSocketClient):
        self._client = client
        self._handlers: Dict[str, Callable] = {
            "execute_command": self._handle_execute_command,
            "download_file": self._handle_download_file,
            "update_config": self._handle_update_config,
            "restart": self._handle_restart,
        }

    async def dispatch(self, action: str, payload: dict, raw_data: dict) -> bool:
        handler = self._handlers.get(action)
        if handler is None:
            logger.warning(f"未知的 action: {action}")
            return True

        try:
            await handler(payload)
        except Exception as e:
            logger.error(f"处理 action {action} 失败: {e}")
            await self._client.websocket_send_json({
                "action": f"{action}_response",
                "data": {"success": False, "error": str(e)},
            })

        return True

    async def _handle_execute_command(self, payload: dict):
        command = payload.get("command")
        timeout = payload.get("timeout", 30)

        if not self._check_command_policy(command):
            raise ValueError(f"命令被安全策略禁止: {command}")

        result = await self._execute_command(command, timeout)

        await self._client.websocket_send_json({
            "action": "execute_command_response",
            "data": {"success": True, "output": result},
        })

async def _message_handler(self) -> None:
    while self._ws and not self._ws.closed:
        msg = await self._ws.receive()

        match msg.type:
            case aiohttp.WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                except json.JSONDecodeError:
                    logger.warning("收到非法 JSON")
                    continue

                action = data.get("action")
                payload = data.get("data")

                if not isinstance(action, str) or not action:
                    logger.warning("收到缺少 action 的消息")
                    continue

                should_continue = await self._dispatcher.dispatch(action, payload, data)
                if not should_continue:
                    return

            case aiohttp.WSMsgType.CLOSE:
                logger.info("连接已断开")
                return

class CommandPolicy:
    def __init__(self, config: dict):
        self._disable_list = config.get("disable_command_list", "")
        self._strict = config.get("strict", True)

    def check(self, command: str) -> bool:
        disabled = [c.strip() for c in self._disable_list.split(",") if c.strip()]

        for pattern in disabled:
            if pattern == "*":
                return False
            if pattern.startswith("*") and command.endswith(pattern[1:]):
                return False
            if pattern.endswith("*") and command.startswith(pattern[:-1]):
                return False
            if command == pattern:
                return False

        return True

class NodeRuntimeServices:
    def __init__(self, client, session, connection, access_token, config_getter):
        self.client = client
        self.session = session
        self.connection = connection
        self.access_token = access_token
        self.config_getter = config_getter

    @classmethod
    def create(cls, client, session, connection, access_token, config_getter, command_policy):
        services = cls(client, session, connection, access_token, config_getter)
        services.terminal_service = TerminalService(session, connection)
        services.file_service = FileService(session, connection)
        services.monitor_service = MonitorService(session, connection)
        return services

    async def close(self):
        await self.terminal_service.close()
        await self.file_service.close()
        await self.monitor_service.close()

class AutomationEngine:
    def __init__(self):
        self._condition_last_results: Dict[str, bool] = {}

    def _evaluate_conditions(self, rule: AutomationRule, event: TriggerEvent) -> bool:
        if not rule.conditions:
            return True

        enabled_conditions = [c for c in rule.conditions if c.enabled]
        if not enabled_conditions:
            return True

        all_pass = True
        for condition in enabled_conditions:
            if condition.condition_type == ConditionType.RANGE:
                if not self._check_range_condition(condition, event):
                    all_pass = False
                    break
            elif condition.condition_type == ConditionType.COMPARE:
                if not self._check_compare_condition(condition, event):
                    all_pass = False
                    break

        last_result = self._condition_last_results.get(rule.rule_id, False)
        self._condition_last_results[rule.rule_id] = all_pass

        if not all_pass:
            return False

        if last_result:
            # 持续满足，跳过
            return False

        # False → True，触发
        return True

class ConditionType(str, Enum):
    RANGE = "range"
    COMPARE = "compare"
    TIME_RANGE = "time_range"
    LOGIC = "logic"
    EXPRESSION = "expression"
    IF_ELSE = "if_else"

# True / False 各自走哪个输出端口
CONDITION_TRUE_OUTPUT = {
    ConditionType.RANGE: "inRange",
    ConditionType.COMPARE: "result",
    ConditionType.IF_ELSE: "true",
}
CONDITION_FALSE_OUTPUT = {
    ConditionType.RANGE: "outRange",
    ConditionType.IF_ELSE: "false",
}

def _resolve_branched_actions(self, rule, event, enabled_conditions):
    condition_results = {}
    for cond in enabled_conditions:
        if cond.condition_type == ConditionType.RANGE:
            condition_results[cond.condition_id] = self._check_range_condition(cond, event)
        elif cond.condition_type == ConditionType.COMPARE:
            condition_results[cond.condition_id] = self._check_compare_condition(cond, event)

    branch_map = {}
    for cid, result in condition_results.items():
        cond = cond_by_id.get(cid)
        ctype = cond.condition_type if cond else ConditionType.EXPRESSION
        if result:
            branch_map[cid] = CONDITION_TRUE_OUTPUT.get(ctype, "result")
        else:
            false_output = CONDITION_FALSE_OUTPUT.get(ctype)
            if false_output:
                branch_map[cid] = false_output

    # 分支也要做边沿检测，否则会持续触发
    state_key = str(sorted(branch_map.items()))
    cache_key = f"branch:{rule.rule_id}"
    last_state = self._condition_last_results.get(cache_key)
    self._condition_last_results[cache_key] = state_key

    if state_key == last_state:
        return None

    executable = []
    for action in rule.actions:
        branch_info = action.config.get("_condition_branch")
        if not branch_info:
            executable.append(action)  # 不挂分支的动作总是执行
            continue
        cid = branch_info.get("condition_id")
        output = branch_info.get("output")
        if cid in branch_map and branch_map[cid] == output:
            executable.append(action)

    return executable

def _bootstrap_startup_entity_evaluation(self):
    if not self.device_manager:
        return

    state_machine = getattr(self.device_manager, "_state_machine", None)
    if not state_machine:
        return

    watch_entity_ids = self._collect_startup_watch_entity_ids()
    if not watch_entity_ids:
        return

    all_states = state_machine.get_all()

    for state_entity_id, state in all_states.items():
        if not state or not getattr(state, "attributes", None):
            continue

        extra = getattr(state.attributes, "extra", None) or {}
        raw_entity_db_id = extra.get("entity_db_id")
        if raw_entity_db_id is None:
            continue

        entity_db_id = int(raw_entity_db_id)
        if entity_db_id not in watch_entity_ids:
            continue

        new_value = getattr(state, "state", None)
        scale = extra.get("scale", 1.0)

        self._on_trigger(TriggerEvent(
            event_type="entity_value_changed",
            entity_id=entity_db_id,
            old_value=None,  # 启动时没有旧值
            new_value=new_value,
            scale=scale,
            additional_data={"bootstrap": True},
        ))

class AutomationEngine:
    def __init__(self):
        self._executing_rules: set = set()

    def _on_trigger(self, event: TriggerEvent):
        with self._lock:
            for rule in self.rules.values():
                if not rule.enabled:
                    continue
                if rule.rule_id in self._executing_rules:
                    continue
                # ... 匹配逻辑

    def _execute_rule_async(self, rule, context, event):
        rule_id = rule.rule_id
        with self._lock:
            self._executing_rules.add(rule_id)
        try:
            for action in actions:
                await self._execute_action(action, context)
        finally:
            with self._lock:
                self._executing_rules.discard(rule_id)

class BaseTrigger:
    def __init__(self):
        self._last_trigger_time: float = 0
        self._debounce_until: float = 0

    def check_debounce(self) -> bool:
        now = time.time()
        if now < self._debounce_until:
            return False
        self._debounce_until = now + self.debounce_seconds
        return True

class AutomationEngine:
    def _check_cooldown(self, rule: AutomationRule) -> bool:
        stats = self.execution_stats.get(rule.rule_id, {})
        last_execution = stats.get('last_execution_at')
        if last_execution is None:
            return True
        return time.time() - last_execution >= rule.cooldown_seconds

    def _check_daily_limit(self, rule: AutomationRule) -> bool:
        if rule.daily_limit is None:
            return True
        stats = self.execution_stats.get(rule.rule_id, {})
        return stats.get('total_executions', 0) < rule.daily_limit

def _rebuild_condition_branches(self, rule: AutomationRule):
    nodes_config = rule.nodes_config
    if not nodes_config:
        return

    connections = nodes_config.get('connections', [])

    node_to_condition = {}
    for cond in rule.conditions:
        node_id = cond.config.get("_node_id") if cond.config else None
        if node_id:
            node_to_condition[node_id] = cond.condition_id

    node_to_action = {}
    for action in rule.actions:
        node_id = action.config.get("_node_id") if action.config else None
        if node_id:
            node_to_action[node_id] = action.action_id

    for conn in connections:
        from_node = conn.get('source')
        from_output = conn.get('sourceOutput')
        to_node = conn.get('target')

        if from_node in condition_node_set and to_node in action_node_set:
            backend_action_id = node_to_action[to_node]
            backend_condition_id = node_to_condition[from_node]
            for action in rule.actions:
                if action.action_id == backend_action_id:
                    action.config["_condition_branch"] = {
                        "condition_id": backend_condition_id,
                        "output": from_output,
                    }

from enum import Enum
from typing import Dict, Type

class DeviceConfigType(str, Enum):
    MODBUS = "modbus"
    MQTT = "mqtt"
    ESPHOME_NATIVE = "esphome_native"
    ESPHOME_MQTT = "esphome_mqtt"
    OMG_GATEWAY = "omg_gateway"

class DeviceManager:
    def __init__(self):
        self.drivers: Dict[DeviceConfigType, BaseDriver] = {
            DeviceConfigType.MODBUS: ModbusDriver(),
            DeviceConfigType.MQTT: MQTTDriver(),
            DeviceConfigType.ESPHOME_NATIVE: ESPHomeNativeDriver(),
            DeviceConfigType.ESPHOME_MQTT: ESPHomeMQTTDriver(),
            DeviceConfigType.OMG_GATEWAY: OMGDriver(),
        }

class BaseDriver(ABC):
    @abstractmethod
    async def connect(self, config: dict) -> bool: ...

    @abstractmethod
    async def read_entities(self, config: dict) -> list: ...

    @abstractmethod
    async def write_entity(self, config: dict, entity_id: str, value: any) -> bool: ...

    @abstractmethod
    async def disconnect(self): ...

from dataclasses import dataclass
from typing import Any, Optional

@dataclass
class Entity:
    entity_id: str
    device_id: int
    name: str
    entity_type: str        # sensor / switch / light / button
    value: Any
    unit: Optional[str]
    readable: bool
    writable: bool
    last_update: float

temp_entity = device_manager.get_entity("greenhouse_1.temperature")
print(f"当前温度: {temp_entity.value} {temp_entity.unit}")

relay_entity = device_manager.get_entity("greenhouse_1.relay_1")
await device_manager.write_entity(relay_entity.entity_id, True)

from dataclasses import dataclass
import time

@dataclass
class CachedMappingInfo:
    entity_id: int
    entity_type: str
    scale: float
    register_address: Optional[int]   # Modbus
    mqtt_topic: Optional[str]         # MQTT
    esphome_key: Optional[str]        # ESPHome

class DeviceManager:
    def __init__(self):
        self.mapping_cache: Dict[int, Dict[int, CachedMappingInfo]] = {}
        self.config_cache: Dict[int, CachedDeviceConfig] = {}
        self._cache_ttl = 300
        self._cache_timestamps: Dict[str, float] = {}

    def get_mapping(self, device_id: int, entity_id: int):
        cache_key = f"mapping_{device_id}"
        if self._is_cache_expired(cache_key):
            self._refresh_mapping_cache(device_id)
        return self.mapping_cache.get(device_id, {}).get(entity_id)

    def _is_cache_expired(self, cache_key: str) -> bool:
        timestamp = self._cache_timestamps.get(cache_key, 0)
        return time.time() - timestamp > self._cache_ttl

from concurrent.futures import ThreadPoolExecutor, Future

class DeviceManager:
    def __init__(self):
        self._modbus_port_workers: Dict[str, ThreadPoolExecutor] = {}
        self._modbus_port_futures: Dict[str, Future] = {}

    def _get_port_worker(self, port_key: str) -> ThreadPoolExecutor:
        if port_key not in self._modbus_port_workers:
            self._modbus_port_workers[port_key] = ThreadPoolExecutor(
                max_workers=1,
                thread_name_prefix=f"modbus-{port_key}",
            )
        return self._modbus_port_workers[port_key]

    def poll_device(self, device: Device):
        if device.config_type == DeviceConfigType.MODBUS:
            port_key = self._get_modbus_port_key(device)
            worker = self._get_port_worker(port_key)
            future = worker.submit(self._poll_modbus_device, device)
            self._modbus_port_futures[port_key] = future
        else:
            self._poll_device(device)

    def _get_modbus_port_key(self, device: Device) -> str:
        config = device.modbus_config
        if config.serial_port:
            return f"serial_{config.serial_port}"
        return f"tcp_{config.host}:{config.port}"

from concurrent.futures import ThreadPoolExecutor

class DeviceManager:
    def __init__(self):
        self._startup_preload_executor = ThreadPoolExecutor(
            max_workers=8,
            thread_name_prefix="startup-preload",
        )
        self._startup_preload_futures: Dict[int, Future] = {}

    def connect_all_devices(self):
        devices = Device.select().where(Device.enabled == True)
        for device in devices:
            future = self._startup_preload_executor.submit(
                self._connect_and_preload, device
            )
            self._startup_preload_futures[device.id] = future

    def _connect_and_preload(self, device: Device):
        try:
            driver = self.drivers[device.config_type]
            driver.connect(device.config)
            entities = driver.read_entities(device.config)
            for entity in entities:
                self._update_entity_state(device.id, entity)
        except Exception as e:
            logger.error(f"设备 {device.name} 连接失败: {e}")

from enum import Enum
from typing import Callable, Dict, List

class EventType(str, Enum):
    STATE_CHANGED = "state_changed"
    DEVICE_CONNECTED = "device_connected"
    DEVICE_DISCONNECTED = "device_disconnected"

class EventBus:
    def __init__(self):
        self._listeners: Dict[EventType, List[Callable]] = {}

    def listen(self, event_type: EventType, callback: Callable):
        if event_type not in self._listeners:
            self._listeners[event_type] = []
        self._listeners[event_type].append(callback)

        def unsubscribe():
            self._listeners[event_type].remove(callback)
        return unsubscribe

    def emit(self, event_type: EventType, data: dict):
        for callback in self._listeners.get(event_type, []):
            try:
                callback(data)
            except Exception as e:
                logger.error(f"事件处理失败: {e}")

import gymnasium as gym
from gymnasium import spaces

class GameExamEnv(gym.Env):
    def __init__(self):
        super().__init__()
        self.global_dim = 60          # 全局状态维度
        self.action_feature_dim = 100 # 每个动作的特征维度
        self.max_actions = 50         # 手牌 + 饮料 + 结束回合

        self.observation_space = spaces.Dict({
            'global': spaces.Box(-20.0, 20.0, shape=(self.global_dim,), dtype=np.float32),
            'action_features': spaces.Box(
                -20.0, 20.0,
                shape=(self.max_actions, self.action_feature_dim),
                dtype=np.float32,
            ),
            'action_mask': spaces.Box(0.0, 1.0, shape=(self.max_actions,), dtype=np.float32),
        })

def _global_observation(self) -> np.ndarray:
    state = self.runtime.state

    global_values = {
        # 进度
        'step_ratio': state['step'] / max(self.scenario.steps, 1),
        'remaining_step_ratio': max(self.scenario.steps - state['step'], 0) / max(self.scenario.steps, 1),
        'audition_progress': state['audition_index'] / max(len(self.scenario.audition_sequence), 1),

        # 资源
        'stamina_ratio': state['stamina'] / max(state['max_stamina'], 1.0),
        'produce_point_ratio': state['produce_points'] / 150.0,
        'fan_vote_ratio': state['fan_votes'] / 5000.0,

        # 三维参数
        'vocal_ratio': state['vocal'] / parameter_scale,
        'dance_ratio': state['dance'] / parameter_scale,
        'visual_ratio': state['visual'] / parameter_scale,

        # 卡组
        'deck_quality': state['deck_quality'] / 20.0,
        'drink_quality': state['drink_quality'] / 10.0,
        'deck_size_ratio': len(self.runtime.deck) / 40.0,
    }

    return np.array(
        [float(global_values[name]) for name in self.global_feature_names],
        dtype=np.float32,
    )

def _candidate_feature(self, candidate) -> np.ndarray:
    action_type_vec = self.taxonomy.encode_actions([candidate.action_type])
    effect_types = self._produce_effect_types(candidate)
    effect_vec = self.taxonomy.encode_produce_effects(effect_types)

    numeric = np.array([
        candidate.stamina_delta / max(state['max_stamina'], 1.0),
        candidate.produce_point_delta / 100.0,
        candidate.success_probability,
        len(candidate.produce_effect_ids) / 8.0,
        1.0 if candidate.available else 0.0,
        state['stamina'] / max(state['max_stamina'], 1.0),
    ], dtype=np.float32)

    return np.concatenate([action_type_vec, effect_vec, numeric]).astype(np.float32)

def action_masks(self) -> np.ndarray:
    return np.array([
        bool(candidate.payload.get('available', False))
        for candidate in self._candidates
    ], dtype=bool)

CURRICULUM_STAGES = [
    "初中间考试",
    "初最终考试",
    "NIA中间考试",
    "NIA最终考试",
    "NIA选拔",
    "初Regular全流程",
    "初Master全流程",
    "NIA Pro全流程",
    "NIA Master全流程",
]

def run_curriculum(self, stages, timesteps_per_stage=131072):
    for stage_name in stages:
        env = self._create_env_for_stage(stage_name)
        model = MaskablePPO("MultiInputPolicy", env, verbose=1)
        model.learn(total_timesteps=timesteps_per_stage)
        quality = self._evaluate_model(model, stage_name)
        self._save_checkpoint(model, stage_name, quality)

def self_bootstrap(self, iterations=3):
    for iteration in range(iterations):
        # 多 checkpoint 跑固定种子
        candidates = []
        for checkpoint in self.checkpoints:
            for seed in self.evaluation_seeds:
                episode = self._run_episode(checkpoint, seed)
                candidates.append(episode)

        # 排序选最优
        candidates.sort(key=lambda c: c.quality_key(), reverse=True)
        best_trajectories = candidates[:10]

        # BC 蒸馏
        bc_model = self._bc_train(best_trajectories)

        # RL 微调
        rl_model = self._rl_finetune(bc_model, timesteps=50000)

        # 谁稳留谁
        if self._evaluate(rl_model) > self._evaluate(bc_model):
            self.best_model = rl_model
        else:
            self.best_model = bc_model

@dataclass
class EpisodeCandidate:
    checkpoint_path: Path
    seed: int
    records: list
    total_reward: float
    terminal_score: float
    invalid_actions: int
    clear_rank: int

    def quality_key(self):
        return (
            -int(self.invalid_actions),
            int(self.clear_rank),
            float(self.terminal_score),
            float(self.total_reward),
        )

from fastapi import FastAPI
from sb3_contrib import MaskablePPO

app = FastAPI()
model = MaskablePPO.load("best_model.zip")

@app.post("/api/inference/predict")
async def predict(request: PredictRequest):
    obs = build_observation(request.state, request.legal_actions)
    action_mask = build_action_mask(request.legal_actions)

    action_value, _ = model.predict(
        obs,
        deterministic=True,
        action_masks=action_mask,
    )

    action_index = int(action_value)
    return {
        "action_index": action_index,
        "action_id": request.legal_actions[action_index]["action_id"],
        "confidence": float(action_value),
    }

class RLInferenceClient:
    def __init__(self, base_url="http://127.0.0.1:8100"):
        self._base_url = base_url
        self._session = requests.Session()

    def predict(self, exam_state, legal_actions, deterministic=True):
        payload = {
            **exam_state,
            "legal_actions": legal_actions,
            "deterministic": deterministic,
        }
        try:
            resp = self._session.post(
                f"{self._base_url}/api/inference/predict",
                json=payload,
                timeout=10.0,
            )
            resp.raise_for_status()
            return resp.json()
        except requests.RequestException as exc:
            logger.warning(f"RL 推理请求失败: {exc}")
            return None

def build_reward_config(mode: str) -> RewardConfig:
    if mode == "score":
        # 拼分数
        return RewardConfig(
            score_weight=1.0,
            clear_bonus=0.5,
            invalid_action_penalty=-0.25,
        )
    elif mode == "clear":
        # 保过线
        return RewardConfig(
            score_weight=0.3,
            clear_bonus=2.0,
            overclear_penalty=-0.1,  # 过线后再硬冲会被惩罚
            invalid_action_penalty=-0.25,
        )

def learn_card(self, app, card_frame, card_list):
    for card in card_list:
        # 先试 CLIP
        existing_id = self._try_clip_identify(app, card.frame)
        if existing_id:
            continue

        # CLIP 不认识，OCR 上
        learned_id = self._learn_via_ocr(app, card.frame)
        if learned_id:
            # OCR 找到了，让 CLIP 把这张图记下来
            self.clip_manager.add_to_memory(card.frame, learned_id)

def _try_clip_identify(self, app, card_frame):
    try:
        result = self.clip_manager.retrieve(card_frame)
        if result is not None:
            return result.payload.id
    except Exception as e:
        logger.debug(f"CLIP identify failed: {e}")
    return None

def _learn_via_ocr(self, app, card_image):
    ocr_result = ocr_service.ocr(card_image)
    if not ocr_result or not ocr_result.results:
        return None

    for item in ocr_result.results:
        if len(item.text) >= 3:
            status, db_result = database.search(item.text)
            if status and db_result:
                return db_result.id
    return None

def identify_element(self, screenshot, element_frame):
    # 置信度够高就直接用
    result = self.clip_manager.retrieve(element_frame)
    if result is not None and result.similarity > 0.96:
        return result.payload

    # 没命中，那就触发一次学习
    self.learn_element(screenshot, element_frame)

    # 再试一次
    result = self.clip_manager.retrieve(element_frame)
    if result is not None:
        return result.payload

    return None

EVENT_ROUTING = "routing"           # 路由分类结果
EVENT_TOOL_STATUS = "tool_status"   # 工具调用状态
EVENT_TOOL_RESULT = "tool_result"   # 工具执行结果
EVENT_CONTENT = "content"           # 正文内容

yield {
    "type": "routing",
    "data": {
        "category": "device_management",
        "confidence": 0.92,
        "reasoning": "用户询问设备状态",
        "secondary_category": None,
    },
}

yield {
    "type": "tool_status",
    "data": {"name": "list_devices", "status": "calling"},
}

yield {
    "type": "tool_status",
    "data": {"name": "list_devices", "status": "completed"},
}

yield {
    "type": "tool_result",
    "data": {
        "name": "list_devices",
        "success": True,
        "data": {"devices": [...]},
    },
}

yield {
    "type": "content",
    "data": {"content": "当前系统中共有"},
}
yield {
    "type": "content",
    "data": {"content": " 12 个设备"},
}

@classmethod
def _extract_stream_chunk_deltas(cls, chunk) -> Dict[str, str]:
    content_delta = cls._message_content_to_text(getattr(chunk, "content", ""))
    thinking_parts = []

    raw_content = getattr(chunk, "content", None)
    if isinstance(raw_content, list):
        for block in raw_content:
            if not isinstance(block, dict):
                continue
            block_type = str(block.get("type", "")).lower()
            if block_type in {"reasoning", "thinking", "reasoning_content"}:
                text = cls._extract_reasoning_text(block)
                if text:
                    thinking_parts.append(text)

    for container in (
        getattr(chunk, "additional_kwargs", None),
        getattr(chunk, "response_metadata", None),
    ):
        if not isinstance(container, dict):
            continue
        for key in ("reasoning_content", "thinking", "reasoning"):
            if key in container:
                text = cls._extract_reasoning_text(container.get(key))
                if text:
                    thinking_parts.append(text)

    return {
        "content_delta": content_delta,
        "thinking_delta": "".join(thinking_parts),
    }

yield {
    "type": "thinking",
    "data": {"content": "用户询问设备状态，我需要调用 list_devices 工具..."},
}

import asyncio
from enum import Enum

class RequestStatus(str, Enum):
    QUEUED = "queued"
    PROCESSING = "processing"
    TOOL_CALLING = "tool_calling"
    COMPLETED = "completed"
    FAILED = "failed"

class QueueManager:
    def __init__(self, max_concurrent: int = 10):
        self._semaphore = asyncio.Semaphore(max_concurrent)
        self._active_requests: Dict[str, RequestTask] = {}

    async def submit(self, request_id: str, handler) -> RequestTask:
        task = RequestTask(request_id=request_id)
        self._active_requests[request_id] = task

        async def _execute():
            async with self._semaphore:
                task.status = RequestStatus.PROCESSING
                try:
                    result = await handler(task)
                    task.result = result
                    task.status = RequestStatus.COMPLETED
                except Exception as e:
                    task.error = str(e)
                    task.status = RequestStatus.FAILED

        asyncio.create_task(_execute())
        return task

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import json

app = FastAPI()

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatRequest):
    if request.stream:
        return StreamingResponse(
            stream_generator(request),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            },
        )
    result = await agent.chat(request.messages)
    return result

async def stream_generator(request):
    async for event in agent.chat_stream(request.messages):
        yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
    yield "data: [DONE]\n\n"

from langgraph.prebuilt import create_react_agent

agent = create_react_agent(model, tools, prompt=SystemMessage(content=system_prompt))
result = await agent.ainvoke({"messages": [HumanMessage(content=user_message)]})

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

def _build_agent_graph(self, model, tools, system_prompt):
    tool_node = ToolNode(tools, awrap_tool_call=self._wrap_tool_call)
    finalize_model = model.bind(tool_choice="none")

    async def agent_node(state): ...
    async def precheck_node(state): ...
    async def enforce_tool_node(state): ...
    async def finalize_node(state): ...
    async def update_flags_node(state): ...

    graph = StateGraph(AgentState)
    graph.add_node("agent", agent_node)
    graph.add_node("precheck", precheck_node)
    graph.add_node("run_tools", tool_node)
    graph.add_node("enforce_tool", enforce_tool_node)
    graph.add_node("finalize", finalize_node)
    graph.add_node("update_flags", update_flags_node)

    graph.add_conditional_edges("agent", after_agent, {
        "precheck": "precheck",
        "enforce_tool": "enforce_tool",
        "finalize": "finalize",
    })
    graph.add_conditional_edges("precheck", after_precheck, {
        "run_tools": "run_tools",
        "agent": "agent",
    })
    graph.add_edge("run_tools", "update_flags")
    graph.add_conditional_edges("update_flags", after_update_flags, {
        "agent": "agent",
        "finalize": "finalize",
    })
    graph.add_edge("enforce_tool", "agent")
    graph.add_edge("finalize", END)
    graph.set_entry_point("agent")

    return graph.compile()

async def _wrap_tool_call(self, request, execute):
    response = await execute(request)
    if not isinstance(response, ToolMessage):
        return response

    parsed = self._parse_tool_output(getattr(response, "content", ""))
    success = getattr(response, "status", "") != "error"

    error_text = None
    if isinstance(parsed, dict):
        if "success" in parsed:
            success = bool(parsed.get("success"))
        if parsed.get("error"):
            error_text = str(parsed.get("error"))

    normalized_payload = {
        "success": success,
        "data": parsed.get("data", parsed) if isinstance(parsed, dict) else parsed,
    }
    if error_text:
        normalized_payload["error"] = error_text

    return ToolMessage(
        content=json.dumps(normalized_payload, ensure_ascii=False),
        name=getattr(response, "name", ""),
        tool_call_id=getattr(response, "tool_call_id", ""),
        status="error" if not success else "success",
    )

async def precheck_node(self, state):
    last_msg = state["messages"][-1]
    if not (isinstance(last_msg, AIMessage) and last_msg.tool_calls):
        return {"_precheck_result": None}

    tool_name = last_msg.tool_calls[0].get("name", "")
    tool_args = last_msg.tool_calls[0].get("args", {})

    precheck_err = self._precheck_tool_call(
        tool_name=tool_name,
        tool_args=tool_args,
        user_message=state.get("user_message"),
    )

    if precheck_err is None:
        return {"_precheck_result": None}

    resolved = await self._resolve_precheck_with_auto_dependency(
        precheck_error=precheck_err,
        tool_name=tool_name,
        tool_args=tool_args,
    )

    if resolved.get("dependency_tool_result") and not resolved["precheck_error"]:
        # 自动执行了依赖工具，把结果塞进消息历史
        return {"messages": [ToolMessage(
            content=str(resolved["dependency_tool_result"]),
            tool_call_id="auto_dep",
        )]}
    else:
        return {"messages": [ToolMessage(
            content=json.dumps(resolved["precheck_error"]),
            tool_call_id=tool_call_id,
        )]}

def _tool_enforcement_instruction(self, routing, user_message, available_tool_names):
    return f"""你必须调用工具来完成用户的请求，不要直接回复。

用户问题: {user_message}

可用工具: {', '.join(available_tool_names)}

请立即调用合适的工具。"""

class AgentState(TypedDict):
    tool_calls_made: int
    tool_failure_streak: int
    last_failed_tool: Optional[str]
    last_tool_error: Optional[str]
    forced_tool_rounds: int

def after_update_flags(self, state):
    # 连续失败 2 次以上，别再试了
    if state.get("tool_failure_streak", 0) >= 2:
        return "finalize"

    # 工具调用次数超过 10 次，强制收尾
    if state.get("tool_calls_made", 0) >= 10:
        return "finalize"

    return "agent"

@_tool(args_schema=SearchToolsArgs)
async def search_tools(query: str, tags: list = None) -> dict:
    """搜索当前未加载的工具"""
    results = registry.search_tools_by_keyword(query, max_results=10)

    if not results:
        return {
            "success": True,
            "data": {
                "tools": [],
                "message": f"未找到与 '{query}' 匹配的工具",
            },
        }

    return {
        "success": True,
        "data": {
            "tools": results,
            "message": f"找到 {len(results)} 个匹配工具，已动态加载",
        },
    }

from enum import Enum

class IntentCategory(str, Enum):
    PEST_DIAGNOSIS = "pest_diagnosis"        # 病虫害诊断（多模态）
    DEVICE_MANAGEMENT = "device_management"  # 设备管理
    DATA_ANALYSIS = "data_analysis"          # 数据分析
    AUTOMATION = "automation"                # 自动化规则
    AGRICULTURE = "agriculture"              # 农业管理
    GENERAL = "general"                      # 通用问答

from pydantic import BaseModel, Field

class RouterOutput(BaseModel):
    category: str = Field(description="意图分类")
    confidence: float = Field(default=0.5, description="置信度 0-1")
    reasoning: str = Field(default="", description="分类理由")
    requires_image: bool = Field(default=False, description="是否需要图片")
    sub_intent: str = Field(default="", description="子意图")

from langchain_openai import ChatOpenAI

class IntentRouter:
    def __init__(self):
        self._router_model = ChatOpenAI(
            base_url="http://localhost:11434/v1",
            model="qwen3:1.7b",
            max_tokens=256,
            temperature=0.1,
            timeout=5.0,
        )

    async def classify(self, query: str, history: list = None) -> RoutingResult:
        try:
            return await self._classify_with_llm(query, history)
        except Exception:
            return self._keyword_fallback(query)

def _keyword_fallback(self, query: str) -> RoutingResult:
    q = query.lower()

    pest_keywords = ["病虫害", "病害", "虫害", "枯萎", "黄叶", "斑点",
                     "诊断", "识别病", "什么病", "怎么治"]
    if any(kw in q for kw in pest_keywords):
        return RoutingResult(
            category=IntentCategory.PEST_DIAGNOSIS,
            confidence=0.75,
            reasoning="关键词匹配: 病虫害相关",
        )

    auto_keywords = ["自动化", "自动", "规则", "触发", "条件", "定时",
                     "创建规则", "编排", "工作流", "联动"]
    if any(kw in q for kw in auto_keywords):
        return RoutingResult(
            category=IntentCategory.AUTOMATION,
            confidence=0.7,
            reasoning="关键词匹配: 自动化相关",
        )

    device_keywords = ["设备", "传感器", "连接", "断开", "在线", "离线", "状态"]
    if any(kw in q for kw in device_keywords):
        return RoutingResult(
            category=IntentCategory.DEVICE_MANAGEMENT,
            confidence=0.7,
            reasoning="关键词匹配: 设备管理相关",
        )

    return RoutingResult(
        category=IntentCategory.GENERAL,
        confidence=0.5,
        reasoning="未匹配到特定领域关键词",
    )

def _select_pipeline(self, routing: RoutingResult):
    registry = self._tool_registry

    if routing.category == IntentCategory.DEVICE_MANAGEMENT:
        tools = registry.get_tools_by_tags(["device", "entity"])
        system_prompt = get_prompt("device_expert")

    elif routing.category == IntentCategory.DATA_ANALYSIS:
        tools = registry.get_tools_by_tags(["data", "device"])
        system_prompt = get_prompt("data_analyst")

    elif routing.category == IntentCategory.AUTOMATION:
        tools = registry.get_tools_by_tags(["automation", "device", "entity"])
        system_prompt = get_prompt("automation_expert")

    elif routing.category == IntentCategory.AGRICULTURE:
        tools = registry.get_tools_by_tags(["agriculture", "diagnosis", "data"])
        system_prompt = get_prompt("agriculture_expert")

    else:
        tools = registry.get_tools_by_tags(["general"])
        system_prompt = get_prompt("general_assistant")

    return model, tools, system_prompt

class ToolRegistry:
    def __init__(self):
        self._tools: Dict[str, ToolRegistration] = {}

    def register(self, tool, tags=None, description_zh="", examples=None):
        reg = ToolRegistration(
            tool=tool,
            tags=set(tags or []),
            description_zh=description_zh or tool.description,
            examples=examples or [],
        )
        self._tools[tool.name] = reg

    def get_tools_by_tags(self, tags: List[str]) -> List[BaseTool]:
        result_set = {}
        tag_set = set(tags)
        for name, reg in self._tools.items():
            if reg.tags & tag_set:
                result_set[name] = reg.tool
        return sorted(result_set.values(), key=lambda t: t.name)