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

from typing import Dict, List, Any

from google.protobuf import text_format

from .protobuf.calibration_pb2 import (
    CalibInformation,
    CameraModel,
    CameraModelType,
    CameraParams,
    Extrinsics,
    Intrinsics,
    Rotation,
    Transforms,
    Translation,
    VehicleInfo,
    VehicleParam,
)

# vehicle_param 中 int32 字段（proto 中为 optional int32）
VEHICLE_PARAM_INT_FIELDS = frozenset({
    "brake_deadzone", "throttle_deadzone", "steer_zero_offset_angle",
    "steer_deadband", "steer_lqr_deadband",
})

VEHICLE_PARAM_FLOAT_FIELDS = [
    "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", "max_abs_speed_when_stopped",
    "min_abs_speed_when_start", "brake_deadzone", "throttle_deadzone", "steer_zero_offset_angle",
    "steer_deadband", "steer_lqr_deadband",
]


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


def message_to_dict(msg: CalibInformation) -> Dict[str, Any]:
    """将 CalibInformation 转为与 _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 _vehicle_param_to_dict(vp: VehicleParam) -> Dict[str, float]:
    out = {}
    for key in VEHICLE_PARAM_FLOAT_FIELDS:
        if key in VEHICLE_PARAM_INT_FIELDS:
            out[key] = float(getattr(vp, key, 0) or 0)
        else:
            out[key] = float(getattr(vp, key, 0) or 0)
    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 if trans.HasField("x") else 0.0,
                "y": trans.y if trans.HasField("y") else 0.0,
                "z": trans.z if trans.HasField("z") else 0.0,
            },
            "rotation": {
                "qx": rot.qx if rot.HasField("qx") else 0.0,
                "qy": rot.qy if rot.HasField("qy") else 0.0,
                "qz": rot.qz if rot.HasField("qz") else 0.0,
                "qw": rot.qw if rot.HasField("qw") else 1.0,
            },
        })
    return result


def _camera_model_to_dict(cm: CameraModel) -> Dict[str, Any]:
    return {
        "width": int(cm.width if cm.HasField("width") else 0),
        "height": int(cm.height if cm.HasField("height") else 0),
        "intrinsic": list(cm.intrinsic),
        "distortion": list(cm.distortion),
    }


def _intrinsics_to_list(intr: Intrinsics) -> List[Dict[str, Any]]:
    result = []
    for cp in intr.camera_params:
        # 优先使用显式的 model_type 枚举（PB 中可能写 model_type: FISHEYE 但数据在 pinhole 块里）
        if cp.HasField("model_type") and cp.model_type == 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"):
            item["pinhole"] = _camera_model_to_dict(cp.pinhole)
        elif cp.HasField("fisheye"):
            item["fisheye"] = _camera_model_to_dict(cp.fisheye)
        else:
            item["pinhole"] = {"width": 0, "height": 0, "intrinsic": [], "distortion": []}
        # 若 model_type 为 FISHEYE 但只有 pinhole 块（PB 格式不统一），补一份 fisheye 引用便于下游一致使用
        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]) -> CalibInformation:
    """将现有 dict 结构（YAML 加载结果）填充为 CalibInformation。"""
    msg = CalibInformation()
    _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_vehicle_param(vp: VehicleParam, d: Dict[str, Any]) -> None:
    for key in VEHICLE_PARAM_FLOAT_FIELDS:
        val = d.get(key)
        if val is None:
            continue
        if key in VEHICLE_PARAM_INT_FIELDS:
            setattr(vp, key, int(val))
        else:
            setattr(vp, key, float(val))


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", 0) or 0)
        tr.rotation.qy = float(rot.get("qy", 0) or 0)
        tr.rotation.qz = float(rot.get("qz", 0) or 0)
        tr.rotation.qw = float(rot.get("qw", 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 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:
            # 默认填 pinhole（与原有逻辑一致）
            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: CalibInformation) -> str:
    """将 CalibInformation 序列化为 PB 文本格式。"""
    return text_format.MessageToString(msg, as_one_line=False)
