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

from __future__ import annotations

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

from control_data_common import (
    CHASSIS_RESPONSE_FIELDS,
    CONTROL_EVALUATION_FIELDS,
    REFERENCE_SIGNAL_FIELDS,
    TRUTH_TRAJECTORY_FIELDS,
    apply_cli_overrides,
    build_dataset_index,
    build_summary,
    has_placeholder,
    load_config,
    read_csv_rows,
    resolve_path,
    validate_config,
    validate_control_dataset_files,
    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("--control-axis", default="", help="覆盖 control_axis。")
    parser.add_argument("--controller-algorithm", default="", help="覆盖 controller_algorithm。")
    parser.add_argument("--control-role", default="", help="覆盖 control_role。")
    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-type", default="", help="覆盖 chassis_type。")
    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="本次采集对应的运控评估任务类型。")
    parser.add_argument("--duration-sec", type=float, default=0.0, help="采集时长；0 表示按 Ctrl+C 结束。")
    parser.add_argument("--max-control-samples", type=int, default=0, help="达到该运控样本数后自动结束；0 表示不限。")
    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 task_type_from_config(config: dict[str, Any], args: argparse.Namespace) -> str:
    if args.task_type:
        return args.task_type
    synthetic = config.get("synthetic", {}) or {}
    if isinstance(synthetic, dict):
        value = str(synthetic.get("task_type", "") or "").strip()
        if value:
            return value
    return "trajectory_tracking"


