#!/usr/bin/env python3
from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import numpy as np

HOST_ROOT = Path(__file__).resolve().parents[1]
PROJECT_ROOT = HOST_ROOT.parents[1]
if str(HOST_ROOT) not in sys.path:
    sys.path.insert(0, str(HOST_ROOT))
if str(PROJECT_ROOT / 'guguji_rl') not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT / 'guguji_rl'))

from guguji_real_bridge.action_mapper import DeployActionMapper
from guguji_real_bridge.config import load_yaml
from guguji_real_bridge.protocol import (
    COMMAND_FLAG_ARM,
    COMMAND_FLAG_CLEAR_FAULTS,
    COMMAND_FLAG_DISARM,
    COMMAND_FLAG_ESTOP,
    COMMAND_FLAG_HOLD,
    TELEMETRY_FAULT_IMU_OFFLINE,
    CommandPacket,
    TelemetryPacket,
    UdpRobotLink,
)
from guguji_real_bridge.safety import TargetLimiter


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description='Run a trained PPO policy on the real guguji robot through WiFi.')
    parser.add_argument('--deploy-config', default=str(HOST_ROOT / 'configs' / 'guguji_real_robot.yaml'))
    parser.add_argument('--rl-config', default=str(PROJECT_ROOT / 'guguji_rl' / 'configs' / 'walk_ppo.yaml'))
    parser.add_argument('--model', required=True, help='训练好的 stable-baselines3 PPO 模型 .zip 路径')
    parser.add_argument('--arm', action='store_true', help='真正使能并下发运动指令；不加时只保持失能')
    parser.add_argument('--clear-faults', action='store_true', help='启动时请求 ESP32S3 清除 RS00 故障')
    parser.add_argument('--max-seconds', type=float, default=0.0, help='运行时长，0 表示一直运行')
    parser.add_argument('--deterministic', action='store_true', help='使用确定性策略输出')
    parser.add_argument('--log-interval', type=float, default=None, help='ESP32S3 遥测日志间隔，单位秒')
    return parser.parse_args()


def load_model(model_path: str | Path):
    try:
        from stable_baselines3 import PPO
    except ImportError as exc:
        raise RuntimeError('缺少 stable-baselines3，请先安装 guguji_real_robot/host/requirements.txt') from exc
    return PPO.load(str(Path(model_path).resolve()))


def _mask_names(mask: int, names: list[str]) -> list[str]:
    return [name for index, name in enumerate(names) if (mask & (1 << index)) != 0]


def format_fault_summary(telemetry: TelemetryPacket, joint_names: list[str]) -> str:
    parts: list[str] = []
    motor_faults = _mask_names(telemetry.motor_fault_mask, joint_names)
    feedback_timeouts = _mask_names(telemetry.feedback_timeout_mask, joint_names)
    if motor_faults:
        parts.append('motor_fault=' + ','.join(motor_faults))
    if feedback_timeouts:
        parts.append('feedback_timeout=' + ','.join(feedback_timeouts))
    if telemetry.imu_offline:
        parts.append('imu_offline')
    if telemetry.mag_offline:
        parts.append('mag_offline')
    return '; '.join(parts) if parts else 'ok'


def print_telemetry_status(telemetry: TelemetryPacket, joint_names: list[str]) -> None:
    rpy_deg = np.degrees(telemetry.rpy_rad)
    mean_abs_velocity = float(np.mean(np.abs(telemetry.joint_velocity_rad_s)))
    print(
        '[ESP32S3] '
        f'state={telemetry.robot_state_name} '
        f'seq={telemetry.sequence} '
        f'uptime={telemetry.uptime_ms / 1000.0:.1f}s '
        f'cmd_age={telemetry.last_command_age_ms}ms '
        f'faults={format_fault_summary(telemetry, joint_names)} '
        f'rpy=({rpy_deg[0]:.1f},{rpy_deg[1]:.1f},{rpy_deg[2]:.1f})deg '
        f'temp_max={telemetry.max_motor_temperature_c:.1f}C '
        f'joint_vel_mean={mean_abs_velocity:.3f}rad/s'
    )


