# -*- coding: utf-8 -*-
"""
PB 文件解析器：按 vehicle_config.proto 标准反序列化 .pb.txt、对比模板缺项、生成顺序清单与分组 YAML。
"""

import copy
import re
import yaml
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple

from ._proto_io import parse_pb_text, message_to_dict


class _FlowList(list):
    """用于 YAML 输出时以方括号 [] 单行形式表示的列表。"""


def _flow_list_representer(dumper, data):
    return dumper.represent_sequence("tag:yaml.org,2002:seq", data, flow_style=True)


class _IntrinsicsYamlDumper(yaml.SafeDumper):
    pass


_IntrinsicsYamlDumper.add_representer(_FlowList, _flow_list_representer)


def _find_matching_brace(text: str, start_pos: int) -> int:
    depth = 0
    for i in range(start_pos, len(text)):
        if text[i] == "{":
            depth += 1
        elif text[i] == "}":
            depth -= 1
            if depth == 0:
                return i
    return -1


def _extract_float(content: str, field_name: str) -> float:
    pattern = rf"{re.escape(field_name)}:\s*([\d.eE+-]+)"
    match = re.search(pattern, content)
    if match:
        return float(match.group(1))
    return 0.0


def _default_proto_templates_dir() -> Path:
    """包内 proto_templates 目录（安装后位于 site-packages/pb_calibration/proto_templates）。"""
    return Path(__file__).resolve().parent / "proto_templates"


