#!/usr/bin/env python3
"""生成传感器标定待提交参数包。"""

from __future__ import annotations

import argparse
import hashlib
import sys
import time
from pathlib import Path
from typing import Any

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


TASK_PARAM_BRANCH = {
    "camera_intrinsic": "camera_intrinsics",
    "imu_intrinsic": "imu_intrinsics",
    "sensor_extrinsic": "sensor_extrinsics",
    "hand_eye": "hand_eye",
}

SUPPORTED_SUBTYPES = {
    "camera_intrinsic": {"front_camera_intrinsic", "downward_camera_intrinsic"},
    "imu_intrinsic": {"imu_intrinsic"},
    "sensor_extrinsic": {
        "front_camera_extrinsic",
        "downward_camera_extrinsic",
        "imu_extrinsic",
        "lidar_2d_extrinsic",
        "lidar_3d_extrinsic",
    },
    "hand_eye": {"eye_in_hand", "eye_to_hand"},
}


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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="把传感器算法输出整理成现场待审批/待提交参数包。")
    parser.add_argument("--dataset-index", required=True, help="本次会话 dataset_index.yaml。")
    parser.add_argument("--estimated-params", required=True, help="算法输出的传感器参数 YAML。")
    parser.add_argument("--output", default="", help="输出待提交参数包；为空时写到会话目录 sensor/pending_sensor_commit.yaml。")
    parser.add_argument("--commit-reason", default="传感器标定算法输出待提交", help="写入参数包的提交原因。")
    parser.add_argument("--operator-id", default="", help="生成待提交包的操作员 ID。")
    parser.add_argument("--previous-parameter-version", default="", help="当前车端已生效传感器参数版本，用于回滚引用。")
    parser.add_argument("--persistent-write", action=argparse.BooleanOptionalAction, default=True, help="最终提交到车端时是否持久化。")
    parser.add_argument("--require-manual-approval", action=argparse.BooleanOptionalAction, default=True, help="是否要求人工审批。")
    parser.add_argument("--update-dataset-index", action=argparse.BooleanOptionalAction, default=True, help="是否把待提交包路径写回 dataset_index.yaml。")
    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 sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


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 map。")
    return value


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


def has_file_values(value: Any) -> bool:
    if value is None:
        return False
    if isinstance(value, list):
        return bool(value)
    return bool(str(value).strip())


