#!/usr/bin/env python3
"""校验传感器标定数据文件，并可选写入 dataset_index.yaml。"""

from __future__ import annotations

import argparse
import csv
import json
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

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


SENSOR_FIELDS = {
    "image_sample_files",
    "imu_sample_files",
    "pointcloud_sample_files",
    "laser_scan_sample_files",
    "target_detection_files",
    "robot_pose_sample_files",
    "synchronized_dataset_files",
}

IMU_SAMPLE_COLUMNS = {
    "stamp_us",
    "sensor_id",
    "frame_id",
    "ax_mps2",
    "ay_mps2",
    "az_mps2",
    "gx_rps",
    "gy_rps",
    "gz_rps",
}

ROBOT_POSE_COLUMNS = {
    "stamp_us",
    "parent_frame_id",
    "child_frame_id",
    "x_m",
    "y_m",
    "z_m",
    "qx",
    "qy",
    "qz",
    "qw",
}


def now_us() -> int:
    return int(datetime.now(timezone.utc).timestamp() * 1_000_000)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="校验车间电脑侧已落盘的传感器标定数据，并可选写入 dataset_index.yaml。"
    )
    parser.add_argument("--dataset-root", required=True, help="本次会话数据集根目录。")
    parser.add_argument("--dataset-index", default="", help="可选 dataset_index.yaml；填写后写入 data_inputs.sensor。")
    parser.add_argument("--session-id", default="", help="本次会话 ID；写 dataset_index 时使用。")
    parser.add_argument("--site-id", default="", help="现场或产线 ID；写 dataset_index 时使用。")
    parser.add_argument("--vehicle-id", default="", help="车辆 ID；写 dataset_index 时使用。")
    parser.add_argument("--start-us", type=int, default=0, help="数据窗口开始时间；不填写时从文件推断。")
    parser.add_argument("--end-us", type=int, default=0, help="数据窗口结束时间；不填写时从文件推断。")
    parser.add_argument("--image-index", action="append", default=[], help="登记并校验 images.yaml，可重复。")
    parser.add_argument("--imu-samples", action="append", default=[], help="登记并校验 imu_samples.csv，可重复。")
    parser.add_argument("--imu-segments", action="append", default=[], help="登记并校验 imu_segments.yaml，可重复。")
    parser.add_argument("--pointcloud-index", action="append", default=[], help="登记并校验 pointclouds.yaml，可重复。")
    parser.add_argument("--laser-scan-index", action="append", default=[], help="登记并校验 laser_scans.yaml，可重复。")
    parser.add_argument("--target-detections", action="append", default=[], help="登记并校验 target_detections.json，可重复。")
    parser.add_argument("--robot-poses", action="append", default=[], help="登记并校验 robot_poses.csv，可重复。")
    parser.add_argument("--synchronized-dataset", action="append", default=[], help="登记并校验 synchronized_sensor_dataset.yaml，可重复。")
    parser.add_argument("--validate-only", action="store_true", help="只校验文件，不写 dataset_index.yaml。")
    parser.add_argument("--overwrite", action="store_true", help="写 dataset_index 时覆盖已有 data_inputs.sensor。")
    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 write_yaml(path: Path, data: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(yaml.safe_dump(data, sort_keys=False, allow_unicode=True), encoding="utf-8")


def resolve_path(raw_path: str, dataset_root: Path) -> Path:
    path = Path(raw_path).expanduser()
    if path.is_absolute():
        return path
    return (dataset_root / path).resolve(strict=False)


def relative_to_root(path: Path, dataset_root: Path) -> str:
    try:
        return str(path.resolve(strict=False).relative_to(dataset_root.resolve(strict=False)))
    except ValueError:
        return str(path.resolve(strict=False))


def require_string(value: Any, field_name: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{field_name} 必须是非空字符串。")
    return value.strip()


def require_map(value: Any, field_name: str) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise ValueError(f"{field_name} 必须是 YAML/JSON map。")
    return value


def require_list(value: Any, field_name: str) -> list[Any]:
    if not isinstance(value, list):
        raise ValueError(f"{field_name} 必须是列表。")
    return value


def read_int(value: Any, field_name: str) -> int:
    try:
        return int(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{field_name} 必须是整数，当前值为 {value!r}。") from exc


def validate_schema_version(data: dict[str, Any], path: Path) -> None:
    if read_int(data.get("schema_version"), f"{path}.schema_version") != 1:
        raise ValueError(f"{path} 的 schema_version 必须为 1。")


def collect_stamp(stamps: list[int], value: Any, field_name: str) -> None:
    stamp_us = read_int(value, field_name)
    if stamp_us <= 0:
        raise ValueError(f"{field_name} 必须大于 0。")
    stamps.append(stamp_us)


def validate_index_frames(
    path: Path,
    data: dict[str, Any],
    list_key: str,
    stamp_field_name: str,
) -> list[int]:
    validate_schema_version(data, path)
    require_string(data.get("sensor_id"), f"{path}.sensor_id")
    require_string(data.get("frame_id"), f"{path}.frame_id")
    frames = require_list(data.get(list_key), f"{path}.{list_key}")
    if not frames:
        raise ValueError(f"{path}.{list_key} 不能为空。")
    stamps: list[int] = []
    for index, frame_value in enumerate(frames):
        frame = require_map(frame_value, f"{path}.{list_key}[{index}]")
        collect_stamp(stamps, frame.get("stamp_us"), f"{path}.{list_key}[{index}].stamp_us")
        require_string(frame.get("file"), f"{path}.{list_key}[{index}].file")
    if stamps != sorted(stamps):
        raise ValueError(f"{path}.{list_key} 必须按 {stamp_field_name} 递增排序。")
    return stamps


def validate_image_index(path: Path) -> list[int]:
    data = load_yaml(path)
    require_string(data.get("sensor_type"), f"{path}.sensor_type")
    return validate_index_frames(path, data, "frames", "stamp_us")


def validate_pointcloud_index(path: Path) -> list[int]:
    data = load_yaml(path)
    return validate_index_frames(path, data, "frames", "stamp_us")


def validate_laser_scan_index(path: Path) -> list[int]:
    data = load_yaml(path)
    return validate_index_frames(path, data, "scans", "stamp_us")


def validate_imu_segments(path: Path) -> list[int]:
    data = load_yaml(path)
    validate_schema_version(data, path)
    require_string(data.get("sensor_id"), f"{path}.sensor_id")
    require_string(data.get("frame_id"), f"{path}.frame_id")
    segments = require_list(data.get("segments"), f"{path}.segments")
    if not segments:
        raise ValueError(f"{path}.segments 不能为空。")
    stamps: list[int] = []
    for index, segment_value in enumerate(segments):
        segment = require_map(segment_value, f"{path}.segments[{index}]")
        segment_type = require_string(segment.get("segment_type"), f"{path}.segments[{index}].segment_type")
        if segment_type not in {"static", "motion"}:
            raise ValueError(f"{path}.segments[{index}].segment_type 必须是 static 或 motion。")
        start_us = read_int(segment.get("start_us"), f"{path}.segments[{index}].start_us")
        end_us = read_int(segment.get("end_us"), f"{path}.segments[{index}].end_us")
        if start_us <= 0 or end_us <= start_us:
            raise ValueError(f"{path}.segments[{index}] 的时间窗口无效。")
        stamps.extend([start_us, end_us])
    return stamps


def validate_json_detections(path: Path) -> list[int]:
    with path.open("r", encoding="utf-8") as stream:
        data = json.load(stream)
    if not isinstance(data, dict):
        raise ValueError(f"{path} 顶层必须是 JSON object。")
    validate_schema_version(data, path)
    require_string(data.get("sensor_id"), f"{path}.sensor_id")
    require_string(data.get("target_type"), f"{path}.target_type")
    detections = require_list(data.get("detections"), f"{path}.detections")
    if not detections:
        raise ValueError(f"{path}.detections 不能为空。")
    stamps: list[int] = []
    for index, detection_value in enumerate(detections):
        detection = require_map(detection_value, f"{path}.detections[{index}]")
        collect_stamp(stamps, detection.get("stamp_us"), f"{path}.detections[{index}].stamp_us")
        if "target_detected" not in detection:
            raise ValueError(f"{path}.detections[{index}].target_detected 不能为空。")
    return stamps


def csv_header(path: Path) -> list[str]:
    with path.open("r", encoding="utf-8", newline="") as stream:
        reader = csv.reader(stream)
        try:
            header = next(reader)
        except StopIteration as exc:
            raise ValueError(f"{path} 不能为空 CSV。") from exc
    return [item.strip() for item in header]


def validate_csv_columns(path: Path, required_columns: set[str]) -> list[int]:
    header = csv_header(path)
    missing = sorted(required_columns - set(header))
    if missing:
        raise ValueError(f"{path} 缺少 CSV 列：{', '.join(missing)}。")
    stamps: list[int] = []
    with path.open("r", encoding="utf-8", newline="") as stream:
        reader = csv.DictReader(stream)
        for row_index, row in enumerate(reader):
            collect_stamp(stamps, row.get("stamp_us"), f"{path}.row[{row_index}].stamp_us")
    if not stamps:
        raise ValueError(f"{path} 至少需要一行数据。")
    if stamps != sorted(stamps):
        raise ValueError(f"{path} 必须按 stamp_us 递增排序。")
    return stamps


def validate_synchronized_dataset(path: Path) -> list[int]:
    data = load_yaml(path)
    validate_schema_version(data, path)
    require_string(data.get("session_id"), f"{path}.session_id")
    require_string(data.get("vehicle_id"), f"{path}.vehicle_id")
    files = require_map(data.get("files"), f"{path}.files")
    if not any(files.get(field_name) for field_name in SENSOR_FIELDS):
        raise ValueError(f"{path}.files 至少需要包含一个传感器数据字段。")
    start_us = read_int(data.get("data_window_start_timestamp_us"), f"{path}.data_window_start_timestamp_us")
    end_us = read_int(data.get("data_window_end_timestamp_us"), f"{path}.data_window_end_timestamp_us")
    if start_us <= 0 or end_us <= start_us:
        raise ValueError(f"{path} 的数据窗口无效。")
    return [start_us, end_us]


def append_file(
    sensor_inputs: dict[str, list[str]],
    field_name: str,
    raw_files: list[str],
    dataset_root: Path,
    validator,
    stamps: list[int],
) -> None:
    for raw_file in raw_files:
        path = resolve_path(raw_file, dataset_root)
        if not path.exists():
            raise FileNotFoundError(f"文件不存在：{path}")
        stamps.extend(validator(path))
        value = relative_to_root(path, dataset_root)
        sensor_inputs.setdefault(field_name, [])
        if value not in sensor_inputs[field_name]:
            sensor_inputs[field_name].append(value)


def load_or_create_index(
    dataset_index_path: Path,
    dataset_root: Path,
    args: argparse.Namespace,
    start_us: int,
    end_us: int,
) -> dict[str, Any]:
    if dataset_index_path.exists():
        index = load_yaml(dataset_index_path)
    else:
        index = {"schema_version": 1, "data_inputs": {}}
    session = index.setdefault("session", {})
    if not isinstance(session, dict):
        raise ValueError("dataset_index.session 必须是 YAML map。")
    session["session_id"] = args.session_id or session.get("session_id") or dataset_root.name
    session["site_id"] = args.site_id or session.get("site_id", "")
    session["vehicle_id"] = args.vehicle_id or session.get("vehicle_id", "")
    session["dataset_root"] = str(dataset_root)
    session["data_window_start_timestamp_us"] = start_us
    session["data_window_end_timestamp_us"] = end_us
    index.setdefault("schema_version", 1)
    index.setdefault("data_inputs", {})
    return index


def merge_sensor_inputs(index: dict[str, Any], sensor_inputs: dict[str, list[str]], overwrite: bool) -> None:
    data_inputs = index.setdefault("data_inputs", {})
    if not isinstance(data_inputs, dict):
        raise ValueError("dataset_index.data_inputs 必须是 YAML map。")
    current = {} if overwrite else data_inputs.setdefault("sensor", {})
    if not isinstance(current, dict):
        raise ValueError("dataset_index.data_inputs.sensor 必须是 YAML map。")
    for field_name, values in sensor_inputs.items():
        existing_value = [] if overwrite else current.get(field_name, []) or []
        if isinstance(existing_value, str):
            merged = [existing_value]
        else:
            merged = list(existing_value)
        for value in values:
            if value not in merged:
                merged.append(value)
        current[field_name] = merged
    data_inputs["sensor"] = current


def main() -> int:
    args = parse_args()
    dataset_root = Path(args.dataset_root).expanduser().resolve(strict=False)
    sensor_inputs: dict[str, list[str]] = {}
    stamps: list[int] = []

    append_file(sensor_inputs, "image_sample_files", args.image_index, dataset_root, validate_image_index, stamps)
    append_file(sensor_inputs, "imu_sample_files", args.imu_samples, dataset_root, lambda path: validate_csv_columns(path, IMU_SAMPLE_COLUMNS), stamps)
    append_file(sensor_inputs, "imu_sample_files", args.imu_segments, dataset_root, validate_imu_segments, stamps)
    append_file(sensor_inputs, "pointcloud_sample_files", args.pointcloud_index, dataset_root, validate_pointcloud_index, stamps)
    append_file(sensor_inputs, "laser_scan_sample_files", args.laser_scan_index, dataset_root, validate_laser_scan_index, stamps)
    append_file(sensor_inputs, "target_detection_files", args.target_detections, dataset_root, validate_json_detections, stamps)
    append_file(sensor_inputs, "robot_pose_sample_files", args.robot_poses, dataset_root, lambda path: validate_csv_columns(path, ROBOT_POSE_COLUMNS), stamps)
    append_file(sensor_inputs, "synchronized_dataset_files", args.synchronized_dataset, dataset_root, validate_synchronized_dataset, stamps)

    if not sensor_inputs:
        raise ValueError("至少需要提供一个传感器数据文件。")

    start_us = args.start_us or min(stamps)
    end_us = args.end_us or max(stamps)
    if end_us <= start_us:
        end_us = start_us + 1

    print(f"[OK] 传感器数据文件校验通过，文件组数：{sum(len(v) for v in sensor_inputs.values())}")
    print(f"[OK] 数据窗口：{start_us} -> {end_us}")

    if args.validate_only:
        return 0

    if not args.dataset_index:
        raise ValueError("未设置 --validate-only 时必须提供 --dataset-index。")
    dataset_index_path = Path(args.dataset_index).expanduser().resolve(strict=False)
    index = load_or_create_index(dataset_index_path, dataset_root, args, start_us, end_us)
    merge_sensor_inputs(index, sensor_inputs, overwrite=args.overwrite)
    write_yaml(dataset_index_path, index)
    print(f"[OK] 已写入 dataset_index.yaml：{dataset_index_path}")
    return 0


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