class PBParser:
    """解析 protobuf 文本格式的 vehicle_config.pb.txt，并输出顺序清单与分组 YAML。"""

    def __init__(self, proto_templates_dir: Optional[str] = None):
        self.proto_templates_dir = Path(proto_templates_dir or _default_proto_templates_dir())

    def parse_file(self, pb_path: str) -> Dict[str, Any]:
        """
        解析 PB 文件为结构化数据。按 vehicle_config.proto 标准反序列化后转为 dict。
        若解析失败则抛出异常。
        """
        with open(pb_path, "r", encoding="utf-8") as f:
            content = f.read()
        msg = parse_pb_text(content)
        return message_to_dict(msg)

    def validate_format(self, pb_path: str) -> Tuple[bool, Optional[str]]:
        """
        通过是否可反序列化判断 PB 文件格式是否正确（与 api.check 一致，均走 proto）。
        返回 (是否成功, 错误信息)。
        """
        try:
            self.parse_file(pb_path)
            return True, None
        except Exception as e:
            return False, str(e)

    def _parse_vehicle_info(self, content: str) -> Dict[str, str]:
        info = {}
        match = re.search(r"vehicle_info\s*\{([^}]*)\}", content, re.DOTALL)
        if not match:
            return info
        block = match.group(1)
        for key in ["model", "id", "vin", "plate", "other_unique_id"]:
            m = re.search(rf'{key}:\s*"([^"]*)"', block)
            if m:
                info[key] = m.group(1)
        return info

    def _parse_vehicle_param(self, content: str) -> Dict[str, float]:
        param = {}
        match = re.search(r"vehicle_param\s*\{([^}]*)\}", content, re.DOTALL)
        if not match:
            return param
        block = match.group(1)
        keys = [
            "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",
        ]
        for key in keys:
            m = re.search(rf"{re.escape(key)}:\s*([\d.eE+-]+)", block)
            if m:
                param[key] = float(m.group(1))
        return param

    def _parse_extrinsics(self, content: str) -> List[Dict[str, Any]]:
        transforms = []
        start = content.find("extrinsics")
        if start == -1:
            return transforms
        brace_start = content.find("{", start)
        if brace_start == -1:
            return transforms
        brace_end = _find_matching_brace(content, brace_start)
        if brace_end == -1:
            return transforms
        extrinsics_content = content[brace_start + 1 : brace_end]
        pos = 0
        while True:
            ts = extrinsics_content.find("tansforms", pos)
            if ts == -1:
                break
            ts_brace = extrinsics_content.find("{", ts)
            ts_end = _find_matching_brace(extrinsics_content, ts_brace)
            if ts_brace == -1 or ts_end == -1:
                break
            block = extrinsics_content[ts_brace + 1 : ts_end]
            t = self._parse_transform_block(block)
            if t:
                transforms.append(t)
            pos = ts_end + 1
        return transforms

    def _parse_transform_block(self, block: str) -> Optional[Dict[str, Any]]:
        sm = re.search(r'source_frame:\s*"([^"]*)"', block)
        tm = re.search(r'target_frame:\s*"([^"]*)"', block)
        if not sm or not tm:
            return None
        trans_m = re.search(r"translation\s*\{([^}]*)\}", block, re.DOTALL)
        rot_m = re.search(r"rotation\s*\{([^}]*)\}", block, re.DOTALL)
        if not trans_m or not rot_m:
            return None
        tx = _extract_float(trans_m.group(1), "x")
        ty = _extract_float(trans_m.group(1), "y")
        tz = _extract_float(trans_m.group(1), "z")
        qx = _extract_float(rot_m.group(1), "qx")
        qy = _extract_float(rot_m.group(1), "qy")
        qz = _extract_float(rot_m.group(1), "qz")
        qw = _extract_float(rot_m.group(1), "qw")
        return {
            "source_frame": sm.group(1),
            "target_frame": tm.group(1),
            "translation": {"x": tx, "y": ty, "z": tz},
            "rotation": {"qx": qx, "qy": qy, "qz": qz, "qw": qw},
        }

    def _parse_intrinsics(self, content: str) -> List[Dict[str, Any]]:
        items = []
        start = content.find("intrinsics")
        if start == -1:
            return items
        brace_start = content.find("{", start)
        if brace_start == -1:
            return items
        brace_end = _find_matching_brace(content, brace_start)
        if brace_end == -1:
            return items
        intrinsics_content = content[brace_start + 1 : brace_end]
        pos = 0
        while True:
            cam = intrinsics_content.find("camera_params", pos)
            if cam == -1:
                break
            block_start = intrinsics_content.find("{", cam)
            block_end = _find_matching_brace(intrinsics_content, block_start)
            if block_start == -1 or block_end == -1:
                break
            block = intrinsics_content[block_start + 1 : block_end]
            item = self._parse_intrinsic_block(block)
            if item:
                items.append(item)
            pos = block_end + 1
        return items

    def _parse_intrinsic_block(self, block: str) -> Optional[Dict[str, Any]]:
        fid = re.search(r'frame_id:\s*"([^"]*)"', block)
        if not fid:
            return None
        model_m = re.search(r"model_type:\s*(\w+)", block)
        model_type = model_m.group(1) if model_m else "PINHOLE"
        pinhole_m = re.search(r"pinhole\s*\{([^{}]*(?:\{[^{}]*\}[^{}]*)*)\}", block, re.DOTALL)
        if not pinhole_m:
            return None
        ph = pinhole_m.group(1)
        width = int(_extract_float(ph, "width"))
        height = int(_extract_float(ph, "height"))
        intrinsic_m = re.search(r"intrinsic:\s*\[(.*?)\]", ph, re.DOTALL)
        intrinsic = []
        if intrinsic_m:
            intrinsic = [float(x.strip()) for x in intrinsic_m.group(1).split(",") if x.strip()]
        distortion_m = re.search(r"distortion:\s*\[(.*?)\]", ph, re.DOTALL)
        distortion = []
        if distortion_m:
            distortion = [float(x.strip()) for x in distortion_m.group(1).split(",") if x.strip()]
        return {
            "frame_id": fid.group(1),
            "model_type": model_type,
            "pinhole": {
                "width": width,
                "height": height,
                "intrinsic": intrinsic,
                "distortion": distortion,
            },
        }

    def compare_with_templates(self, data: Dict[str, Any]) -> Dict[str, List[Any]]:
        """对比反序列化结果与 proto_templates，返回缺项/漏项。"""
        report = {
            "vehicle_info": [],
            "vehicle_param": [],
            "extrinsics": [],
            "intrinsics": [],
        }

        def load_template(name: str) -> Dict:
            p = self.proto_templates_dir / f"{name}.yaml"
            if not p.exists():
                return {}
            with open(p, "r", encoding="utf-8") as f:
                return yaml.safe_load(f) or {}

        t_vi = load_template("vehicle_info")
        for field in (t_vi.get("required_fields") or []):
            if field not in data.get("vehicle_info", {}):
                report["vehicle_info"].append(field)

        t_vp = load_template("vehicle_param")
        for field in (t_vp.get("required_fields") or []):
            if field not in data.get("vehicle_param", {}):
                report["vehicle_param"].append(field)

        t_ext = load_template("extrinsic")
        req = set(t_ext.get("required_fields") or [])
        trans_f = set(t_ext.get("translation_fields") or [])
        rot_f = set(t_ext.get("rotation_fields") or [])
        for i, ext in enumerate(data.get("extrinsics") or []):
            missing = []
            for f in req:
                if f not in ext:
                    missing.append(f)
            if "translation" in ext:
                for f in trans_f:
                    if f not in ext["translation"]:
                        missing.append(f"translation.{f}")
            if "rotation" in ext:
                for f in rot_f:
                    if f not in ext["rotation"]:
                        missing.append(f"rotation.{f}")
            if missing:
                report["extrinsics"].append({"index": i, "missing": missing})

        t_int = load_template("intrinsic")
        req_int = set(t_int.get("required_fields") or [])
        ph_f = set(t_int.get("pinhole_fields") or [])
        for i, int_item in enumerate(data.get("intrinsics") or []):
            missing = []
            for f in req_int:
                if f not in int_item:
                    missing.append(f)
            if "pinhole" in int_item:
                for f in ph_f:
                    if f not in int_item["pinhole"]:
                        missing.append(f"pinhole.{f}")
            if missing:
                report["intrinsics"].append({"index": i, "frame_id": int_item.get("frame_id", ""), "missing": missing})

        return report

    def generate_order_manifest(self, data: Dict[str, Any], output_dir: str) -> Dict[str, Any]:
        """生成顺序清单（YAML 存储，用于按顺序重组）。"""
        manifest = {
            "root_sections": ["vehicle_info", "vehicle_param", "extrinsics", "intrinsics"],
            "vehicle_info_file": "vehicle_info.yaml",
            "vehicle_param_file": "vehicle_param.yaml",
            "extrinsics": [],
            "intrinsics": [],
        }
        for i, _ in enumerate(data.get("extrinsics") or []):
            manifest["extrinsics"].append({"index": i, "file": f"extrinsics/{i:03d}.yaml"})
        for item in (data.get("intrinsics") or []):
            fid = item.get("frame_id", "unknown")
            safe_name = fid.replace("/", "_")
            manifest["intrinsics"].append({"frame_id": fid, "file": f"intrinsics/{safe_name}.yaml"})
        return manifest

    def dump_to_yaml_dir(
        self,
        data: Dict[str, Any],
        output_dir: str,
        order_manifest: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """将解析结果写出：order_manifest.yaml、vehicle_info/vehicle_param、extrinsics/*、intrinsics/*。"""
        out = Path(output_dir)
        out.mkdir(parents=True, exist_ok=True)
        (out / "extrinsics").mkdir(exist_ok=True)
        (out / "intrinsics").mkdir(exist_ok=True)

        if order_manifest is None:
            order_manifest = self.generate_order_manifest(data, str(out))

        def write_yaml(path: Path, obj: Any):
            with open(path, "w", encoding="utf-8") as f:
                yaml.dump(obj, f, allow_unicode=True, default_flow_style=False, sort_keys=False)

        write_yaml(out / "order_manifest.yaml", order_manifest)
        if data.get("vehicle_info"):
            write_yaml(out / "vehicle_info.yaml", data["vehicle_info"])
        if data.get("vehicle_param"):
            write_yaml(out / "vehicle_param.yaml", data["vehicle_param"])
        for i, ext in enumerate(data.get("extrinsics") or []):
            write_yaml(out / "extrinsics" / f"{i:03d}.yaml", ext)
        for item in (data.get("intrinsics") or []):
            fid = item.get("frame_id", "unknown")
            safe_name = fid.replace("/", "_")
            out_item = copy.deepcopy(item)
            for model_key in ("pinhole", "fisheye"):
                ph = out_item.get(model_key) or {}
                if "intrinsic" in ph and isinstance(ph["intrinsic"], list):
                    ph["intrinsic"] = _FlowList(ph["intrinsic"])
                if "distortion" in ph and isinstance(ph["distortion"], list):
                    ph["distortion"] = _FlowList(ph["distortion"])
            with open(out / "intrinsics" / f"{safe_name}.yaml", "w", encoding="utf-8") as f:
                yaml.dump(out_item, f, allow_unicode=True, default_flow_style=False, sort_keys=False, Dumper=_IntrinsicsYamlDumper)

        return order_manifest

    def parse_and_dump(self, pb_path: str, output_dir: str) -> Tuple[Dict[str, Any], Dict[str, List[Any]]]:
        """解析 PB 文件，生成顺序清单并写出所有分组 YAML；同时返回缺项报告。"""
        data = self.parse_file(pb_path)
        missing_report = self.compare_with_templates(data)
        manifest = self.generate_order_manifest(data, output_dir)
        self.dump_to_yaml_dir(data, output_dir, manifest)
        return data, missing_report