def average_module_steer_rad(modules: Any) -> float:
    values: list[float] = []
    for module in modules:
        values.append(math.radians(float_value(getattr(module, "steer_angle_deg", 0.0))))
    if not values:
        return 0.0
    return sum(values) / len(values)


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 ControlSessionCapture:
    def __init__(
        self,
        config: dict[str, Any],
        paths: dict[str, Path],
        args: argparse.Namespace,
    ) -> None:
        self.config = config
        self.paths = paths
        self.args = args
        self.task_type = task_type_from_config(config, args)
        self.chassis_type = str(config.get("chassis_type", "ackermann") or "ackermann")
        self.control_axis = str(config.get("control_axis", "combined") or "combined")
        self.controller_algorithm = str(config.get("controller_algorithm", "mpc") or "mpc")
        self.control_role = str(config.get("control_role", "path_tracking_outer_loop") or "path_tracking_outer_loop")
        self.reference_index = 0
        self.control_csv = CsvAppender(paths["control_evaluation_data_file"], CONTROL_EVALUATION_FIELDS)
        self.reference_csv = CsvAppender(paths["reference_signal_file"], REFERENCE_SIGNAL_FIELDS)
        self.chassis_csv = CsvAppender(paths["chassis_response_file"], CHASSIS_RESPONSE_FIELDS)
        self.truth_csv = CsvAppender(paths["truth_trajectory_file"], TRUTH_TRAJECTORY_FIELDS)

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

    def on_control_telemetry(self, msg: Any) -> None:
        stamp = int_value(getattr(msg, "hardware_timestamp_us", 0), now_us())
        measured_x = float_value(getattr(msg, "odom_x_m", 0.0))
        measured_y = float_value(getattr(msg, "odom_y_m", 0.0))
        measured_yaw = float_value(getattr(msg, "odom_yaw_rad", 0.0))
        measured_speed = float_value(getattr(msg, "linear_velocity_ms", 0.0))
        lateral_error = float_value(getattr(msg, "lateral_error_m", 0.0))
        heading_error = float_value(getattr(msg, "heading_error_rad", 0.0))
        speed_error = float_value(getattr(msg, "speed_error_ms", 0.0))
        target_velocity = measured_speed + speed_error
        target_y = measured_y - lateral_error
        target_yaw = measured_yaw - heading_error
        active_job_id = str(getattr(msg, "active_job_id", ""))

        self.control_csv.write_row({
            "hardware_timestamp_us": stamp,
            "chassis_type": self.chassis_type,
            "task_type": self.task_type,
            "control_axis": self.control_axis,
            "controller_algorithm": self.controller_algorithm,
            "control_role": self.control_role,
            "target_x_m": measured_x,
            "target_y_m": target_y,
            "target_yaw_rad": target_yaw,
            "target_velocity_ms": target_velocity,
            "measured_x_m": measured_x,
            "measured_y_m": measured_y,
            "measured_yaw_rad": measured_yaw,
            "measured_velocity_ms": measured_speed,
            "lateral_error_m": lateral_error,
            "heading_error_rad": heading_error,
            "speed_error_ms": speed_error,
            "control_output_steer_rad": float_value(getattr(msg, "steering_output", 0.0)),
            "control_output_accel_ms2": float_value(getattr(msg, "throttle_output", 0.0)),
            "control_output_brake": float_value(getattr(msg, "brake_output", 0.0)),
            "saturation_flag": int(bool(getattr(msg, "saturation_flag", False))),
            "active_job_id": active_job_id,
            "diagnostics_json": json.dumps({
                "parameter_version": str(getattr(msg, "parameter_version", "")),
                "pose_source_name": str(getattr(msg, "pose_source_name", "")),
                "pose_from_external_truth": bool(getattr(msg, "pose_from_external_truth", False)),
            }, ensure_ascii=False),
        })
        self.reference_csv.write_row({
            "reference_timestamp_us": stamp,
            "task_type": self.task_type,
            "sequence_index": self.reference_index,
            "target_x_m": measured_x,
            "target_y_m": target_y,
            "target_yaw_rad": target_yaw,
            "target_velocity_ms": target_velocity,
            "target_accel_ms2": synthetic_value(self.config, "target_accel_ms2", 0.0),
            "curvature_1pm": 0.0,
            "stop_required": 1 if self.task_type == "stop_accuracy" else 0,
            "active_job_id": active_job_id,
        })
        self.reference_index += 1

    def on_chassis_telemetry(self, msg: Any) -> None:
        self.chassis_csv.write_row({
            "hardware_timestamp_us": int_value(getattr(msg, "hardware_timestamp_us", 0), now_us()),
            "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)),
            "steering_angle_rad": average_module_steer_rad(getattr(msg, "modules", [])),
            "accel_ms2": float_value(getattr(msg, "accel_ms2", 0.0)),
            "brake_state": int(bool(getattr(msg, "brake_state", False))),
            "driver_error_code": int_value(getattr(msg, "driver_error_code", 0)),
            "active_job_id": str(getattr(msg, "active_job_id", "")),
        })

    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 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 {
        "control_evaluation_data_file": resolve_path(session_dir, str(files["control_evaluation_data_file"])),
        "reference_signal_file": resolve_path(session_dir, str(files["reference_signal_file"])),
        "chassis_response_file": resolve_path(session_dir, str(files["chassis_response_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_control_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", ""),
            "control_axis": config.get("control_axis", ""),
            "controller_algorithm": config.get("controller_algorithm", ""),
            "control_role": config.get("control_role", ""),
            "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

    control_rows = read_csv_rows(target_paths["control_evaluation_data_file"])
    reference_rows = read_csv_rows(target_paths["reference_signal_file"])
    chassis_rows = read_csv_rows(target_paths["chassis_response_file"])
    truth_rows = read_csv_rows(target_paths["truth_trajectory_file"])
    summary = build_summary(control_rows, reference_rows, chassis_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)
    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_control_interfaces.msg import ControlTelemetry
        from calibration_external_localization_interfaces.msg import ExternalLocalizationTelemetry
    except ImportError as exc:
        print(f"[错误] 缺少 ROS 2 Python 依赖或标定消息包，无法现场采集: {exc}", file=sys.stderr)
        return 1

    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.init()
    node = rclpy.create_node("workshop_control_session_capture")
    capture = ControlSessionCapture(config, target_paths, args)
    node.create_subscription(ControlTelemetry, control_topic, capture.on_control_telemetry, 50)
    node.create_subscription(ChassisTelemetry, chassis_topic, capture.on_chassis_telemetry, 50)
    node.create_subscription(ExternalLocalizationTelemetry, external_topic, capture.on_external_pose, 50)

    print(f"[*] 运控遥测采集: {control_topic}", file=sys.stderr)
    print(f"[*] 底盘响应采集: {chassis_topic}", file=sys.stderr)
    print(f"[*] 外部真值采集: {external_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_control_samples > 0 and capture.control_csv.count >= args.max_control_samples:
                break
    except KeyboardInterrupt:
        pass
    finally:
        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)
