#!/usr/bin/env python3
"""按底盘动作 profile 执行动作，并同步采集现场数据。"""

from __future__ import annotations

import argparse
import math
import sys
import time
from pathlib import Path
from typing import Any

SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from chassis_action_profile import load_yaml as load_action_profile_yaml
from chassis_action_profile import validate_profile as validate_action_profile
from chassis_data_common import apply_cli_overrides, load_config, validate_config
from capture_chassis_session import (
    ChassisSessionCapture,
    config_command_source,
    config_topic,
    finalize_capture,
    make_target_paths,
    now_us,
)


PRIMITIVE_TYPE_VALUES = {
    "straight_line": "STRAIGHT_LINE",
    "arc": "ARC",
    "in_place_rotation": "IN_PLACE_ROTATION",
    "steering_sweep": "STEERING_SWEEP",
    "lateral_translation": "LATERAL_TRANSLATION",
    "diagonal_motion": "DIAGONAL_MOTION",
    "module_alignment": "MODULE_ALIGNMENT",
    "coordinated_steering": "COORDINATED_STEERING",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="按现场底盘动作 profile 顺序发送动作原语，同时采集底盘、命令和外部真值数据。"
    )
    parser.add_argument("--config", required=True, help="底盘标定现场数据配置文件。")
    parser.add_argument(
        "--action-profile",
        default="src/site_deployment/workshop_chassis_calibration_real/config/chassis_action_profile.yaml",
        help="底盘动作 profile YAML 路径。",
    )
    parser.add_argument("--chassis-type", default="", help="覆盖 chassis_type，并选择对应动作序列。")
    parser.add_argument("--task-code", default="", help="只执行指定 task_code；为空时执行该底盘类型全部动作。")
    parser.add_argument("--session-id", default="", help="覆盖 session_id。")
    parser.add_argument("--site-id", default="", help="覆盖 site_id。")
    parser.add_argument("--vehicle-id", default="", help="覆盖 vehicle_id。")
    parser.add_argument("--operator-id", default="workshop_operator", help="写入请求头的操作员 ID。")
    parser.add_argument("--workshop-host", default="workshop_chassis_profile_capture", help="写入请求头的车间主机名。")
    parser.add_argument("--session-dir", default="", help="覆盖 session_dir。")
    parser.add_argument("--dataset-index-path", default="", help="覆盖 dataset_index_path。")
    parser.add_argument("--chassis-telemetry-topic", default="", help="底盘遥测话题。")
    parser.add_argument("--external-pose-topic", default="", help="外部真值位姿话题。")
    parser.add_argument("--ackermann-command-topic", default="", help="阿克曼内部命令话题。")
    parser.add_argument(
        "--command-source",
        choices=["auto", "topic", "planned", "none"],
        default="",
        help="执行器命令 CSV 来源：auto 表示优先话题，缺失时按动作 profile 写计划命令。",
    )
    parser.add_argument("--action-server", default="/chassis/execute_motion_primitive", help="底盘动作 action 名称。")
    parser.add_argument("--service-timeout-sec", type=float, default=10.0, help="等待 action server 和 goal 响应的超时。")
    parser.add_argument("--action-timeout-sec", type=float, default=0.0, help="覆盖每段动作等待结果超时；0 表示使用 profile timeout。")
    parser.add_argument("--action-result-grace-sec", type=float, default=5.0, help="profile timeout 之外额外等待时间。")
    parser.add_argument("--start-before-motion-sec", type=float, default=-1.0, help="动作前提前采集时间；负数表示读取 profile。")
    parser.add_argument("--stop-after-motion-sec", type=float, default=-1.0, help="动作后延迟采集时间；负数表示读取 profile。")
    parser.add_argument("--stop-on-action-failure", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--allow-incomplete", action="store_true", help="样本不足时仍写入索引和诊断文件。")
    return parser.parse_args()


def as_float(value: Any, default: float = 0.0) -> float:
    if value in (None, ""):
        return default
    return float(value)


def as_bool(value: Any, default: bool = False) -> bool:
    if value in (None, ""):
        return default
    if isinstance(value, bool):
        return value
    return str(value).strip().lower() in {"1", "true", "yes", "on"}


def as_string_list(value: Any) -> list[str]:
    if isinstance(value, list):
        return [str(item) for item in value if str(item)]
    result: list[str] = []
    token = ""
    for ch in str(value):
        if ch in ",; \t\n":
            if token:
                result.append(token)
                token = ""
            continue
        token += ch
    if token:
        result.append(token)
    return result


def spin_for(rclpy_module: Any, node: Any, seconds: float) -> None:
    deadline = time.monotonic() + max(0.0, seconds)
    while rclpy_module.ok() and time.monotonic() < deadline:
        rclpy_module.spin_once(node, timeout_sec=0.05)


def spin_until_future(
    rclpy_module: Any,
    node: Any,
    future: Any,
    timeout_sec: float,
    wait_name: str,
) -> Any:
    deadline = time.monotonic() + timeout_sec
    while rclpy_module.ok() and not future.done() and time.monotonic() < deadline:
        rclpy_module.spin_once(node, timeout_sec=0.05)
    if not future.done():
        raise TimeoutError(f"等待 {wait_name} 超时。")
    return future.result()


def select_actions(profile: dict[str, Any], chassis_type: str, task_code: str) -> list[dict[str, Any]]:
    section = profile["chassis_profiles"][chassis_type]
    actions = list(section["actions"])
    if not task_code:
        return actions
    selected = [action for action in actions if action.get("task_code") == task_code]
    if not selected:
        raise ValueError(f"动作 profile 中没有 task_code={task_code!r}。")
    return selected


def fill_common_goal(goal: Any, config: dict[str, Any], action: dict[str, Any], args: argparse.Namespace) -> None:
    task_code = str(action["task_code"])
    metadata = action["metadata"]
    goal.goal.header.session_id = str(config["session_id"])
    goal.goal.header.task_id = task_code
    goal.goal.header.vehicle_id = str(config["vehicle_id"])
    goal.goal.header.request_id = f"{task_code}_{now_us()}"
    goal.goal.header.client_send_timestamp_us = now_us()
    goal.goal.header.operator_id = args.operator_id
    goal.goal.header.workshop_host = args.workshop_host
    goal.goal.test_case_id = task_code
    goal.goal.task_purpose.value = goal.goal.task_purpose.DATA_COLLECTION
    goal.goal.brake_when_finished = as_bool(metadata.get("brake_when_finished"), True)
    goal.goal.timeout_sec = as_float(metadata.get("timeout_sec"), 60.0)
    goal.goal.source_iteration_id = str(config["session_id"])


def fill_primitive_goal(goal: Any, action: dict[str, Any]) -> None:
    metadata = action["metadata"]
    primitive_type = str(action["primitive_type"])
    goal.goal.selected_primitive.value = getattr(goal.goal.selected_primitive, PRIMITIVE_TYPE_VALUES[primitive_type])

    if primitive_type == "straight_line":
        goal.goal.straight_line.target_distance_m = as_float(metadata["straight_line.target_distance_m"])
        goal.goal.straight_line.target_speed_ms = as_float(metadata["straight_line.target_speed_ms"])
        goal.goal.straight_line.reverse = as_bool(metadata.get("straight_line.reverse"), False)
    elif primitive_type == "arc":
        goal.goal.arc.target_speed_ms = as_float(metadata["arc.target_speed_ms"])
        goal.goal.arc.radius_m = as_float(metadata["arc.radius_m"])
        goal.goal.arc.sweep_angle_deg = as_float(metadata["arc.sweep_angle_deg"])
        goal.goal.arc.clockwise = as_bool(metadata.get("arc.clockwise"), False)
    elif primitive_type == "in_place_rotation":
        goal.goal.in_place_rotation.target_yaw_deg = as_float(metadata["in_place_rotation.target_yaw_deg"])
        goal.goal.in_place_rotation.target_angular_vel_deg_s = as_float(
            metadata["in_place_rotation.target_angular_vel_deg_s"]
        )
    elif primitive_type == "steering_sweep":
        goal.goal.steering_sweep.target_angle_deg = as_float(metadata["steering_sweep.target_angle_deg"])
        goal.goal.steering_sweep.sweep_amplitude_deg = as_float(metadata["steering_sweep.sweep_amplitude_deg"])
        goal.goal.steering_sweep.sweep_frequency_hz = as_float(metadata["steering_sweep.sweep_frequency_hz"])
        goal.goal.steering_sweep.duration_sec = as_float(metadata["steering_sweep.duration_sec"])
    elif primitive_type == "lateral_translation":
        goal.goal.lateral_translation.target_speed_ms = as_float(metadata["lateral_translation.target_speed_ms"])
        goal.goal.lateral_translation.target_distance_m = as_float(metadata["lateral_translation.target_distance_m"])
        goal.goal.lateral_translation.move_left = as_bool(metadata.get("lateral_translation.move_left"), False)
    elif primitive_type == "diagonal_motion":
        goal.goal.diagonal_motion.target_speed_ms = as_float(metadata["diagonal_motion.target_speed_ms"])
        goal.goal.diagonal_motion.target_distance_m = as_float(metadata["diagonal_motion.target_distance_m"])
        goal.goal.diagonal_motion.heading_deg = as_float(metadata["diagonal_motion.heading_deg"])
    elif primitive_type == "module_alignment":
        goal.goal.module_alignment.module_ids = as_string_list(metadata["module_alignment.module_ids"])
        goal.goal.module_alignment.target_zero_deg = as_float(metadata["module_alignment.target_zero_deg"])
        goal.goal.module_alignment.tolerance_deg = as_float(metadata["module_alignment.tolerance_deg"])
    elif primitive_type == "coordinated_steering":
        goal.goal.coordinated_steering.module_ids = as_string_list(metadata["coordinated_steering.module_ids"])
        goal.goal.coordinated_steering.target_angle_deg = as_float(metadata["coordinated_steering.target_angle_deg"])
        goal.goal.coordinated_steering.hold_time_sec = as_float(metadata["coordinated_steering.hold_time_sec"])
    else:
        raise ValueError(f"不支持的动作原语：{primitive_type}")


def build_goal(execute_task_type: Any, config: dict[str, Any], action: dict[str, Any], args: argparse.Namespace) -> Any:
    goal = execute_task_type.Goal()
    fill_common_goal(goal, config, action, args)
    fill_primitive_goal(goal, action)
    return goal


def planned_command_values(action: dict[str, Any]) -> dict[str, float]:
    metadata = action["metadata"]
    primitive_type = str(action["primitive_type"])
    speed = 0.0
    distance = 0.0
    steering_angle_rad = 0.0

    if primitive_type == "straight_line":
        speed = as_float(metadata["straight_line.target_speed_ms"])
        distance = as_float(metadata["straight_line.target_distance_m"])
    elif primitive_type == "arc":
        speed = as_float(metadata["arc.target_speed_ms"])
        distance = as_float(metadata["arc.radius_m"]) * abs(math.radians(as_float(metadata["arc.sweep_angle_deg"])))
    elif primitive_type == "lateral_translation":
        speed = as_float(metadata["lateral_translation.target_speed_ms"])
        distance = as_float(metadata["lateral_translation.target_distance_m"])
    elif primitive_type == "diagonal_motion":
        speed = as_float(metadata["diagonal_motion.target_speed_ms"])
        distance = as_float(metadata["diagonal_motion.target_distance_m"])
    elif primitive_type == "steering_sweep":
        steering_angle_rad = math.radians(as_float(metadata["steering_sweep.target_angle_deg"]))
    elif primitive_type == "module_alignment":
        steering_angle_rad = math.radians(as_float(metadata["module_alignment.target_zero_deg"]))
    elif primitive_type == "coordinated_steering":
        steering_angle_rad = math.radians(as_float(metadata["coordinated_steering.target_angle_deg"]))

    return {
        "target_speed_ms": speed,
        "target_distance_m": distance,
        "target_steering_angle_rad": steering_angle_rad,
    }


def write_planned_command_row(capture: ChassisSessionCapture, action: dict[str, Any], command_timestamp_us: int) -> None:
    values = planned_command_values(action)
    capture.command_csv.write_row({
        "command_timestamp_us": command_timestamp_us,
        "command_id": str(action["task_code"]),
        "source": "chassis_action_profile",
        "target_speed_ms": values["target_speed_ms"],
        "target_distance_m": values["target_distance_m"],
        "target_accel_ms2": 0.0,
        "target_steering_angle_rad": values["target_steering_angle_rad"],
        "target_steering_rate_rads": 0.0,
        "brake_command": 0.0,
        "throttle_command": 0.0,
        "command_timeout_sec": as_float(action["metadata"].get("timeout_sec"), 0.0),
    })


def update_capture_args_for_action(capture: ChassisSessionCapture, action: dict[str, Any]) -> None:
    values = planned_command_values(action)
    capture.args.command_id = str(action["task_code"])
    capture.args.target_speed_ms = values["target_speed_ms"]
    capture.args.target_distance_m = values["target_distance_m"]
    capture.args.target_steering_angle_rad = values["target_steering_angle_rad"]
    capture.args.brake_command = 0.0


def execute_action(
    rclpy_module: Any,
    node: Any,
    action_client: Any,
    execute_task_type: Any,
    config: dict[str, Any],
    action: dict[str, Any],
    args: argparse.Namespace,
) -> Any:
    if not action_client.wait_for_server(timeout_sec=args.service_timeout_sec):
        raise TimeoutError(f"底盘动作 action server 不可用: {args.action_server}")

    goal = build_goal(execute_task_type, config, action, args)
    send_future = action_client.send_goal_async(goal)
    goal_handle = spin_until_future(
        rclpy_module,
        node,
        send_future,
        args.service_timeout_sec,
        "发送底盘动作 goal",
    )
    if not goal_handle.accepted:
        raise RuntimeError(f"底盘动作被拒绝: {action['task_code']}")

    result_timeout = args.action_timeout_sec
    if result_timeout <= 0.0:
        result_timeout = goal.goal.timeout_sec + max(0.0, args.action_result_grace_sec)
    result_future = goal_handle.get_result_async()
    return spin_until_future(
        rclpy_module,
        node,
        result_future,
        result_timeout,
        f"底盘动作结果 {action['task_code']}",
    )


def import_ros_dependencies() -> dict[str, Any]:
    try:
        import rclpy
        from rclpy.action import ActionClient
        from calibration_chassis_interfaces.action import ExecuteMotionPrimitive
        from calibration_chassis_interfaces.msg import ChassisTelemetry
        from calibration_external_localization_interfaces.msg import ExternalLocalizationTelemetry
    except ImportError as exc:
        raise RuntimeError(f"缺少 ROS 2 Python 依赖或标定消息包: {exc}") from exc

    try:
        from vehicle_internal_interfaces.msg import AckermannDriveCommand
    except ImportError:
        AckermannDriveCommand = None

    return {
        "rclpy": rclpy,
        "ActionClient": ActionClient,
        "ExecuteMotionPrimitive": ExecuteMotionPrimitive,
        "ChassisTelemetry": ChassisTelemetry,
        "ExternalLocalizationTelemetry": ExternalLocalizationTelemetry,
        "AckermannDriveCommand": AckermannDriveCommand,
    }


def main() -> int:
    args = parse_args()
    config_path = Path(args.config).expanduser().resolve(strict=False)
    config = load_config(config_path)
    apply_cli_overrides(config, args)
    if args.chassis_type:
        config["chassis_type"] = args.chassis_type
    if not args.command_source:
        args.command_source = config_command_source(config)

    validation = validate_config(config, config_path)
    if not validation.ok:
        for error in validation.errors:
            print(f"[错误] {error}", file=sys.stderr)
        return 1

    profile_path = Path(args.action_profile).expanduser().resolve(strict=False)
    action_profile = validate_action_profile(load_action_profile_yaml(profile_path))
    chassis_type = str(config["chassis_type"])
    actions = select_actions(action_profile, chassis_type, args.task_code)
    data_capture = action_profile.get("data_capture", {}) or {}
    start_before_sec = (
        as_float(data_capture.get("start_before_motion_sec"), 1.0)
        if args.start_before_motion_sec < 0.0 else args.start_before_motion_sec
    )
    stop_after_sec = (
        as_float(data_capture.get("stop_after_motion_sec"), 1.0)
        if args.stop_after_motion_sec < 0.0 else args.stop_after_motion_sec
    )

    try:
        deps = import_ros_dependencies()
    except RuntimeError as exc:
        print(f"[错误] {exc}", file=sys.stderr)
        return 1

    rclpy_module = deps["rclpy"]
    target_paths = make_target_paths(config)
    chassis_topic = args.chassis_telemetry_topic or config_topic(config, "chassis_telemetry_topic", "/chassis/telemetry")
    external_topic = args.external_pose_topic or config_topic(
        config,
        "external_pose_topic",
        "/workshop/external_localization/vehicle/pose",
    )
    command_topic = args.ackermann_command_topic or config_topic(config, "ackermann_command_topic", "")

    rclpy_module.init()
    node = rclpy_module.create_node("workshop_chassis_profile_capture")
    capture = ChassisSessionCapture(node, config, target_paths, args)
    action_failed = False
    finalize_code = 1

    try:
        node.create_subscription(deps["ChassisTelemetry"], chassis_topic, capture.on_chassis_telemetry, 50)
        node.create_subscription(deps["ExternalLocalizationTelemetry"], external_topic, capture.on_external_pose, 50)
        if deps["AckermannDriveCommand"] is not None and command_topic and args.command_source in ("auto", "topic"):
            node.create_subscription(deps["AckermannDriveCommand"], command_topic, capture.on_ackermann_command, 50)
        elif args.command_source == "topic":
            raise RuntimeError("已要求从命令话题采集，但命令消息包或话题配置不可用。")

        action_client = deps["ActionClient"](node, deps["ExecuteMotionPrimitive"], args.action_server)
        print(f"[*] 底盘动作 profile: {profile_path}", file=sys.stderr)
        print(f"[*] 底盘类型: {chassis_type}", file=sys.stderr)
        print(f"[*] 动作数量: {len(actions)}", file=sys.stderr)
        print(f"[*] 底盘遥测采集: {chassis_topic}", file=sys.stderr)
        print(f"[*] 外部真值采集: {external_topic}", file=sys.stderr)
        if command_topic:
            print(f"[*] 执行器命令采集: {command_topic}", file=sys.stderr)

        for index, action in enumerate(actions, start=1):
            update_capture_args_for_action(capture, action)
            before_command_count = capture.command_csv.count
            command_timestamp_us = now_us()
            print(f"[*] 执行动作 {index}/{len(actions)}: {action['task_code']}", file=sys.stderr)
            spin_for(rclpy_module, node, start_before_sec)
            if args.command_source == "planned":
                write_planned_command_row(capture, action, command_timestamp_us)
            wrapped_result = execute_action(
                rclpy_module,
                node,
                action_client,
                deps["ExecuteMotionPrimitive"],
                config,
                action,
                args,
            )
            spin_for(rclpy_module, node, stop_after_sec)
            if args.command_source == "auto" and capture.command_csv.count == before_command_count:
                write_planned_command_row(capture, action, command_timestamp_us)

            job_result = wrapped_result.result.result
            if not job_result.success:
                action_failed = True
                print(f"[错误] 动作失败 {action['task_code']}: {job_result.message}", file=sys.stderr)
                if args.stop_on_action_failure:
                    break
            else:
                message = job_result.message or "动作完成。"
                print(f"[*] 动作完成 {action['task_code']}: {message}", file=sys.stderr)
    except Exception as exc:
        action_failed = True
        print(f"[错误] {exc}", file=sys.stderr)
    finally:
        capture.close()
        node.destroy_node()
        rclpy_module.shutdown()
        finalize_code = finalize_capture(config, target_paths, config_path, args)

    if action_failed:
        return 2 if finalize_code == 0 else finalize_code
    return finalize_code


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