def require_number(value: Any, field_name: str) -> float:
    if value in (None, ""):
        raise ValueError(f"{field_name} 不能为空。")
    try:
        return float(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{field_name} 必须是数字，当前值为 {value!r}。") from exc


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


def require_number_list(value: Any, field_name: str, length: int | None = None) -> None:
    values = require_list(value, field_name)
    if not values:
        raise ValueError(f"{field_name} 不能为空。")
    if length is not None and len(values) != length:
        raise ValueError(f"{field_name} 必须包含 {length} 个数字。")
    for index, item in enumerate(values):
        require_number(item, f"{field_name}[{index}]")


def require_number_matrix(value: Any, field_name: str, row_count: int, column_count: int) -> None:
    rows = require_list(value, field_name)
    if len(rows) != row_count:
        raise ValueError(f"{field_name} 必须包含 {row_count} 行。")
    for row_index, row in enumerate(rows):
        require_number_list(row, f"{field_name}[{row_index}]", column_count)


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


def validate_dataset_index(index: dict[str, Any]) -> dict[str, Any]:
    if int(index.get("schema_version", 0) or 0) != 1:
        raise ValueError("dataset_index.schema_version 必须为 1。")
    session = require_map(index.get("session"), "dataset_index.session")
    for field_name in ("session_id", "vehicle_id", "dataset_root"):
        require_string(session.get(field_name), f"dataset_index.session.{field_name}")
    sensor_inputs = require_map(
        require_map(index.get("data_inputs"), "dataset_index.data_inputs").get("sensor"),
        "dataset_index.data_inputs.sensor",
    )
    if not any(has_file_values(value) for value in sensor_inputs.values()):
        raise ValueError("dataset_index.data_inputs.sensor 至少需要一个传感器数据文件引用。")
    return session


def validate_camera_intrinsics(params: dict[str, Any]) -> None:
    require_integer(params.get("image_width"), "estimated_params.camera_intrinsics.image_width")
    require_integer(params.get("image_height"), "estimated_params.camera_intrinsics.image_height")
    require_string(params.get("camera_model"), "estimated_params.camera_intrinsics.camera_model")
    require_string(params.get("distortion_model"), "estimated_params.camera_intrinsics.distortion_model")
    for field_name in ("fx", "fy", "cx", "cy"):
        require_number(params.get(field_name), f"estimated_params.camera_intrinsics.{field_name}")
    require_number_list(
        params.get("distortion_coefficients"),
        "estimated_params.camera_intrinsics.distortion_coefficients",
    )


def validate_imu_intrinsics(params: dict[str, Any]) -> None:
    require_number_list(params.get("gyro_bias_xyz_rads"), "estimated_params.imu_intrinsics.gyro_bias_xyz_rads", 3)
    require_number_list(params.get("accel_bias_xyz_ms2"), "estimated_params.imu_intrinsics.accel_bias_xyz_ms2", 3)
    for field_name in (
        "gyro_noise_density_rads_sqrt_hz",
        "accel_noise_density_ms2_sqrt_hz",
        "gyro_random_walk_rads_s_sqrt_hz",
        "accel_random_walk_ms3_sqrt_hz",
    ):
        require_number(params.get(field_name), f"estimated_params.imu_intrinsics.{field_name}")
    if "scale_matrix" in params:
        require_number_matrix(params["scale_matrix"], "estimated_params.imu_intrinsics.scale_matrix", 3, 3)


def validate_transform_params(params: dict[str, Any], field_prefix: str) -> None:
    require_number_list(params.get("translation_xyz_m"), f"{field_prefix}.translation_xyz_m", 3)
    require_number_list(params.get("rotation_xyzw"), f"{field_prefix}.rotation_xyzw", 4)


def validate_sensor_extrinsics(params: dict[str, Any]) -> None:
    require_string(params.get("parent_frame_id"), "estimated_params.sensor_extrinsics.parent_frame_id")
    require_string(params.get("child_frame_id"), "estimated_params.sensor_extrinsics.child_frame_id")
    validate_transform_params(params, "estimated_params.sensor_extrinsics")


def validate_hand_eye(params: dict[str, Any]) -> None:
    require_string(params.get("base_frame_id"), "estimated_params.hand_eye.base_frame_id")
    require_string(params.get("tool_frame_id"), "estimated_params.hand_eye.tool_frame_id")
    require_string(params.get("camera_frame_id"), "estimated_params.hand_eye.camera_frame_id")
    validate_transform_params(params, "estimated_params.hand_eye")


def validate_quality(params: dict[str, Any]) -> None:
    quality = require_map(params.get("quality"), "estimated_params.quality")
    if not bool(quality.get("data_quality_passed", False)):
        raise ValueError("estimated_params.quality.data_quality_passed 必须为 true。")
    if not bool(quality.get("suitable_for_commit", False)):
        raise ValueError("estimated_params.quality.suitable_for_commit 必须为 true。")
    validation_summary = require_map(
        quality.get("validation_summary"),
        "estimated_params.quality.validation_summary",
    )
    if "auto_acceptance_passed" in validation_summary and not bool(validation_summary["auto_acceptance_passed"]):
        raise ValueError("estimated_params.quality.validation_summary.auto_acceptance_passed 必须为 true。")


def validate_estimated_params(
    params: dict[str, Any],
    session: dict[str, Any],
) -> tuple[str, str, str, str, dict[str, Any]]:
    if int(params.get("schema_version", 0) or 0) != 1:
        raise ValueError("estimated_params.schema_version 必须为 1。")
    parameter_version = require_string(params.get("parameter_version"), "estimated_params.parameter_version")
    selected_task = require_string(params.get("selected_task"), "estimated_params.selected_task")
    if selected_task not in TASK_PARAM_BRANCH:
        raise ValueError(f"estimated_params.selected_task 必须是 {sorted(TASK_PARAM_BRANCH)} 之一。")
    task_subtype = require_string(params.get("task_subtype"), "estimated_params.task_subtype")
    if task_subtype not in SUPPORTED_SUBTYPES[selected_task]:
        raise ValueError(f"estimated_params.task_subtype={task_subtype} 不适用于 selected_task={selected_task}。")
    sensor_id = require_string(params.get("sensor_id"), "estimated_params.sensor_id")
    vehicle_id = str(params.get("vehicle_id", "") or "").strip()
    if vehicle_id and vehicle_id != session.get("vehicle_id"):
        raise ValueError("estimated_params.vehicle_id 与 dataset_index.session.vehicle_id 不一致。")

    estimated = require_map(params.get("estimated_params"), "estimated_params.estimated_params")
    branch_name = TASK_PARAM_BRANCH[selected_task]
    selected_params = require_map(estimated.get(branch_name), f"estimated_params.estimated_params.{branch_name}")
    if selected_task == "camera_intrinsic":
        validate_camera_intrinsics(selected_params)
    elif selected_task == "imu_intrinsic":
        validate_imu_intrinsics(selected_params)
    elif selected_task == "sensor_extrinsic":
        validate_sensor_extrinsics(selected_params)
    elif selected_task == "hand_eye":
        validate_hand_eye(selected_params)
    validate_quality(params)
    return parameter_version, selected_task, task_subtype, sensor_id, {branch_name: selected_params}


def resolve_output_path(args: argparse.Namespace, session: dict[str, Any]) -> Path:
    if args.output:
        return Path(args.output).expanduser().resolve(strict=False)
    dataset_root = Path(str(session["dataset_root"])).expanduser().resolve(strict=False)
    return dataset_root / "sensor" / "pending_sensor_commit.yaml"


def build_commit_package(
    dataset_index_path: Path,
    estimated_params_path: Path,
    index: dict[str, Any],
    params: dict[str, Any],
    args: argparse.Namespace,
    output_path: Path,
) -> dict[str, Any]:
    session = index["session"]
    parameter_version, selected_task, task_subtype, sensor_id, estimated = validate_estimated_params(params, session)
    dataset_root = Path(str(session["dataset_root"])).expanduser().resolve(strict=False)
    created_us = now_us()
    return {
        "schema_version": 1,
        "commit_state": "pending_manual_approval" if args.require_manual_approval else "pending_vehicle_commit",
        "created_timestamp_us": created_us,
        "session": {
            "session_id": session["session_id"],
            "site_id": session.get("site_id", ""),
            "vehicle_id": session["vehicle_id"],
            "dataset_root": session["dataset_root"],
        },
        "source": {
            "dataset_index_file": relative_to_root(dataset_index_path, dataset_root),
            "estimated_params_file": relative_to_root(estimated_params_path, dataset_root),
            "estimated_params_digest": {
                "checksum_type": "sha256",
                "checksum_value": sha256_file(estimated_params_path),
            },
        },
        "approval": {
            "required": bool(args.require_manual_approval),
            "approved": False,
            "operator_id": args.operator_id,
            "approved_timestamp_us": 0,
        },
        "commit_request": {
            "parameter_version": parameter_version,
            "selected_task": selected_task,
            "task_subtype": task_subtype,
            "sensor_id": sensor_id,
            "commit_reason": args.commit_reason,
            "persistent_write": bool(args.persistent_write),
            "estimated_params": estimated,
        },
        "rollback": {
            "previous_parameter_version": args.previous_parameter_version,
            "rollback_required_on_vehicle_commit_failure": True,
            "rollback_verified": False,
        },
        "trace": {
            "pending_commit_file": relative_to_root(output_path, dataset_root),
            "generated_by": "stage_sensor_parameter_commit.py",
        },
    }


def update_dataset_index(index_path: Path, index: dict[str, Any], output_path: Path) -> None:
    session = index["session"]
    dataset_root = Path(str(session["dataset_root"])).expanduser().resolve(strict=False)
    metadata = index.setdefault("metadata", {})
    commit_package = load_yaml(output_path)
    source = require_map(commit_package.get("source"), "pending_commit.source")
    digest = require_map(source.get("estimated_params_digest"), "pending_commit.source.estimated_params_digest")
    approval = require_map(commit_package.get("approval"), "pending_commit.approval")
    commit_request = require_map(commit_package.get("commit_request"), "pending_commit.commit_request")
    metadata["sensor_estimated_params_file"] = source.get("estimated_params_file", "")
    metadata["sensor_estimated_params_checksum_type"] = digest.get("checksum_type", "")
    metadata["sensor_estimated_params_checksum_value"] = digest.get("checksum_value", "")
    metadata["sensor_pending_commit_file"] = relative_to_root(output_path, dataset_root)
    metadata["sensor_pending_commit_state"] = commit_package.get("commit_state", "")
    metadata["sensor_pending_parameter_version"] = commit_request.get("parameter_version", "")
    metadata["sensor_pending_commit_selected_task"] = commit_request.get("selected_task", "")
    metadata["sensor_pending_commit_task_subtype"] = commit_request.get("task_subtype", "")
    metadata["sensor_pending_commit_sensor_id"] = commit_request.get("sensor_id", "")
    metadata["sensor_pending_commit_approval_required"] = bool(approval.get("required", False))
    metadata["sensor_pending_commit_approved"] = bool(approval.get("approved", False))
    write_yaml(index_path, index)


def main() -> int:
    args = parse_args()
    dataset_index_path = Path(args.dataset_index).expanduser().resolve(strict=False)
    estimated_params_path = Path(args.estimated_params).expanduser().resolve(strict=False)
    index = load_yaml(dataset_index_path)
    params = load_yaml(estimated_params_path)
    session = validate_dataset_index(index)
    output_path = resolve_output_path(args, session)
    commit_package = build_commit_package(
        dataset_index_path,
        estimated_params_path,
        index,
        params,
        args,
        output_path,
    )
    write_yaml(output_path, commit_package)
    if args.update_dataset_index:
        update_dataset_index(dataset_index_path, index, output_path)
    print(f"[OK] 已生成传感器待提交参数包: {output_path}")
    print(f"[OK] commit_state={commit_package['commit_state']}")
    print(f"[OK] parameter_version={commit_package['commit_request']['parameter_version']}")
    return 0


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