#!/usr/bin/env python3
"""本地验证运控评估 profile 执行和采集落盘链路。"""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
import tempfile
import threading
import time
from pathlib import Path
from typing import Any

try:
    import yaml
except ImportError as exc:
    raise SystemExit("缺少 PyYAML，请先在当前 Python 环境安装 yaml 模块。") from exc

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

from control_data_common import read_csv_rows


def now_us() -> int:
    return int(time.time() * 1_000_000)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="启动假的运控 action server 和遥测话题，验证 run_control_profile_capture.py。"
    )
    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="ackermann",
        choices=["ackermann", "differential", "single_steer_wheel", "multi_steer_wheel"],
        help="本次验证使用的底盘类型。",
    )
    parser.add_argument(
        "--task-code",
        default="",
        help="只验证指定任务；为空时默认取该底盘类型的第一个任务。",
    )
    parser.add_argument("--session-dir", default="", help="输出目录；为空时使用临时目录。")
    parser.add_argument("--keep-output", action="store_true", help="保留临时输出目录。")
    parser.add_argument("--verbose", action="store_true", help="打印 profile runner 的完整输出。")
    parser.add_argument("--timeout-sec", type=float, default=35.0, help="等待 profile runner 结束的超时。")
    parser.add_argument("--rmw-implementation", default="rmw_fastrtps_cpp", help="本地 smoke 使用的 RMW 实现。")
    return parser.parse_args()


