#!/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


CHASSIS_TYPES = {"ackermann", "differential", "single_steer_wheel", "multi_steer_wheel"}

REQUIRED_PARAM_FIELDS = {
    "ackermann": [
        "ackermann.front_left_steer_zero_offset_deg",
        "ackermann.front_right_steer_zero_offset_deg",
        "ackermann.rear_left_wheel_radius_m",
        "ackermann.rear_right_wheel_radius_m",
        "ackermann.steering_ratio",
    ],
    "differential": [
        "differential.left_wheel_radius_m",
        "differential.right_wheel_radius_m",
        "differential.axle_track_width_m",
    ],
    "single_steer_wheel": [
        "single_steer.drive_wheel_radius_m",
        "single_steer.steer_zero_offset_deg",
        "single_steer.steering_ratio",
        "single_steer.drive_encoder_scale",
    ],
    "multi_steer_wheel": [
        "multi_steer.modules",
    ],
}


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="输出待提交参数包；为空时写到会话目录 chassis/pending_chassis_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 dotted_get(data: dict[str, Any], dotted_key: str) -> Any:
    current: Any = data
    for part in dotted_key.split("."):
        if not isinstance(current, dict) or part not in current:
            return None
        current = current[part]
    return current


def has_value(value: Any) -> bool:
    if value in (None, ""):
        return False
    if isinstance(value, list):
        return bool(value)
    return True


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


def dataset_chassis_type(index: dict[str, Any]) -> str:
    metadata = index.get("metadata", {})
    if isinstance(metadata, dict):
        summary = metadata.get("chassis_replay_summary", {})
        if isinstance(summary, dict):
            value = str(summary.get("chassis_type", "") or "").strip()
            if value:
                return value
    return ""


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}")
    chassis_inputs = require_map(
        require_map(index.get("data_inputs"), "dataset_index.data_inputs").get("chassis"),
        "dataset_index.data_inputs.chassis",
    )
    for field_name in ("chassis_motion_data_files", "actuator_command_files", "truth_trajectory_files"):
        values = chassis_inputs.get(field_name)
        if not isinstance(values, list) or not values:
            raise ValueError(f"dataset_index.data_inputs.chassis.{field_name} 不能为空。")
    return session


def validate_estimated_params(
    params: dict[str, Any],
    session: dict[str, Any],
    index: dict[str, Any],
) -> tuple[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")
    chassis_type = require_string(params.get("chassis_type"), "estimated_params.chassis_type")
    if chassis_type not in CHASSIS_TYPES:
        raise ValueError(f"estimated_params.chassis_type 必须是 {sorted(CHASSIS_TYPES)} 之一。")
    index_chassis_type = dataset_chassis_type(index)
    if index_chassis_type and index_chassis_type != chassis_type:
        raise ValueError("estimated_params.chassis_type 与 dataset_index 中的底盘类型不一致。")
    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")
    for field_name in REQUIRED_PARAM_FIELDS[chassis_type]:
        value = dotted_get(estimated, field_name)
        if field_name == "multi_steer.modules":
            if not isinstance(value, list) or not value:
                raise ValueError("estimated_params.estimated_params.multi_steer.modules 不能为空。")
            continue
        require_number(value, f"estimated_params.estimated_params.{field_name}")
    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。")
    return parameter_version, chassis_type, estimated


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 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 / "chassis" / "pending_chassis_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, chassis_type, estimated = validate_estimated_params(params, session, index)
    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,
            "chassis_type": chassis_type,
            "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_chassis_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["chassis_estimated_params_file"] = source.get("estimated_params_file", "")
    metadata["chassis_estimated_params_checksum_type"] = digest.get("checksum_type", "")
    metadata["chassis_estimated_params_checksum_value"] = digest.get("checksum_value", "")
    metadata["chassis_pending_commit_file"] = relative_to_root(output_path, dataset_root)
    metadata["chassis_pending_commit_state"] = commit_package.get("commit_state", "")
    metadata["chassis_pending_parameter_version"] = commit_request.get("parameter_version", "")
    metadata["chassis_pending_commit_chassis_type"] = commit_request.get("chassis_type", "")
    metadata["chassis_pending_commit_approval_required"] = bool(approval.get("required", False))
    metadata["chassis_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)
