#!/usr/bin/env python3
"""按运控评估 profile 执行任务，并同步采集现场数据。"""

from __future__ import annotations

import argparse
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 capture_control_session import (
    ControlSessionCapture,
    config_topic,
    finalize_capture,
    make_target_paths,
    now_us,
)
from control_data_common import apply_cli_overrides, load_config, validate_config
from control_evaluation_profile import (
    TASK_PAYLOAD_KEYS,
    TASK_TYPE_TO_METADATA,
    load_yaml as load_evaluation_profile_yaml,
    validate_profile as validate_evaluation_profile,
)


TASK_TYPE_VALUES = {
    "trajectory_tracking": "TRAJECTORY_TRACKING",
    "velocity_step": "VELOCITY_STEP",
    "acceleration_deceleration": "ACCELERATION_DECELERATION",
    "accel_decel": "ACCELERATION_DECELERATION",
    "stop_accuracy": "STOP_ACCURACY",
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="按现场运控评估 profile 顺序发送控制评估任务，同时采集运控、底盘和外部真值数据。"
    )
    parser.add_argument("--config", required=True, help="运控标定现场数据配置文件。")
    parser.add_argument(
        "--evaluation-profile",
        default="src/site_deployment/workshop_control_calibration_real/config/control_evaluation_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_control_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("--control-telemetry-topic", default="", help="运控遥测话题。")
    parser.add_argument("--chassis-telemetry-topic", default="", help="底盘遥测话题。")
    parser.add_argument("--external-pose-topic", default="", help="外部真值位姿话题。")
    parser.add_argument("--task-type", default="", help=argparse.SUPPRESS)
    parser.add_argument("--action-server", default="/control/execute_controller_evaluation", 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 推导。")
    parser.add_argument("--action-result-grace-sec", type=float, default=8.0, help="profile 推导时长之外额外等待时间。")
    parser.add_argument("--start-before-task-sec", type=float, default=-1.0, help="任务前提前采集时间；负数表示读取 profile。")
    parser.add_argument("--stop-after-task-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_uint(value: Any, default: int = 0) -> int:
    if value in (None, ""):
        return default
    return max(0, int(value))


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 normalized_task_type(action: dict[str, Any]) -> str:
    return TASK_TYPE_TO_METADATA[str(action["selected_task"])]


def action_payload(action: dict[str, Any]) -> dict[str, Any]:
    payload_key = TASK_PAYLOAD_KEYS[str(action["selected_task"])]
    return action[payload_key]


def task_timeout_sec(action: dict[str, Any], args: argparse.Namespace) -> float:
    if args.action_timeout_sec > 0.0:
        return args.action_timeout_sec

    selected_task = str(action["selected_task"])
    payload = action_payload(action)
    if selected_task == "trajectory_tracking":
        base = as_float(payload.get("timeout_sec"), 60.0)
    elif selected_task == "stop_accuracy":
        base = as_float(payload.get("timeout_sec"), 45.0)
    elif selected_task == "velocity_step":
        base = as_float(payload.get("settle_before_step_sec"), 0.0) + as_float(payload.get("hold_time_sec"), 0.0)
    else:
        start_speed = as_float(payload.get("start_velocity_ms"), 0.0)
        target_speed = as_float(payload.get("target_velocity_ms"), 0.0)
        target_accel = abs(as_float(payload.get("target_accel_ms2"), 0.1))
        ramp_sec = abs(target_speed - start_speed) / target_accel if target_accel > 0.0 else 0.0
        base = ramp_sec + as_float(payload.get("hold_time_sec"), 0.0)
    return base + max(0.0, args.action_result_grace_sec)


def fill_common_goal(goal: Any, config: dict[str, Any], action: dict[str, Any], args: argparse.Namespace) -> None:
    task_code = str(action["task_code"])
    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.selected_task.value = getattr(goal.goal.selected_task, TASK_TYPE_VALUES[str(action["selected_task"])])
    goal.goal.source_iteration_id = str(config["session_id"])


def fill_trajectory_goal(goal: Any, action: dict[str, Any], trajectory_point_type: Any) -> None:
    payload = action_payload(action)
    goal.goal.trajectory_tracking.stop_at_end = as_bool(payload.get("stop_at_end"), True)
    goal.goal.trajectory_tracking.timeout_sec = as_float(payload.get("timeout_sec"), 60.0)
    goal.goal.trajectory_tracking.required_external_pose_source_id = str(
        payload.get("required_external_pose_source_id", "") or ""
    )
    goal.goal.trajectory_tracking.max_external_pose_age_ms = as_float(payload.get("max_external_pose_age_ms"), 0.0)
    goal.goal.trajectory_tracking.min_external_pose_quality_score = as_float(
        payload.get("min_external_pose_quality_score"),
        0.0,
    )
    goal.goal.trajectory_tracking.trajectory_id = str(payload.get("trajectory_id", "") or "")
    goal.goal.trajectory_tracking.segment_index = as_uint(payload.get("segment_index"), 0)
    goal.goal.trajectory_tracking.total_segments = as_uint(payload.get("total_segments"), 0)
    goal.goal.trajectory_tracking.is_final_segment = as_bool(payload.get("is_final_segment"), False)

    for point in payload["path"]:
        point_msg = trajectory_point_type()
        point_msg.x_m = as_float(point["x_m"])
        point_msg.y_m = as_float(point["y_m"])
        point_msg.yaw_rad = as_float(point["yaw_rad"])
        point_msg.target_speed_ms = as_float(point["target_speed_ms"])
        goal.goal.trajectory_tracking.path.append(point_msg)


def fill_velocity_step_goal(goal: Any, action: dict[str, Any]) -> None:
    payload = action_payload(action)
    goal.goal.velocity_step.target_velocity_ms = as_float(payload["target_velocity_ms"])
    goal.goal.velocity_step.hold_time_sec = as_float(payload["hold_time_sec"])
    goal.goal.velocity_step.settle_before_step_sec = as_float(payload["settle_before_step_sec"])


def fill_accel_decel_goal(goal: Any, action: dict[str, Any]) -> None:
    payload = action_payload(action)
    goal.goal.accel_decel.start_velocity_ms = as_float(payload["start_velocity_ms"])
    goal.goal.accel_decel.target_velocity_ms = as_float(payload["target_velocity_ms"])
    goal.goal.accel_decel.target_accel_ms2 = as_float(payload["target_accel_ms2"])
    goal.goal.accel_decel.hold_time_sec = as_float(payload["hold_time_sec"])


def fill_stop_accuracy_goal(goal: Any, action: dict[str, Any]) -> None:
    payload = action_payload(action)
    goal.goal.stop_accuracy.target_stop_x_m = as_float(payload["target_stop_x_m"])
    goal.goal.stop_accuracy.target_stop_y_m = as_float(payload["target_stop_y_m"])
    goal.goal.stop_accuracy.target_stop_yaw_rad = as_float(payload["target_stop_yaw_rad"])
    goal.goal.stop_accuracy.timeout_sec = as_float(payload["timeout_sec"])


def build_goal(
    execute_task_type: Any,
    trajectory_point_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)
    selected_task = str(action["selected_task"])
    if selected_task == "trajectory_tracking":
        fill_trajectory_goal(goal, action, trajectory_point_type)
    elif selected_task == "velocity_step":
        fill_velocity_step_goal(goal, action)
    elif selected_task in {"acceleration_deceleration", "accel_decel"}:
        fill_accel_decel_goal(goal, action)
    elif selected_task == "stop_accuracy":
        fill_stop_accuracy_goal(goal, action)
    else:
        raise ValueError(f"不支持的运控评估任务: {selected_task}")
    return goal


def update_capture_for_action(capture: ControlSessionCapture, action: dict[str, Any], chassis_type: str) -> None:
    capture.task_type = normalized_task_type(action)
    capture.chassis_type = chassis_type
    capture.control_axis = str(action["control_axis"])
    capture.controller_algorithm = str(action["controller_algorithm"])
    capture.control_role = str(action["control_role"])


def execute_action(
    rclpy_module: Any,
    node: Any,
    action_client: Any,
    execute_task_type: Any,
    trajectory_point_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, trajectory_point_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_future = goal_handle.get_result_async()
    return spin_until_future(
        rclpy_module,
        node,
        result_future,
        task_timeout_sec(action, args),
        f"运控评估结果 {action['task_code']}",
    )


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

    return {
        "rclpy": rclpy,
        "ActionClient": ActionClient,
        "ExecuteControllerEvaluation": ExecuteControllerEvaluation,
        "ControlTelemetry": ControlTelemetry,
        "TrajectoryPoint": TrajectoryPoint,
        "ChassisTelemetry": ChassisTelemetry,
        "ExternalLocalizationTelemetry": ExternalLocalizationTelemetry,
    }


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

    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.evaluation_profile).expanduser().resolve(strict=False)
    evaluation_profile = validate_evaluation_profile(load_evaluation_profile_yaml(profile_path))
    chassis_type = str(config["chassis_type"])
    actions = select_actions(evaluation_profile, chassis_type, args.task_code)
    data_capture = evaluation_profile.get("data_capture", {}) or {}
    start_before_sec = (
        as_float(data_capture.get("start_before_task_sec"), 1.0)
        if args.start_before_task_sec < 0.0 else args.start_before_task_sec
    )
    stop_after_sec = (
        as_float(data_capture.get("stop_after_task_sec"), 1.0)
        if args.stop_after_task_sec < 0.0 else args.stop_after_task_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)
    control_topic = args.control_telemetry_topic or config_topic(config, "control_telemetry_topic", "/control/telemetry")
    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",
    )

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

    try:
        node.create_subscription(deps["ControlTelemetry"], control_topic, capture.on_control_telemetry, 50)
        node.create_subscription(deps["ChassisTelemetry"], chassis_topic, capture.on_chassis_telemetry, 50)
        node.create_subscription(deps["ExternalLocalizationTelemetry"], external_topic, capture.on_external_pose, 50)

        action_client = deps["ActionClient"](node, deps["ExecuteControllerEvaluation"], 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"[*] 运控遥测采集: {control_topic}", file=sys.stderr)
        print(f"[*] 底盘响应采集: {chassis_topic}", file=sys.stderr)
        print(f"[*] 外部真值采集: {external_topic}", file=sys.stderr)

        for index, action in enumerate(actions, start=1):
            update_capture_for_action(capture, action, chassis_type)
            print(f"[*] 执行运控任务 {index}/{len(actions)}: {action['task_code']}", file=sys.stderr)
            spin_for(rclpy_module, node, start_before_sec)
            wrapped_result = execute_action(
                rclpy_module,
                node,
                action_client,
                deps["ExecuteControllerEvaluation"],
                deps["TrajectoryPoint"],
                config,
                action,
                args,
            )
            spin_for(rclpy_module, node, stop_after_sec)

            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())
