#!/usr/bin/env python3
"""把真实采集数据集索引转换为算法服务可读取的 ROS 2 参数文件。"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path
from typing import Any

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


URI_SCHEME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9+.-]*://")

SERVICE_INPUT_FIELDS = {
    "external_localization_service": {
        "index_section": "external",
        "fields": [
            "external_observation_files",
            "marker_observation_files",
            "time_sync_diagnostic_files",
        ],
    },
    "chassis_calibration_service": {
        "index_section": "chassis",
        "fields": [
            "chassis_motion_data_files",
            "actuator_command_files",
            "truth_trajectory_files",
        ],
    },
    "control_calibration_service": {
        "index_section": "control",
        "fields": [
            "control_evaluation_data_files",
            "reference_signal_files",
            "chassis_response_files",
            "truth_trajectory_files",
        ],
    },
    "sensor_calibration_service": {
        "index_section": "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",
        ],
    },
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="根据真实采集 dataset_index.yaml 生成 site_data_input.yaml。"
    )
    parser.add_argument(
        "dataset_index",
        help="真实采集或 ingest 生成的数据集索引文件。",
    )
    parser.add_argument(
        "-o",
        "--output",
        help="输出的 ROS 2 参数文件；不填写时输出到标准输出。",
    )
    parser.add_argument(
        "--dataset-root",
        help="覆盖索引文件中的 session.dataset_root，用于现场挂载路径重定向。",
    )
    parser.add_argument(
        "--no-index-reference",
        action="store_true",
        help="不自动把 dataset_index.yaml 写入 sensor 的 synchronized_dataset_files。",
    )
    parser.add_argument(
        "--strict",
        action="store_true",
        help="严格检查：要求时间窗口有效，并且至少存在一个数据文件引用。",
    )
    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)
    if data is None:
        return {}
    if not isinstance(data, dict):
        raise ValueError(f"{path} 的顶层结构必须是 YAML map。")
    return data


def is_uri(value: str) -> bool:
    return bool(URI_SCHEME_RE.match(value))


def normalize_dataset_root(index_path: Path, raw_root: str | None) -> Path:
    if not raw_root:
        return index_path.parent
    root = Path(raw_root).expanduser()
    if root.is_absolute():
        return root
    return (index_path.parent / root).resolve(strict=False)


def resolve_file_uri(raw_uri: str, dataset_root: Path) -> str:
    uri = raw_uri.strip()
    if not uri:
        return ""
    if is_uri(uri):
        return uri
    path = Path(uri).expanduser()
    if path.is_absolute():
        return str(path)
    return str((dataset_root / path).resolve(strict=False))


def extract_file_uri(entry: Any) -> str:
    if isinstance(entry, str):
        return entry
    if isinstance(entry, dict):
        for key in ("file_uri", "uri", "path", "file"):
            value = entry.get(key)
            if isinstance(value, str):
                return value
    raise ValueError(f"文件引用条目必须是字符串，或包含 file_uri/uri/path/file 的 map：{entry!r}")


def read_file_list(value: Any, dataset_root: Path) -> list[str]:
    if value is None:
        return []
    entries = value if isinstance(value, list) else [value]
    result: list[str] = []
    for entry in entries:
        file_uri = resolve_file_uri(extract_file_uri(entry), dataset_root)
        if file_uri and file_uri not in result:
            result.append(file_uri)
    return result


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


def module_data(index: dict[str, Any], section_name: str) -> dict[str, Any]:
    data_inputs = index.get("data_inputs")
    if data_inputs is None:
        data_inputs = index.get("files", {})
    if not isinstance(data_inputs, dict):
        raise ValueError("data_inputs 必须是 YAML map。")
    section = data_inputs.get(section_name, {})
    if section is None:
        return {}
    if not isinstance(section, dict):
        raise ValueError(f"data_inputs.{section_name} 必须是 YAML map。")
    return section


def window_value(session: dict[str, Any], section: dict[str, Any], key: str) -> int:
    return read_int(section.get(key, session.get(key, 0)), key)


def build_ros_params(
    index_path: Path,
    index: dict[str, Any],
    dataset_root_override: str | None,
    add_index_reference: bool,
) -> dict[str, Any]:
    session = index.get("session", {})
    if session is None:
        session = {}
    if not isinstance(session, dict):
        raise ValueError("session 必须是 YAML map。")

    dataset_root = normalize_dataset_root(
        index_path,
        dataset_root_override or session.get("dataset_root"),
    )

    output: dict[str, Any] = {}
    for service_name, config in SERVICE_INPUT_FIELDS.items():
        section_name = str(config["index_section"])
        section = module_data(index, section_name)
        params: dict[str, Any] = {
            "data_input.data_window_start_timestamp_us": window_value(
                session, section, "data_window_start_timestamp_us"
            ),
            "data_input.data_window_end_timestamp_us": window_value(
                session, section, "data_window_end_timestamp_us"
            ),
        }
        for field_name in config["fields"]:
            files = read_file_list(section.get(field_name), dataset_root)
            if (
                add_index_reference
                and section_name == "sensor"
                and field_name == "synchronized_dataset_files"
            ):
                index_uri = str(index_path.resolve(strict=False))
                if index_uri not in files:
                    files.append(index_uri)
            if files:
                params[f"data_input.{field_name}"] = files
        output[service_name] = {"ros__parameters": params}
    return output


def validate_output(output: dict[str, Any], strict: bool) -> None:
    if not strict:
        return
    file_count = 0
    for service_name, node_config in output.items():
        params = node_config["ros__parameters"]
        start_us = params["data_input.data_window_start_timestamp_us"]
        end_us = params["data_input.data_window_end_timestamp_us"]
        if start_us <= 0 or end_us <= 0 or end_us <= start_us:
            raise ValueError(f"{service_name} 的数据时间窗口无效。")
        file_count += sum(
            len(value)
            for key, value in params.items()
            if key.startswith("data_input.") and isinstance(value, list)
        )
    if file_count == 0:
        raise ValueError("严格检查要求至少存在一个数据文件引用。")


def main() -> int:
    args = parse_args()
    index_path = Path(args.dataset_index).expanduser().resolve(strict=False)
    index = load_yaml(index_path)
    output = build_ros_params(
        index_path,
        index,
        args.dataset_root,
        not args.no_index_reference,
    )
    validate_output(output, args.strict)
    rendered = yaml.safe_dump(output, sort_keys=False, allow_unicode=True)

    if args.output:
        output_path = Path(args.output).expanduser()
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(rendered, encoding="utf-8")
        print(f"已生成 ROS 2 数据输入参数文件：{output_path}", file=sys.stderr)
    else:
        print(rendered, end="")
    return 0


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