# -*- coding: utf-8 -*-
"""
PB 与 dict 的转换层：严格按 vehicle_config.proto 定义做序列化/反序列化。
- PB 文本 <-> VehicleConfig (Message)
- VehicleConfig <-> dict（与现有 YAML 结构一致）
"""

from typing import Dict, List, Any

from google.protobuf import text_format

from .protobuf.vehicle_config_pb2 import (
    VehicleConfig,
    VehicleInfo,
    VehicleParam,
    LatencyParam,
    Extrinsics,
    Transform,
    Point3D,
    Quaternion,
    Intrinsics,
    CameraParam,
    PinholeCamera,
    FisheyeCamera,
    CameraModelType,
)

# vehicle_config.proto VehicleParam 标量字段（不含嵌套 LatencyParam）
VEHICLE_PARAM_SCALAR_FIELDS = [
    "top_edge_to_center", "bottom_edge_to_center", "front_edge_to_center", "back_edge_to_center",
    "left_edge_to_center", "right_edge_to_center", "length", "width", "height", "max_speed",
    "min_turn_radius", "max_acceleration", "max_deceleration", "max_engine_pedal", "max_brake_pedal",
    "max_steer_angle", "max_steer_angle_rate", "min_steer_angle_rate", "steer_ratio", "wheel_base",
    "wheel_rolling_radius", "drive_ratio", "total_vehicle_mass", "max_abs_speed_when_stopped",
    "min_abs_speed_when_start", "brake_deadzone", "throttle_deadzone",
    "steer_zero_offset_angle", "steer_deadband", "steer_lqr_deadband",
    "steer_ratio_a0", "steer_ratio_a1", "steer_ratio_a2", "steer_ratio_a3",
    "steer_ratio_a4", "steer_ratio_a5", "steer_ratio_a6", "steer_ratio_a7",
]

LATENCY_PARAM_FIELDS = ["dead_time", "rise_time", "peak_time", "settling_time"]


def parse_pb_text(content: str) -> VehicleConfig:
    """将 PB 文本格式反序列化为 VehicleConfig。"""
    msg = VehicleConfig()
    text_format.Parse(content, msg)
    return msg


def message_to_dict(msg: VehicleConfig) -> Dict[str, Any]:
    """将 VehicleConfig 转为与 _parser 一致的 dict 结构，供写 YAML 使用。"""
    data = {
        "vehicle_info": _vehicle_info_to_dict(msg.vehicle_info),
        "vehicle_param": _vehicle_param_to_dict(msg.vehicle_param),
        "extrinsics": _extrinsics_to_list(msg.extrinsics),
        "intrinsics": _intrinsics_to_list(msg.intrinsics),
    }
    return data


def _vehicle_info_to_dict(vi: VehicleInfo) -> Dict[str, str]:
    return {
        "model": vi.model or "",
        "id": vi.id or "",
        "vin": vi.vin or "",
        "plate": vi.plate or "",
        "other_unique_id": vi.other_unique_id or "",
    }


def _latency_param_to_dict(lp: LatencyParam) -> Dict[str, float]:
    return {f: float(getattr(lp, f, 0) or 0) for f in LATENCY_PARAM_FIELDS}


def _vehicle_param_to_dict(vp: VehicleParam) -> Dict[str, Any]:
    out = {}
    for key in VEHICLE_PARAM_SCALAR_FIELDS:
        out[key] = float(getattr(vp, key, 0) or 0)
    for key in ("steering_latency_param", "throttle_latency_param", "brake_latency_param"):
        sub = getattr(vp, key, None)
        if sub is not None:
            out[key] = _latency_param_to_dict(sub)
        else:
            out[key] = {f: 0.0 for f in LATENCY_PARAM_FIELDS}
    return out


def _extrinsics_to_list(ext: Extrinsics) -> List[Dict[str, Any]]:
    # proto 中字段名为 tansforms（拼写保留）
    result = []
    for t in ext.tansforms:
        trans = t.translation
        rot = t.rotation
        result.append({
            "source_frame": t.source_frame or "",
            "target_frame": t.target_frame or "",
            "translation": {"x": trans.x, "y": trans.y, "z": trans.z},
            "rotation": {"qx": rot.qx, "qy": rot.qy, "qz": rot.qz, "qw": rot.qw},
        })
    return result


def _camera_model_to_dict(width: int, height: int, intrinsic: List[float], distortion: List[float]) -> Dict[str, Any]:
    return {
        "width": width,
        "height": height,
        "intrinsic": list(intrinsic),
        "distortion": list(distortion),
    }


def _intrinsics_to_list(intr: Intrinsics) -> List[Dict[str, Any]]:
    result = []
    for cp in intr.camera_params:
        if getattr(cp, "model_type", CameraModelType.PINHOLE) == CameraModelType.FISHEYE:
            model_type_str = "FISHEYE"
        else:
            model_type_str = "FISHEYE" if cp.WhichOneof("model") == "fisheye" else "PINHOLE"
        item = {"frame_id": cp.frame_id or "", "model_type": model_type_str}
        if cp.HasField("pinhole"):
            ph = cp.pinhole
            item["pinhole"] = _camera_model_to_dict(ph.width, ph.height, list(ph.intrinsic), list(ph.distortion))
        elif cp.HasField("fisheye"):
            fe = cp.fisheye
            item["fisheye"] = _camera_model_to_dict(fe.width, fe.height, list(fe.intrinsic), list(fe.distortion))
        else:
            item["pinhole"] = {"width": 0, "height": 0, "intrinsic": [], "distortion": []}
        if model_type_str == "FISHEYE" and "fisheye" not in item and "pinhole" in item:
            item["fisheye"] = item["pinhole"]
        result.append(item)
    return result