def main() -> int:
    args = parse_args()
    deploy_config = load_yaml(args.deploy_config)
    network_config = deploy_config['network']
    control_config = deploy_config['control']
    safety_config = deploy_config['safety']
    logging_config = deploy_config.get('logging', {})
    log_interval_s = args.log_interval
    if log_interval_s is None:
        log_interval_s = float(logging_config.get('status_interval_s', 1.0))

    mapper = DeployActionMapper(rl_config_path=args.rl_config, deploy_config=deploy_config['state_estimation'])
    limiter = TargetLimiter(
        max_step_rad=float(safety_config['max_joint_step_rad']),
        lower=mapper.joint_lower,
        upper=mapper.joint_upper,
    )
    model = load_model(args.model)
    link = UdpRobotLink(
        command_host=str(network_config['esp32_host']),
        command_port=int(network_config['command_port']),
        telemetry_bind_host=str(network_config.get('telemetry_bind_host', '0.0.0.0')),
        telemetry_port=int(network_config['telemetry_port']),
        timeout_s=float(network_config.get('telemetry_timeout_s', 1.0)),
    )

    kp = np.asarray(control_config['kp'], dtype=np.float32)
    kd = np.asarray(control_config['kd'], dtype=np.float32)
    zero_velocity = np.zeros(len(mapper.joint_names), dtype=np.float32)
    zero_torque = np.zeros(len(mapper.joint_names), dtype=np.float32)
    sequence = 0
    start_time = time.monotonic()
    next_tick = start_time
    first_telemetry = True
    last_log_time = 0.0
    last_state: int | None = None
    last_fault_mask: int | None = None
    last_safety_message: str | None = None

    print(
        '主机部署配置: '
        f'esp32={network_config["esp32_host"]}:{int(network_config["command_port"])} '
        f'telemetry_port={int(network_config["telemetry_port"])} '
        f'rl_config={Path(args.rl_config).resolve()} '
        f'model={Path(args.model).resolve()} '
        f'arm={args.arm}'
    )

    # 固件只有收到第一条 UDP 命令后，才知道遥测应该回传到哪个主机和端口。
    # 所以这里先发几帧“失能 + 保持”的安全包，建立回传路径。
    for _ in range(5):
        link.send_command(
            CommandPacket(
                sequence=sequence,
                positions_rad=mapper.nominal_joint_targets.astype(np.float32),
                velocities_rad_s=zero_velocity,
                kp=kp,
                kd=kd,
                torque_ff_nm=zero_torque,
                flags=COMMAND_FLAG_DISARM | COMMAND_FLAG_HOLD,
            )
        )
        sequence = (sequence + 1) & 0xFFFFFFFF
        time.sleep(0.05)

    print('等待 ESP32S3 遥测...')
    try:
        while True:
            try:
                telemetry = link.receive_telemetry()
            except TimeoutError:
                link.send_command(
                    CommandPacket(
                        sequence=sequence,
                        positions_rad=mapper.nominal_joint_targets.astype(np.float32),
                        velocities_rad_s=zero_velocity,
                        kp=kp,
                        kd=kd,
                        torque_ff_nm=zero_torque,
                        flags=COMMAND_FLAG_DISARM | COMMAND_FLAG_HOLD | COMMAND_FLAG_ESTOP,
                    )
                )
                sequence = (sequence + 1) & 0xFFFFFFFF
                print('等待遥测超时，已发送失能安全包，继续等待...')
                continue

            if first_telemetry:
                limiter.reset(telemetry.joint_position_rad)
                mapper.reset()
                first_telemetry = False
                start_time = time.monotonic()
                next_tick = start_time
                print('已收到遥测，开始部署循环。')
                print_telemetry_status(telemetry, mapper.joint_names)
                last_log_time = time.monotonic()
                last_state = telemetry.robot_state
                last_fault_mask = telemetry.fault_mask

            elapsed = time.monotonic() - start_time
            if args.max_seconds > 0 and elapsed >= args.max_seconds:
                break

            now = time.monotonic()
            should_log = (
                telemetry.robot_state != last_state
                or telemetry.fault_mask != last_fault_mask
                or (log_interval_s > 0 and now - last_log_time >= log_interval_s)
            )
            if should_log:
                print_telemetry_status(telemetry, mapper.joint_names)
                last_log_time = now
                last_state = telemetry.robot_state
                last_fault_mask = telemetry.fault_mask

            observation = mapper.build_observation(
                joint_position=telemetry.joint_position_rad,
                joint_velocity=telemetry.joint_velocity_rad_s,
                roll=telemetry.roll,
                pitch=telemetry.pitch,
            )
            action, _ = model.predict(observation, deterministic=args.deterministic)
            target = mapper.action_to_joint_targets(action)
            target = limiter.apply(target)

            flags = 0
            if args.clear_faults and sequence < 10:
                flags |= COMMAND_FLAG_CLEAR_FAULTS
            if args.arm:
                flags |= COMMAND_FLAG_ARM
            else:
                flags |= COMMAND_FLAG_DISARM | COMMAND_FLAG_HOLD

            safety_messages: list[str] = []
            if telemetry.motor_fault_mask != 0:
                flags |= COMMAND_FLAG_ESTOP
                safety_messages.append(f'检测到电机故障: {format_fault_summary(telemetry, mapper.joint_names)}')

            if args.arm and telemetry.feedback_timeout_mask != 0:
                flags |= COMMAND_FLAG_ESTOP
                safety_messages.append(f'电机反馈超时: {format_fault_summary(telemetry, mapper.joint_names)}')

            if args.arm and (telemetry.fault_mask & TELEMETRY_FAULT_IMU_OFFLINE) != 0:
                flags |= COMMAND_FLAG_ESTOP
                safety_messages.append('QMI8658A IMU 离线，roll/pitch 安全保护不可用')

            if abs(telemetry.roll) > float(safety_config['max_roll_rad']):
                flags |= COMMAND_FLAG_ESTOP
                safety_messages.append(f'roll 超限 {telemetry.roll:.3f} rad')
            if abs(telemetry.pitch) > float(safety_config['max_pitch_rad']):
                flags |= COMMAND_FLAG_ESTOP
                safety_messages.append(f'pitch 超限 {telemetry.pitch:.3f} rad')

            safety_message = '；'.join(safety_messages)
            if safety_message and safety_message != last_safety_message:
                print(f'{safety_message}，已请求急停。')
                last_safety_message = safety_message
            elif not safety_message:
                last_safety_message = None

            link.send_command(
                CommandPacket(
                    sequence=sequence,
                    positions_rad=target,
                    velocities_rad_s=zero_velocity,
                    kp=kp,
                    kd=kd,
                    torque_ff_nm=zero_torque,
                    flags=flags,
                )
            )
            sequence = (sequence + 1) & 0xFFFFFFFF

            next_tick += mapper.control_dt
            sleep_s = next_tick - time.monotonic()
            if sleep_s > 0.0:
                time.sleep(sleep_s)

    except KeyboardInterrupt:
        print('收到 Ctrl+C，发送失能保持命令后退出。')
    finally:
        hold = mapper.nominal_joint_targets.astype(np.float32)
        for _ in range(5):
            link.send_command(
                CommandPacket(
                    sequence=sequence,
                    positions_rad=hold,
                    velocities_rad_s=zero_velocity,
                    kp=kp,
                    kd=kd,
                    torque_ff_nm=zero_torque,
                    flags=COMMAND_FLAG_DISARM | COMMAND_FLAG_HOLD,
                )
            )
            sequence = (sequence + 1) & 0xFFFFFFFF
            time.sleep(0.02)
        link.close()

    return 0


if __name__ == '__main__':
    raise SystemExit(main())
