#!/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"}
CONTROL_AXES = {"lateral_control", "longitudinal_control", "combined"}
CONTROLLER_ALGORITHMS = {"pid", "mpc", "lqr", "pure_pursuit"}
SINGLE_PID_CONTROL_AXES = {"lateral_control", "longitudinal_control"}
CONTROL_ROLES = {
    "path_tracking_outer_loop",
    "speed_loop",
    "acceleration_loop",
    "heading_loop",
    "yaw_rate_loop",
    "steering_angle_inner_loop",
    "module_steering_inner_loop",
    "wheel_speed_inner_loop",
    "multi_loop_pid",
}
CONTROL_ROLE_AXES = {
    "path_tracking_outer_loop": {"lateral_control", "combined"},
    "speed_loop": {"longitudinal_control"},
    "acceleration_loop": {"longitudinal_control"},
    "heading_loop": {"lateral_control"},
    "yaw_rate_loop": {"lateral_control"},
    "steering_angle_inner_loop": {"lateral_control"},
    "module_steering_inner_loop": {"lateral_control"},
    "wheel_speed_inner_loop": {"longitudinal_control"},
    "multi_loop_pid": {"combined"},
}
CONTROL_ROLE_ALGORITHMS = {
    "path_tracking_outer_loop": {"mpc", "lqr", "pure_pursuit", "pid"},
    "speed_loop": {"pid", "mpc"},
    "acceleration_loop": {"pid", "mpc", "lqr"},
    "heading_loop": {"pid", "lqr"},
    "yaw_rate_loop": {"pid"},
    "steering_angle_inner_loop": {"pid"},
    "module_steering_inner_loop": {"pid"},
    "wheel_speed_inner_loop": {"pid"},
    "multi_loop_pid": {"pid"},
}
CHASSIS_ROLE_COMPATIBILITY = {
    "ackermann": {
        "path_tracking_outer_loop",
        "speed_loop",
        "acceleration_loop",
        "heading_loop",
        "yaw_rate_loop",
        "steering_angle_inner_loop",
        "wheel_speed_inner_loop",
        "multi_loop_pid",
    },
    "differential": {
        "path_tracking_outer_loop",
        "speed_loop",
        "acceleration_loop",
        "heading_loop",
        "yaw_rate_loop",
        "wheel_speed_inner_loop",
        "multi_loop_pid",
    },
    "single_steer_wheel": {
        "path_tracking_outer_loop",
        "speed_loop",
        "acceleration_loop",
        "heading_loop",
        "yaw_rate_loop",
        "steering_angle_inner_loop",
        "wheel_speed_inner_loop",
        "multi_loop_pid",
    },
    "multi_steer_wheel": {
        "path_tracking_outer_loop",
        "speed_loop",
        "acceleration_loop",
        "heading_loop",
        "yaw_rate_loop",
        "module_steering_inner_loop",
        "wheel_speed_inner_loop",
        "multi_loop_pid",
    },
}

