#!/usr/bin/env python3
"""现场底盘标定数据采集落盘入口。"""

from __future__ import annotations

import argparse
import csv
import json
import sys
import time
from pathlib import Path
from typing import Any

from chassis_data_common import (
    ACTUATOR_COMMAND_FIELDS,
    CHASSIS_MOTION_FIELDS,
    TRUTH_TRAJECTORY_FIELDS,
    apply_cli_overrides,
    build_dataset_index,
    build_summary,
    has_placeholder,
    load_config,
    read_csv_rows,
    resolve_path,
    normalize_chassis_type_name,
    validate_chassis_dataset_files,
    validate_config,
    validate_summary,
    write_yaml,
)


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


def int_value(value: Any, default: int = 0) -> int:
    if value in (None, ""):
        return default
    return int(value)


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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="从现场 ROS 话题采集底盘遥测、执行器命令和外部真值轨迹，并写入 dataset_index.yaml。"
    )
    parser.add_argument("--config", required=True, help="底盘标定现场数据配置文件。")
    parser.add_argument("--chassis-type", default="", help="覆盖 chassis_type。")
    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("--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 表示优先话题，缺失时写入计划命令。",
    )
    parser.add_argument("--duration-sec", type=float, default=0.0, help="采集时长；0 表示按 Ctrl+C 结束。")
    parser.add_argument("--max-chassis-samples", type=int, default=0, help="达到该底盘样本数后自动结束；0 表示不限。")
    parser.add_argument("--target-speed-ms", type=float, default=0.0, help="计划命令目标速度。")
    parser.add_argument("--target-distance-m", type=float, default=0.0, help="计划命令目标距离。")
    parser.add_argument("--target-steering-angle-rad", type=float, default=0.0, help="计划命令目标转角。")
    parser.add_argument("--brake-command", type=float, default=0.0, help="计划命令制动值。")
    parser.add_argument("--command-id", default="", help="计划命令 ID。")
    parser.add_argument("--allow-incomplete", action="store_true", help="样本不足时仍写入索引和诊断文件。")
    return parser.parse_args()


def config_topic(config: dict[str, Any], key: str, default: str) -> str:
    capture = config.get("capture", {}) or {}
    if not isinstance(capture, dict):
        return default
    value = str(capture.get(key, "") or "").strip()
    if not value or has_placeholder(value):
        return default
    return value


def synthetic_value(config: dict[str, Any], key: str, default: float) -> float:
    synthetic = config.get("synthetic", {}) or {}
    if not isinstance(synthetic, dict):
        return default
    return float_value(synthetic.get(key), default)


def config_command_source(config: dict[str, Any]) -> str:
    capture = config.get("capture", {}) or {}
    if not isinstance(capture, dict):
        return "auto"
    value = str(capture.get("command_source", "auto") or "auto").strip()
    return value if value in {"auto", "topic", "planned", "none"} else "auto"


def module_rows(modules: Any) -> list[dict[str, Any]]:
    result: list[dict[str, Any]] = []
    for module in modules:
        result.append({
            "module_id": str(getattr(module, "module_id", "")),
            "encoder_ticks": int_value(getattr(module, "encoder_ticks", 0)),
            "wheel_speed_rpm": float_value(getattr(module, "wheel_speed_rpm", 0.0)),
            "steer_angle_deg": float_value(getattr(module, "steer_angle_deg", 0.0)),
            "motor_current_amp": float_value(getattr(module, "motor_current_amp", 0.0)),
        })
    return result


class CsvAppender:
    def __init__(self, path: Path, fieldnames: list[str]) -> None:
        self.path = path
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.stream = self.path.open("w", newline="", encoding="utf-8")
        self.writer = csv.DictWriter(self.stream, fieldnames=fieldnames)
        self.writer.writeheader()
        self.count = 0

    def write_row(self, row: dict[str, Any]) -> None:
        self.writer.writerow(row)
        self.stream.flush()
        self.count += 1

    def close(self) -> None:
        self.stream.close()