def load_yaml(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as stream:
        data = yaml.safe_load(stream) or {}
    if not isinstance(data, dict):
        raise ValueError(f"{path} 顶层必须是 YAML map。")
    return data


def first_task_code(profile_path: Path, chassis_type: str) -> str:
    profile = load_yaml(profile_path)
    actions = profile["chassis_profiles"][chassis_type]["actions"]
    if not actions:
        raise ValueError(f"{chassis_type} 没有配置运控评估任务。")
    return str(actions[0]["task_code"])


def write_smoke_config(path: Path, session_dir: Path, chassis_type: str) -> None:
    config = {
        "csv_contract_version": 1,
        "session_id": "control_profile_capture_smoke",
        "site_id": "local_smoke_site",
        "vehicle_id": "smoke_agv_001",
        "session_dir": str(session_dir),
        "dataset_index_path": str(session_dir / "dataset_index.yaml"),
        "chassis_type": chassis_type,
        "control_axis": "combined",
        "controller_algorithm": "pure_pursuit",
        "control_role": "path_tracking_outer_loop",
        "files": {
            "control_evaluation_data_file": "control/control_eval.csv",
            "reference_signal_file": "control/reference_signal.csv",
            "chassis_response_file": "control/chassis_response.csv",
            "truth_trajectory_file": "external/truth_trajectory.csv",
            "diagnostics_file": "control/control_diagnostics.json",
        },
        "capture": {
            "control_telemetry_topic": "/smoke/control/telemetry",
            "chassis_telemetry_topic": "/smoke/chassis/telemetry",
            "external_pose_topic": "/smoke/external_localization/vehicle/pose",
        },
        "validation": {
            "min_control_evaluation_samples": 2,
            "min_reference_signal_samples": 2,
            "min_chassis_response_samples": 2,
            "min_truth_trajectory_samples": 2,
            "max_time_gap_ms": 500.0,
        },
        "synthetic": {
            "enabled": False,
        },
    }
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(yaml.safe_dump(config, sort_keys=False, allow_unicode=True), encoding="utf-8")


class FakeControlFixture:
    def __init__(self, node: Any, action_name: str, chassis_type_name: str) -> None:
        from rclpy.action import ActionServer
        from calibration_chassis_interfaces.msg import ChassisTelemetry
        from calibration_common_interfaces.msg import ErrorCode
        from calibration_control_interfaces.action import ExecuteControllerEvaluation
        from calibration_control_interfaces.msg import ControlTelemetry
        from calibration_control_interfaces.msg import ControllerEvaluationTaskType
        from calibration_external_localization_interfaces.msg import ExternalLocalizationTelemetry
        from calibration_vehicle_profile_interfaces.msg import ChassisType
        from calibration_vehicle_profile_interfaces.msg import ControllerAlgorithmType

        self.node = node
        self.ExecuteControllerEvaluation = ExecuteControllerEvaluation
        self.ControlTelemetry = ControlTelemetry
        self.ControllerEvaluationTaskType = ControllerEvaluationTaskType
        self.ChassisTelemetry = ChassisTelemetry
        self.ExternalLocalizationTelemetry = ExternalLocalizationTelemetry
        self.ErrorCode = ErrorCode
        self.ChassisType = ChassisType
        self.ControllerAlgorithmType = ControllerAlgorithmType
        self.chassis_type_value = {
            "ackermann": ChassisType.ACKERMANN,
            "differential": ChassisType.DIFFERENTIAL,
            "single_steer_wheel": ChassisType.SINGLE_STEER_WHEEL,
            "multi_steer_wheel": ChassisType.MULTI_STEER_WHEEL,
        }[chassis_type_name]
        self.lock = threading.Lock()
        self.x_m = 0.0
        self.y_m = 0.0
        self.yaw_rad = 0.0
        self.speed_ms = 0.0
        self.target_speed_ms = 0.1
        self.active_job_id = ""
        self.control_pub = node.create_publisher(ControlTelemetry, "/smoke/control/telemetry", 20)
        self.chassis_pub = node.create_publisher(ChassisTelemetry, "/smoke/chassis/telemetry", 20)
        self.truth_pub = node.create_publisher(
            ExternalLocalizationTelemetry,
            "/smoke/external_localization/vehicle/pose",
            20,
        )
        self.timer = node.create_timer(0.05, self.publish_telemetry)
        self.action_server = ActionServer(
            node,
            ExecuteControllerEvaluation,
            action_name,
            execute_callback=self.execute_control,
        )

    def publish_telemetry(self) -> None:
        stamp = now_us()
        with self.lock:
            x_m = self.x_m
            y_m = self.y_m
            yaw_rad = self.yaw_rad
            speed_ms = self.speed_ms
            target_speed_ms = self.target_speed_ms
            active_job_id = self.active_job_id

        control_msg = self.ControlTelemetry()
        control_msg.hardware_timestamp_us = stamp
        control_msg.odom_x_m = x_m
        control_msg.odom_y_m = y_m
        control_msg.odom_yaw_rad = yaw_rad
        control_msg.linear_velocity_ms = speed_ms
        control_msg.angular_velocity_rads = 0.02
        control_msg.lateral_error_m = 0.005
        control_msg.heading_error_rad = 0.003
        control_msg.speed_error_ms = target_speed_ms - speed_ms
        control_msg.steering_output = 0.02
        control_msg.throttle_output = 0.1
        control_msg.brake_output = 0.0
        control_msg.saturation_flag = False
        control_msg.active_job_id = active_job_id
        control_msg.parameter_version = "smoke_fake_control"
        control_msg.active_lateral_algorithm.value = self.ControllerAlgorithmType.PURE_PURSUIT
        control_msg.active_longitudinal_algorithm.value = self.ControllerAlgorithmType.PID
        control_msg.pose_source_name = "smoke_fake_truth"
        control_msg.pose_from_external_truth = True
        self.control_pub.publish(control_msg)

        chassis_msg = self.ChassisTelemetry()
        chassis_msg.hardware_timestamp_us = stamp
        chassis_msg.chassis_type.value = self.chassis_type_value
        chassis_msg.odom_x_m = x_m
        chassis_msg.odom_y_m = y_m
        chassis_msg.odom_yaw_rad = yaw_rad
        chassis_msg.linear_velocity_ms = speed_ms
        chassis_msg.angular_velocity_rads = 0.02
        chassis_msg.estop_engaged = False
        chassis_msg.driver_error_code = 0
        chassis_msg.active_job_id = active_job_id
        chassis_msg.lateral_slip_estimate = 0.0
        chassis_msg.curvature_estimate = 0.0
        self.chassis_pub.publish(chassis_msg)

        truth_msg = self.ExternalLocalizationTelemetry()
        truth_msg.hardware_timestamp_us = stamp
        truth_msg.pose_valid = True
        truth_msg.workshop_pose.x_m = x_m
        truth_msg.workshop_pose.y_m = y_m
        truth_msg.workshop_pose.z_m = 0.0
        truth_msg.workshop_pose.roll_rad = 0.0
        truth_msg.workshop_pose.pitch_rad = 0.0
        truth_msg.workshop_pose.yaw_rad = yaw_rad
        truth_msg.position_stddev_m = 0.001
        truth_msg.yaw_stddev_rad = 0.001
        truth_msg.tracking_loss_ratio = 0.0
        truth_msg.time_sync_offset_ms = 1.0
        truth_msg.quality_score = 1.0
        truth_msg.observed_target_count = 4
        truth_msg.reference_source_name = "smoke_fake_truth"
        truth_msg.active_job_id = active_job_id
        self.truth_pub.publish(truth_msg)

    def execute_control(self, goal_handle: Any) -> Any:
        request = goal_handle.request.goal
        with self.lock:
            self.active_job_id = request.header.request_id

        steps = 10
        for step_index in range(steps):
            self.advance_state(request, step_index, steps)
            self.publish_telemetry()
            time.sleep(0.05)

        result = self.ExecuteControllerEvaluation.Result()
        result.result.success = True
        result.result.error_code.code = self.ErrorCode.OK
        result.result.message = "本地 smoke 假运控任务完成。"
        result.result.job_id = request.header.request_id
        result.result.data_quality_passed = True
        result.result.suitable_for_commit = False
        result.result.recommended_parameter_version = "smoke_fake_control"
        goal_handle.succeed()
        with self.lock:
            self.active_job_id = ""
        return result

    def advance_state(self, request: Any, step_index: int, steps: int) -> None:
        selected_task = request.selected_task.value
        with self.lock:
            if selected_task == self.ControllerEvaluationTaskType.TRAJECTORY_TRACKING:
                path = list(request.trajectory_tracking.path)
                if path:
                    point = path[min(len(path) - 1, int(step_index * len(path) / steps))]
                    self.x_m += (point.x_m - self.x_m) * 0.5
                    self.y_m += (point.y_m - self.y_m) * 0.5
                    self.yaw_rad += (point.yaw_rad - self.yaw_rad) * 0.5
                    self.target_speed_ms = point.target_speed_ms
                    self.speed_ms += (point.target_speed_ms - self.speed_ms) * 0.5
            elif selected_task == self.ControllerEvaluationTaskType.VELOCITY_STEP:
                self.target_speed_ms = request.velocity_step.target_velocity_ms
                self.speed_ms += (self.target_speed_ms - self.speed_ms) * 0.4
                self.x_m += self.speed_ms * 0.05
            elif selected_task == self.ControllerEvaluationTaskType.ACCELERATION_DECELERATION:
                self.target_speed_ms = request.accel_decel.target_velocity_ms
                self.speed_ms += request.accel_decel.target_accel_ms2 * 0.05
                if self.speed_ms > self.target_speed_ms:
                    self.speed_ms = self.target_speed_ms
                self.x_m += self.speed_ms * 0.05
            elif selected_task == self.ControllerEvaluationTaskType.STOP_ACCURACY:
                self.x_m += (request.stop_accuracy.target_stop_x_m - self.x_m) * 0.4
                self.y_m += (request.stop_accuracy.target_stop_y_m - self.y_m) * 0.4
                self.yaw_rad += (request.stop_accuracy.target_stop_yaw_rad - self.yaw_rad) * 0.4
                self.target_speed_ms = 0.0
                self.speed_ms *= 0.5


def run_profile_capture(
    repo_root: Path,
    config_path: Path,
    evaluation_profile_path: Path,
    session_dir: Path,
    args: argparse.Namespace,
    action_name: str,
    task_code: str,
) -> subprocess.CompletedProcess[str]:
    command = [
        sys.executable,
        str(SCRIPT_DIR / "run_control_profile_capture.py"),
        "--config",
        str(config_path),
        "--evaluation-profile",
        str(evaluation_profile_path),
        "--chassis-type",
        args.chassis_type,
        "--task-code",
        task_code,
        "--session-id",
        "control_profile_capture_smoke",
        "--site-id",
        "local_smoke_site",
        "--vehicle-id",
        "smoke_agv_001",
        "--session-dir",
        str(session_dir),
        "--dataset-index-path",
        str(session_dir / "dataset_index.yaml"),
        "--control-telemetry-topic",
        "/smoke/control/telemetry",
        "--chassis-telemetry-topic",
        "/smoke/chassis/telemetry",
        "--external-pose-topic",
        "/smoke/external_localization/vehicle/pose",
        "--action-server",
        action_name,
        "--start-before-task-sec",
        "0.2",
        "--stop-after-task-sec",
        "0.2",
        "--service-timeout-sec",
        "5.0",
        "--action-result-grace-sec",
        "5.0",
    ]
    return subprocess.run(
        command,
        cwd=repo_root,
        env=os.environ.copy(),
        text=True,
        capture_output=True,
        timeout=args.timeout_sec,
        check=False,
    )


def validate_outputs(session_dir: Path) -> None:
    dataset_index_path = session_dir / "dataset_index.yaml"
    if not dataset_index_path.exists():
        raise RuntimeError(f"未生成 dataset_index.yaml: {dataset_index_path}")
    index = load_yaml(dataset_index_path)
    control_inputs = index.get("data_inputs", {}).get("control", {})
    for field_name in (
        "control_evaluation_data_files",
        "reference_signal_files",
        "chassis_response_files",
        "truth_trajectory_files",
    ):
        values = control_inputs.get(field_name)
        if not values:
            raise RuntimeError(f"dataset_index 缺少 data_inputs.control.{field_name}")

    control_rows = read_csv_rows(session_dir / "control/control_eval.csv")
    reference_rows = read_csv_rows(session_dir / "control/reference_signal.csv")
    chassis_rows = read_csv_rows(session_dir / "control/chassis_response.csv")
    truth_rows = read_csv_rows(session_dir / "external/truth_trajectory.csv")
    if len(control_rows) < 2 or len(reference_rows) < 2 or len(chassis_rows) < 2 or len(truth_rows) < 2:
        raise RuntimeError(
            "采集样本数不足: "
            f"control={len(control_rows)}, reference={len(reference_rows)}, "
            f"chassis={len(chassis_rows)}, truth={len(truth_rows)}"
        )


def main() -> int:
    args = parse_args()
    repo_root = Path.cwd()
    evaluation_profile_path = Path(args.evaluation_profile).expanduser().resolve(strict=False)
    task_code = args.task_code or first_task_code(evaluation_profile_path, args.chassis_type)
    action_name = f"/smoke/control/execute_controller_evaluation_{now_us()}"

    temp_dir: Any = None
    if args.session_dir:
        session_dir = Path(args.session_dir).expanduser().resolve(strict=False)
        session_dir.mkdir(parents=True, exist_ok=True)
    else:
        temp_dir = tempfile.TemporaryDirectory(prefix="agv_control_profile_capture_smoke_")
        session_dir = Path(temp_dir.name)

    if args.rmw_implementation:
        os.environ["RMW_IMPLEMENTATION"] = args.rmw_implementation

    try:
        import rclpy
        from rclpy.executors import MultiThreadedExecutor
    except ImportError as exc:
        print(f"[FAIL] 缺少 ROS2 Python 环境，请先 source install/setup.bash: {exc}", file=sys.stderr)
        if temp_dir is not None:
            temp_dir.cleanup()
        return 1

    config_path = session_dir / "control_data.yaml"
    write_smoke_config(config_path, session_dir, args.chassis_type)
    ros_log_dir = session_dir / "ros_log"
    ros_log_dir.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("ROS_LOG_DIR", str(ros_log_dir))

    rclpy.init(args=None)
    node = rclpy.create_node("smoke_control_profile_capture_fixture")
    fixture = FakeControlFixture(node, action_name, args.chassis_type)
    executor = MultiThreadedExecutor(num_threads=2)
    executor.add_node(node)
    thread = threading.Thread(target=executor.spin, daemon=True)
    thread.start()

    try:
        process = run_profile_capture(
            repo_root,
            config_path,
            evaluation_profile_path,
            session_dir,
            args,
            action_name,
            task_code,
        )
        if args.verbose or process.returncode != 0:
            print(process.stdout, end="")
            print(process.stderr, end="", file=sys.stderr)
        if process.returncode != 0:
            print(f"[FAIL] run_control_profile_capture.py 返回 {process.returncode}", file=sys.stderr)
            return process.returncode
        validate_outputs(session_dir)
        print(f"[PASS] 运控 profile 采集 smoke 通过: {session_dir}")
        print(f"[PASS] task_code={task_code}")
        return 0
    finally:
        executor.shutdown()
        fixture.action_server.destroy()
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()
        thread.join(timeout=2.0)
        if temp_dir is not None and not args.keep_output:
            temp_dir.cleanup()


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