def dict_to_message(data: Dict[str, Any]) -> VehicleConfig:
    """将现有 dict 结构（YAML 加载结果）填充为 VehicleConfig。"""
    msg = VehicleConfig()
    _fill_vehicle_info(msg.vehicle_info, data.get("vehicle_info") or {})
    _fill_vehicle_param(msg.vehicle_param, data.get("vehicle_param") or {})
    _fill_extrinsics(msg.extrinsics, data.get("extrinsics") or [])
    _fill_intrinsics(msg.intrinsics, data.get("intrinsics") or [])
    return msg


def _fill_vehicle_info(vi: VehicleInfo, d: Dict[str, str]) -> None:
    vi.model = d.get("model", "") or ""
    vi.id = d.get("id", "") or ""
    vi.vin = d.get("vin", "") or ""
    vi.plate = d.get("plate", "") or ""
    vi.other_unique_id = d.get("other_unique_id", "") or ""


def _fill_latency_param(lp: LatencyParam, d: Dict[str, Any]) -> None:
    for f in LATENCY_PARAM_FIELDS:
        setattr(lp, f, float(d.get(f, 0) or 0))


def _fill_vehicle_param(vp: VehicleParam, d: Dict[str, Any]) -> None:
    for key in VEHICLE_PARAM_SCALAR_FIELDS:
        val = d.get(key)
        if val is not None:
            setattr(vp, key, float(val))
    for key in ("steering_latency_param", "throttle_latency_param", "brake_latency_param"):
        sub = d.get(key)
        if isinstance(sub, dict):
            _fill_latency_param(getattr(vp, key), sub)


def _fill_extrinsics(ext: Extrinsics, items: List[Dict[str, Any]]) -> None:
    for t in items:
        tr = ext.tansforms.add()
        tr.source_frame = t.get("source_frame", "") or ""
        tr.target_frame = t.get("target_frame", "") or ""
        trans = t.get("translation") or {}
        tr.translation.x = float(trans.get("x", 0) or 0)
        tr.translation.y = float(trans.get("y", 0) or 0)
        tr.translation.z = float(trans.get("z", 0) or 0)
        rot = t.get("rotation") or {}
        tr.rotation.qx = float(rot.get("qx", rot.get("x", 0)) or 0)
        tr.rotation.qy = float(rot.get("qy", rot.get("y", 0)) or 0)
        tr.rotation.qz = float(rot.get("qz", rot.get("z", 0)) or 0)
        tr.rotation.qw = float(rot.get("qw", rot.get("w", 1)) or 1)


def _fill_intrinsics(intr: Intrinsics, items: List[Dict[str, Any]]) -> None:
    for it in items:
        cp = intr.camera_params.add()
        cp.frame_id = it.get("frame_id", "") or ""
        model_type = (it.get("model_type") or "PINHOLE").upper()
        if model_type == "FISHEYE":
            cp.model_type = CameraModelType.FISHEYE
        else:
            cp.model_type = CameraModelType.PINHOLE
        ph = it.get("pinhole")
        fe = it.get("fisheye")
        if model_type == "FISHEYE" and fe is not None:
            cp.fisheye.width = int(fe.get("width", 0) or 0)
            cp.fisheye.height = int(fe.get("height", 0) or 0)
            cp.fisheye.intrinsic.extend([float(x) for x in (fe.get("intrinsic") or [])])
            cp.fisheye.distortion.extend([float(x) for x in (fe.get("distortion") or [])])
        else:
            ph = ph or {}
            cp.pinhole.width = int(ph.get("width", 0) or 0)
            cp.pinhole.height = int(ph.get("height", 0) or 0)
            cp.pinhole.intrinsic.extend([float(x) for x in (ph.get("intrinsic") or [])])
            cp.pinhole.distortion.extend([float(x) for x in (ph.get("distortion") or [])])


def message_to_pb_text(msg: VehicleConfig) -> str:
    """将 VehicleConfig 序列化为 PB 文本格式。proto3 会省略默认枚举，导致 model_type: PINHOLE 不写出；此处为每个 camera_params 在 frame_id 后、pinhole/fisheye 前补全 model_type 行，顺序与已有 FISHEYE 块一致。"""
    import re
    s = text_format.MessageToString(msg, as_one_line=False)
    parts = s.split("  camera_params {")
    if len(parts) <= 1:
        return s
    result = [parts[0]]
    for part in parts[1:]:
        if "model_type:" in part:
            result.append("  camera_params {" + part)
            continue
        # 在 frame_id 行之后、pinhole/fisheye 块之前插入 model_type，保留缩进且不增加多余换行（\n+\s+ 拆成两组，插入时仅用 \n 一次）
        if "fisheye {" in part:
            part = re.sub(r"(\n)(\s+)(fisheye \{)", r"\1\2model_type: FISHEYE\n\2\3", part, count=1)
        elif "pinhole {" in part:
            part = re.sub(r"(\n)(\s+)(pinhole \{)", r"\1\2model_type: PINHOLE\n\2\3", part, count=1)
        else:
            part = re.sub(r"(\n)(\s+)(pinhole \{)", r"\1\2model_type: PINHOLE\n\2\3", part, count=1)
        result.append("  camera_params {" + part)
    return "".join(result)