class ChassisSessionCapture:
    def __init__(
        self,
        node: Any,
        config: dict[str, Any],
        paths: dict[str, Path],
        args: argparse.Namespace,
    ) -> None:
        self.node = node
        self.config = config
        self.paths = paths
        self.args = args
        self.started_us = now_us()
        self.chassis_csv = CsvAppender(paths["chassis_motion_data_file"], CHASSIS_MOTION_FIELDS)
        self.command_csv = CsvAppender(paths["actuator_command_file"], ACTUATOR_COMMAND_FIELDS)
        self.truth_csv = CsvAppender(paths["truth_trajectory_file"], TRUTH_TRAJECTORY_FIELDS)

    def close(self) -> None:
        self.chassis_csv.close()
        self.command_csv.close()
        self.truth_csv.close()

    def write_planned_command_if_needed(self) -> None:
        if self.args.command_source not in ("auto", "planned"):
            return
        if self.command_csv.count > 0:
            return
        self.command_csv.write_row({
            "command_timestamp_us": self.started_us,
            "command_id": self.args.command_id or f"{self.config.get('session_id', 'session')}_planned_command",
            "source": "planned_command",
            "target_speed_ms": self.args.target_speed_ms or synthetic_value(self.config, "target_speed_ms", 0.1),
            "target_distance_m": self.args.target_distance_m or synthetic_value(self.config, "target_distance_m", 0.2),
            "target_accel_ms2": 0.0,
            "target_steering_angle_rad": self.args.target_steering_angle_rad,
            "target_steering_rate_rads": 0.0,
            "brake_command": self.args.brake_command,
            "throttle_command": 0.0,
            "command_timeout_sec": 0.0,
        })

    def on_chassis_telemetry(self, msg: Any) -> None:
        chassis_type = normalize_chassis_type_name(getattr(getattr(msg, "chassis_type", None), "value", ""))
        self.chassis_csv.write_row({
            "hardware_timestamp_us": int_value(getattr(msg, "hardware_timestamp_us", 0), now_us()),
            "chassis_type": chassis_type,
            "odom_x_m": float_value(getattr(msg, "odom_x_m", 0.0)),
            "odom_y_m": float_value(getattr(msg, "odom_y_m", 0.0)),
            "odom_yaw_rad": float_value(getattr(msg, "odom_yaw_rad", 0.0)),
            "linear_velocity_ms": float_value(getattr(msg, "linear_velocity_ms", 0.0)),
            "angular_velocity_rads": float_value(getattr(msg, "angular_velocity_rads", 0.0)),
            "estop_engaged": int(bool(getattr(msg, "estop_engaged", False))),
            "driver_error_code": int_value(getattr(msg, "driver_error_code", 0)),
            "active_job_id": str(getattr(msg, "active_job_id", "")),
            "lateral_slip_estimate": float_value(getattr(msg, "lateral_slip_estimate", 0.0)),
            "curvature_estimate": float_value(getattr(msg, "curvature_estimate", 0.0)),
            "module_states_json": json.dumps(module_rows(getattr(msg, "modules", [])), ensure_ascii=False),
        })

    def on_external_pose(self, msg: Any) -> None:
        pose = getattr(msg, "workshop_pose", None)
        self.truth_csv.write_row({
            "hardware_timestamp_us": int_value(getattr(msg, "hardware_timestamp_us", 0), now_us()),
            "pose_valid": int(bool(getattr(msg, "pose_valid", False))),
            "x_m": float_value(getattr(pose, "x_m", 0.0)),
            "y_m": float_value(getattr(pose, "y_m", 0.0)),
            "z_m": float_value(getattr(pose, "z_m", 0.0)),
            "roll_rad": float_value(getattr(pose, "roll_rad", 0.0)),
            "pitch_rad": float_value(getattr(pose, "pitch_rad", 0.0)),
            "yaw_rad": float_value(getattr(pose, "yaw_rad", 0.0)),
            "position_stddev_m": float_value(getattr(msg, "position_stddev_m", 0.0)),
            "yaw_stddev_rad": float_value(getattr(msg, "yaw_stddev_rad", 0.0)),
            "tracking_loss_ratio": float_value(getattr(msg, "tracking_loss_ratio", 0.0)),
            "time_sync_offset_ms": float_value(getattr(msg, "time_sync_offset_ms", 0.0)),
            "quality_score": float_value(getattr(msg, "quality_score", 0.0)),
            "observed_target_count": int_value(getattr(msg, "observed_target_count", 0)),
            "reference_source_name": str(getattr(msg, "reference_source_name", "")),
            "active_job_id": str(getattr(msg, "active_job_id", "")),
        })

    def on_ackermann_command(self, msg: Any) -> None:
        if self.args.command_source == "planned":
            return
        self.command_csv.write_row({
            "command_timestamp_us": int_value(getattr(msg, "command_timestamp_us", 0), now_us()),
            "command_id": str(getattr(msg, "command_id", "")),
            "source": str(getattr(msg, "source", "")),
            "target_speed_ms": float_value(getattr(msg, "target_speed_ms", 0.0)),
            "target_distance_m": self.args.target_distance_m,
            "target_accel_ms2": float_value(getattr(msg, "target_accel_ms2", 0.0)),
            "target_steering_angle_rad": float_value(getattr(msg, "target_steering_angle_rad", 0.0)),
            "target_steering_rate_rads": float_value(getattr(msg, "target_steering_rate_rads", 0.0)),
            "brake_command": float_value(getattr(msg, "brake_command", 0.0)),
            "throttle_command": float_value(getattr(msg, "throttle_command", 0.0)),
            "command_timeout_sec": float_value(getattr(msg, "command_timeout_sec", 0.0)),
        })


def make_target_paths(config: dict[str, Any]) -> dict[str, Path]:
    session_dir = Path(str(config["session_dir"])).expanduser().resolve(strict=False)
    files = config["files"]
    return {
        "chassis_motion_data_file": resolve_path(session_dir, str(files["chassis_motion_data_file"])),
        "actuator_command_file": resolve_path(session_dir, str(files["actuator_command_file"])),
        "truth_trajectory_file": resolve_path(session_dir, str(files["truth_trajectory_file"])),
        "diagnostics_file": resolve_path(session_dir, str(files["diagnostics_file"])),
    }