PID_FIELDS = ["kp", "ki", "kd"]
LATERAL_MPC_FIELDS = [
    "lateral_mpc.prediction_horizon",
    "lateral_mpc.control_horizon",
    "lateral_mpc.model_dt_s",
    "lateral_mpc.q_lateral",
    "lateral_mpc.q_heading",
    "lateral_mpc.r_steering",
    "lateral_mpc.r_steering_rate",
]
LONGITUDINAL_MPC_FIELDS = [
    "longitudinal_mpc.prediction_horizon",
    "longitudinal_mpc.control_horizon",
    "longitudinal_mpc.model_dt_s",
    "longitudinal_mpc.q_speed",
    "longitudinal_mpc.q_accel",
    "longitudinal_mpc.r_throttle",
    "longitudinal_mpc.r_brake",
    "longitudinal_mpc.r_jerk",
]
LQR_LIST_FIELDS = ["lqr.q_state_weights", "lqr.r_input_weights"]
PURE_PURSUIT_FIELDS = ["pure_pursuit.lookahead_m"]


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="输出待提交参数包；为空时写到会话目录 control/pending_control_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 require_integer(value: Any, field_name: str) -> int:
    if not has_value(value):
        raise ValueError(f"{field_name} 不能为空。")
    try:
        integer = int(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{field_name} 必须是整数，当前值为 {value!r}。") from exc
    if integer <= 0:
        raise ValueError(f"{field_name} 必须大于 0。")
    return integer


def require_number_list(value: Any, field_name: str) -> None:
    if not isinstance(value, list) or not value:
        raise ValueError(f"{field_name} 必须是非空数字数组。")
    for index, item in enumerate(value):
        require_number(item, f"{field_name}[{index}]")


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


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


def validate_pid_loop(
    loop: dict[str, Any],
    field_prefix: str,
    expected_control_axis: str | None = None,
    expected_control_role: str | None = None,
    chassis_type: str | None = None,
) -> None:
    loop_name = require_string(loop.get("loop_name"), f"{field_prefix}.loop_name")
    control_role = require_string(loop.get("control_role"), f"{field_prefix}.control_role")
    if control_role not in CONTROL_ROLES:
        raise ValueError(f"{field_prefix}.control_role 必须是 {sorted(CONTROL_ROLES)} 之一。")
    if control_role == "multi_loop_pid":
        raise ValueError(f"{field_prefix}.control_role 不能写 multi_loop_pid，必须描述具体 PID 回路。")
    if expected_control_role is not None and control_role != expected_control_role:
        raise ValueError(f"{field_prefix}.control_role 必须与顶层 control_role 一致。")
    if expected_control_axis is not None:
        loop_axis = str(loop.get("control_axis", expected_control_axis) or "").strip()
        if loop_axis != expected_control_axis:
            raise ValueError(f"{field_prefix}.control_axis 必须与顶层 control_axis 一致。")
    elif str(loop.get("control_axis", "") or "").strip() not in SINGLE_PID_CONTROL_AXES:
        raise ValueError(f"{field_prefix}.control_axis 必须是 {sorted(SINGLE_PID_CONTROL_AXES)} 之一。")
    loop_axis = str(loop.get("control_axis", expected_control_axis) or "").strip()
    validate_control_role("pid", loop_axis, control_role, field_prefix, chassis_type)
    if loop_name == "combined":
        raise ValueError(f"{field_prefix}.loop_name 不能写 combined，必须明确具体 PID 回路。")
    for field_name in PID_FIELDS:
        require_number(loop.get(field_name), f"{field_prefix}.{field_name}")


def validate_pid_params(
    estimated: dict[str, Any],
    control_axis: str,
    control_role: str,
    chassis_type: str,
) -> None:
    if control_axis in SINGLE_PID_CONTROL_AXES:
        validate_pid_loop(
            require_map(estimated.get("pid"), "estimated_params.estimated_params.pid"),
            "estimated_params.estimated_params.pid",
            control_axis,
            control_role,
            chassis_type,
        )
        return

    pid_loops = estimated.get("pid_loops")
    if not isinstance(pid_loops, list) or not pid_loops:
        raise ValueError(
            "estimated_params.control_axis=combined 且 controller_algorithm=pid 时，"
            "必须使用 estimated_params.pid_loops 列表分别描述每个 PID 回路。"
        )
    for index, loop in enumerate(pid_loops):
        validate_pid_loop(
            require_map(loop, f"estimated_params.estimated_params.pid_loops[{index}]"),
            f"estimated_params.estimated_params.pid_loops[{index}]",
            None,
            None,
            chassis_type,
        )


def validate_control_role(
    controller_algorithm: str,
    control_axis: str,
    control_role: str,
    field_prefix: str,
    chassis_type: str | None = None,
) -> None:
    if control_role not in CONTROL_ROLES:
        raise ValueError(f"{field_prefix}.control_role 必须是 {sorted(CONTROL_ROLES)} 之一。")
    allowed_axes = CONTROL_ROLE_AXES[control_role]
    if control_axis not in allowed_axes:
        raise ValueError(
            f"{field_prefix}.control_role={control_role} 不适用于 control_axis={control_axis}。"
        )
    allowed_algorithms = CONTROL_ROLE_ALGORITHMS[control_role]
    if controller_algorithm not in allowed_algorithms:
        raise ValueError(
            f"{field_prefix}.control_role={control_role} 不适用于 controller_algorithm={controller_algorithm}。"
        )
    if chassis_type:
        allowed_roles = CHASSIS_ROLE_COMPATIBILITY[chassis_type]
        if control_role not in allowed_roles:
            raise ValueError(
                f"{field_prefix}.control_role={control_role} 不适用于 chassis_type={chassis_type}。"
            )


def validate_mpc_params(estimated: dict[str, Any], control_axis: str) -> None:
    require_lateral = control_axis in {"lateral_control", "combined"}
    require_longitudinal = control_axis in {"longitudinal_control", "combined"}
    if require_lateral:
        for field_name in LATERAL_MPC_FIELDS:
            value = dotted_get(estimated, field_name)
            if field_name.endswith("prediction_horizon") or field_name.endswith("control_horizon"):
                require_integer(value, f"estimated_params.estimated_params.{field_name}")
            else:
                require_number(value, f"estimated_params.estimated_params.{field_name}")
    if require_longitudinal:
        for field_name in LONGITUDINAL_MPC_FIELDS:
            value = dotted_get(estimated, field_name)
            if field_name.endswith("prediction_horizon") or field_name.endswith("control_horizon"):
                require_integer(value, f"estimated_params.estimated_params.{field_name}")
            else:
                require_number(value, f"estimated_params.estimated_params.{field_name}")


def validate_lqr_params(estimated: dict[str, Any]) -> None:
    for field_name in LQR_LIST_FIELDS:
        require_number_list(
            dotted_get(estimated, field_name),
            f"estimated_params.estimated_params.{field_name}",
        )


def validate_pure_pursuit_params(estimated: dict[str, Any]) -> None:
    for field_name in PURE_PURSUIT_FIELDS:
        require_number(dotted_get(estimated, field_name), f"estimated_params.estimated_params.{field_name}")


def validate_estimated_params(
    params: dict[str, Any],
    session: dict[str, Any],
    index: dict[str, Any],
) -> tuple[str, 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")
    chassis_type = require_string(params.get("chassis_type"), "estimated_params.chassis_type")
    control_axis = require_string(params.get("control_axis"), "estimated_params.control_axis")
    controller_algorithm = require_string(params.get("controller_algorithm"), "estimated_params.controller_algorithm")
    control_role = require_string(params.get("control_role"), "estimated_params.control_role")
    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 中的底盘类型不一致。")
    if control_axis not in CONTROL_AXES:
        raise ValueError(f"estimated_params.control_axis 必须是 {sorted(CONTROL_AXES)} 之一。")
    if controller_algorithm not in CONTROLLER_ALGORITHMS:
        raise ValueError(f"estimated_params.controller_algorithm 必须是 {sorted(CONTROLLER_ALGORITHMS)} 之一。")
    validate_control_role(controller_algorithm, control_axis, control_role, "estimated_params", chassis_type)
    if controller_algorithm == "pid" and control_axis == "combined" and control_role != "multi_loop_pid":
        raise ValueError("combined + pid 必须使用 control_role=multi_loop_pid，并在 pid_loops 中列出各 PID 回路。")
    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")

    if controller_algorithm == "pid":
        validate_pid_params(estimated, control_axis, control_role, chassis_type)
    elif controller_algorithm == "mpc":
        validate_mpc_params(estimated, control_axis)
    elif controller_algorithm == "lqr":
        validate_lqr_params(estimated)
    elif controller_algorithm == "pure_pursuit":
        validate_pure_pursuit_params(estimated)

    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, control_axis, controller_algorithm, control_role, estimated


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 / "control" / "pending_control_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, control_axis, controller_algorithm, control_role, 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,
            "control_axis": control_axis,
            "controller_algorithm": controller_algorithm,
            "control_role": control_role,
            "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_control_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["control_estimated_params_file"] = source.get("estimated_params_file", "")
    metadata["control_estimated_params_checksum_type"] = digest.get("checksum_type", "")
    metadata["control_estimated_params_checksum_value"] = digest.get("checksum_value", "")
    metadata["control_pending_commit_file"] = relative_to_root(output_path, dataset_root)
    metadata["control_pending_commit_state"] = commit_package.get("commit_state", "")
    metadata["control_pending_parameter_version"] = commit_request.get("parameter_version", "")
    metadata["control_pending_commit_chassis_type"] = commit_request.get("chassis_type", "")
    metadata["control_pending_commit_control_axis"] = commit_request.get("control_axis", "")
    metadata["control_pending_commit_controller_algorithm"] = commit_request.get("controller_algorithm", "")
    metadata["control_pending_commit_control_role"] = commit_request.get("control_role", "")
    metadata["control_pending_commit_approval_required"] = bool(approval.get("required", False))
    metadata["control_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)