def finalize_capture(
    config: dict[str, Any],
    target_paths: dict[str, Path],
    config_path: Path,
    args: argparse.Namespace,
) -> int:
    contract_validation = validate_chassis_dataset_files(target_paths)
    if not contract_validation.ok:
        summary = {
            "session_id": config.get("session_id", ""),
            "site_id": config.get("site_id", ""),
            "vehicle_id": config.get("vehicle_id", ""),
            "chassis_type": config.get("chassis_type", ""),
            "contract_errors": contract_validation.errors,
            "contract_warnings": contract_validation.warnings,
        }
        target_paths["diagnostics_file"].parent.mkdir(parents=True, exist_ok=True)
        target_paths["diagnostics_file"].write_text(
            json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
        for error in contract_validation.errors:
            print(f"[错误] {error}", file=sys.stderr)
        print(f"[错误] 已写入诊断文件: {target_paths['diagnostics_file']}", file=sys.stderr)
        return 1

    chassis_rows = read_csv_rows(target_paths["chassis_motion_data_file"])
    command_rows = read_csv_rows(target_paths["actuator_command_file"])
    truth_rows = read_csv_rows(target_paths["truth_trajectory_file"])
    summary = build_summary(chassis_rows, command_rows, truth_rows, config)

    summary_validation = validate_summary(summary, config)
    summary["contract_errors"] = contract_validation.errors
    summary["contract_warnings"] = contract_validation.warnings
    summary["validation_errors"] = summary_validation.errors
    summary["validation_warnings"] = summary_validation.warnings
    target_paths["diagnostics_file"].parent.mkdir(parents=True, exist_ok=True)
    target_paths["diagnostics_file"].write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    if summary_validation.errors and not args.allow_incomplete:
        for error in summary_validation.errors:
            print(f"[错误] {error}", file=sys.stderr)
        print(f"[错误] 已写入诊断文件: {target_paths['diagnostics_file']}", file=sys.stderr)
        return 1

    for warning in contract_validation.warnings:
        print(f"[WARN] {warning}", file=sys.stderr)
    for warning in summary_validation.warnings:
        print(f"[WARN] {warning}", file=sys.stderr)

    dataset_index = build_dataset_index(config, summary, target_paths, config_path)
    dataset_index_path = Path(str(config["dataset_index_path"])).expanduser().resolve(strict=False)
    write_yaml(dataset_index_path, dataset_index)
    print(json.dumps(summary, ensure_ascii=False, separators=(",", ":")))
    print(f"[OK] 已写入底盘采集数据集索引: {dataset_index_path}", file=sys.stderr)
    return 0


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

    try:
        import rclpy
        from calibration_chassis_interfaces.msg import ChassisTelemetry
        from calibration_external_localization_interfaces.msg import ExternalLocalizationTelemetry
    except ImportError as exc:
        print(f"[错误] 缺少 ROS 2 Python 依赖或标定消息包，无法现场采集: {exc}", file=sys.stderr)
        return 1

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

    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.init()
    node = rclpy.create_node("workshop_chassis_session_capture")
    capture = ChassisSessionCapture(node, config, target_paths, args)
    node.create_subscription(ChassisTelemetry, chassis_topic, capture.on_chassis_telemetry, 50)
    node.create_subscription(ExternalLocalizationTelemetry, external_topic, capture.on_external_pose, 50)
    if AckermannDriveCommand is not None and command_topic and args.command_source in ("auto", "topic"):
        node.create_subscription(AckermannDriveCommand, command_topic, capture.on_ackermann_command, 50)
    elif args.command_source == "topic":
        print("[错误] 已要求从命令话题采集，但命令消息包或话题配置不可用。", file=sys.stderr)
        capture.close()
        node.destroy_node()
        rclpy.shutdown()
        return 1

    print(f"[*] 底盘遥测采集: {chassis_topic}", file=sys.stderr)
    print(f"[*] 外部真值采集: {external_topic}", file=sys.stderr)
    if command_topic:
        print(f"[*] 执行器命令采集: {command_topic}", file=sys.stderr)

    deadline = time.monotonic() + args.duration_sec if args.duration_sec > 0.0 else None
    try:
        while rclpy.ok():
            rclpy.spin_once(node, timeout_sec=0.1)
            if deadline is not None and time.monotonic() >= deadline:
                break
            if args.max_chassis_samples > 0 and capture.chassis_csv.count >= args.max_chassis_samples:
                break
    except KeyboardInterrupt:
        pass
    finally:
        capture.write_planned_command_if_needed()
        capture.close()
        node.destroy_node()
        rclpy.shutdown()

    return finalize_capture(config, target_paths, config_path, args)


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(f"[错误] {exc}", file=sys.stderr)
        raise SystemExit(1)
