# -*- coding: utf-8 -*-
"""
往返测试：读真实 pb.txt → 解析到 temp YAML → 重组 pb.txt → 与原文做语义对比。
修改参数后往返：修改 YAML 再 build，断言修改在重组 pb 中一致且与原文不等。

运行方式（在 V0.0.7 目录下，先 conda activate py310 && pip install -e .）：
  pytest test/test_roundtrip.py -v
  或
  python test/test_roundtrip.py
"""

import tempfile
import yaml
from pathlib import Path

import pytest

from pb_calibration import parse, build
from pb_calibration._proto_io import parse_pb_text, message_to_dict

TOL = 1e-9


def _deep_eq(a, b, path="", tol=1e-9):
    """递归比较两个结构（dict/list/标量），浮点数用 tol 容差。返回 (是否相等, 差异描述)。"""
    if type(a) != type(b):
        return False, f"{path}: 类型不同 {type(a).__name__} vs {type(b).__name__}"
    if isinstance(a, dict):
        keys_a, keys_b = set(a), set(b)
        if keys_a != keys_b:
            return False, f"{path}: 键不同, 仅左有 {keys_a - keys_b}, 仅右有 {keys_b - keys_a}"
        for k in keys_a:
            ok, msg = _deep_eq(a[k], b[k], path=f"{path}.{k}", tol=tol)
            if not ok:
                return False, msg
        return True, ""
    if isinstance(a, list):
        if len(a) != len(b):
            return False, f"{path}: 列表长度不同 {len(a)} vs {len(b)}"
        for i, (xa, xb) in enumerate(zip(a, b)):
            ok, msg = _deep_eq(xa, xb, path=f"{path}[{i}]", tol=tol)
            if not ok:
                return False, msg
        return True, ""
    if isinstance(a, float):
        if abs(a - b) > tol:
            return False, f"{path}: 浮点不同 {a!r} vs {b!r}"
        return True, ""
    if a != b:
        return False, f"{path}: 值不同 {a!r} vs {b!r}"
    return True, ""


def test_roundtrip_pb_to_yaml_to_pb():
    """原始 pb.txt → parse → temp YAML → build → 重组 pb.txt，与原文语义一致。"""
    test_dir = Path(__file__).resolve().parent
    pb_path = test_dir / "vehicle_config.pb.txt"
    if not pb_path.exists():
        pytest.skip(f"测试输入不存在: {pb_path}")

    with tempfile.TemporaryDirectory(prefix="pb_cal_roundtrip_") as tmpdir:
        yaml_dir = tmpdir
        rebuilt_pb_path = Path(tmpdir) / "rebuilt.pb.txt"

        parse(str(pb_path), yaml_dir)

        # S1: order_manifest 含预期 key
        with open(Path(yaml_dir) / "order_manifest.yaml", "r", encoding="utf-8") as f:
            manifest = yaml.safe_load(f) or {}
        for key in ("vehicle_info_file", "vehicle_param_file", "extrinsics", "intrinsics"):
            assert key in manifest, f"order_manifest 缺少 key: {key}"

        build(yaml_dir, str(rebuilt_pb_path))

        with open(pb_path, "r", encoding="utf-8") as f:
            orig_content = f.read()
        with open(rebuilt_pb_path, "r", encoding="utf-8") as f:
            rebuilt_content = f.read()

        orig_msg = parse_pb_text(orig_content)
        rebuilt_msg = parse_pb_text(rebuilt_content)

        # 语义一致：二进制序列化应相同
        orig_bin = orig_msg.SerializeToString()
        rebuilt_bin = rebuilt_msg.SerializeToString()
        assert orig_bin == rebuilt_bin, (
            "往返后 pb 与原文语义不一致（SerializeToString 不同）。"
            " 可检查 parse→yaml→build 过程中是否有字段丢失或精度变化。"
        )

        # 可选：dict 深度比较，便于失败时定位差异（浮点容差）
        orig_dict = message_to_dict(orig_msg)
        rebuilt_dict = message_to_dict(rebuilt_msg)
        ok, diff_msg = _deep_eq(orig_dict, rebuilt_dict, path="root", tol=TOL)
        assert ok, f"往返后 dict 与原文不一致: {diff_msg}"

        # S2: 外参/内参数量一致
        assert len(orig_dict.get("extrinsics") or []) == len(rebuilt_dict.get("extrinsics") or [])
        assert len(orig_dict.get("intrinsics") or []) == len(rebuilt_dict.get("intrinsics") or [])


def test_roundtrip_modified_yaml_reflected_in_rebuilt():
    """parse → 修改 YAML 若干字段 → build → 解析重组 pb，断言修改后的值一致且与原文不等。"""
    test_dir = Path(__file__).resolve().parent
    pb_path = test_dir / "vehicle_config.pb.txt"
    if not pb_path.exists():
        pytest.skip(f"测试输入不存在: {pb_path}")

    # 修改后的期望值（浮点用容差 1e-9 比较）
    new_model = "TEST_MODEL"
    new_vin = "TEST_VIN"
    new_length = 999.0
    new_max_speed = 88.8
    new_translation_x = 1.234
    new_rotation_qw = 0.5
    new_pinhole_width = 1920
    new_intrinsic_0 = 1000.0

    with tempfile.TemporaryDirectory(prefix="pb_cal_roundtrip_") as tmpdir:
        yaml_dir = Path(tmpdir)
        rebuilt_pb_path = yaml_dir / "rebuilt.pb.txt"

        parse(str(pb_path), str(yaml_dir))

        # 修改 vehicle_info.yaml
        vi_path = yaml_dir / "vehicle_info.yaml"
        with open(vi_path, "r", encoding="utf-8") as f:
            vi = yaml.safe_load(f) or {}
        vi["model"] = new_model
        vi["vin"] = new_vin
        with open(vi_path, "w", encoding="utf-8") as f:
            yaml.dump(vi, f, allow_unicode=True, default_flow_style=False, sort_keys=False)

        # 修改 vehicle_param.yaml
        vp_path = yaml_dir / "vehicle_param.yaml"
        with open(vp_path, "r", encoding="utf-8") as f:
            vp = yaml.safe_load(f) or {}
        vp["length"] = new_length
        vp["max_speed"] = new_max_speed
        with open(vp_path, "w", encoding="utf-8") as f:
            yaml.dump(vp, f, allow_unicode=True, default_flow_style=False, sort_keys=False)

        # 修改 extrinsics/000.yaml
        ext0_path = yaml_dir / "extrinsics" / "000.yaml"
        if ext0_path.exists():
            with open(ext0_path, "r", encoding="utf-8") as f:
                ext0 = yaml.safe_load(f) or {}
            if "translation" not in ext0:
                ext0["translation"] = {}
            ext0["translation"]["x"] = new_translation_x
            ext0["rotation"] = ext0.get("rotation") or {}
            ext0["rotation"]["qw"] = new_rotation_qw
            with open(ext0_path, "w", encoding="utf-8") as f:
                yaml.dump(ext0, f, allow_unicode=True, default_flow_style=False, sort_keys=False)

        # 修改第一个 intrinsics 文件（从 order_manifest 取）
        with open(yaml_dir / "order_manifest.yaml", "r", encoding="utf-8") as f:
            manifest = yaml.safe_load(f) or {}
        int_files = manifest.get("intrinsics") or []
        if int_files:
            first_int_file = int_files[0].get("file")
            if first_int_file:
                int_path = yaml_dir / first_int_file
                if int_path.exists():
                    with open(int_path, "r", encoding="utf-8") as f:
                        int0 = yaml.safe_load(f) or {}
                    ph = int0.get("pinhole") or int0.get("fisheye") or {}
                    ph = dict(ph)
                    ph["width"] = new_pinhole_width
                    if "intrinsic" in ph and isinstance(ph["intrinsic"], list) and len(ph["intrinsic"]) > 0:
                        ph["intrinsic"] = [new_intrinsic_0] + list(ph["intrinsic"][1:])
                    else:
                        ph["intrinsic"] = [new_intrinsic_0]
                    int0["pinhole"] = ph
                    if "fisheye" in int0:
                        int0["fisheye"] = ph
                    with open(int_path, "w", encoding="utf-8") as f:
                        yaml.dump(int0, f, allow_unicode=True, default_flow_style=False, sort_keys=False)

        build(str(yaml_dir), str(rebuilt_pb_path))

        with open(pb_path, "r", encoding="utf-8") as f:
            orig_content = f.read()
        with open(rebuilt_pb_path, "r", encoding="utf-8") as f:
            rebuilt_content = f.read()

        orig_msg = parse_pb_text(orig_content)
        rebuilt_msg = parse_pb_text(rebuilt_content)

        # 重组结果必须与原文不等（证明修改被应用）
        assert orig_msg.SerializeToString() != rebuilt_msg.SerializeToString(), (
            "修改 YAML 后重组应与原文不等，否则修改未生效。"
        )

        rebuilt_dict = message_to_dict(rebuilt_msg)

        # 断言修改后的值在重组结果中一致（浮点容差）
        assert rebuilt_dict.get("vehicle_info", {}).get("model") == new_model
        assert rebuilt_dict.get("vehicle_info", {}).get("vin") == new_vin
        assert abs(rebuilt_dict.get("vehicle_param", {}).get("length", 0) - new_length) <= TOL
        assert abs(rebuilt_dict.get("vehicle_param", {}).get("max_speed", 0) - new_max_speed) <= TOL

        extrinsics = rebuilt_dict.get("extrinsics") or []
        if extrinsics:
            t0 = extrinsics[0]
            trans = t0.get("translation") or {}
            assert abs(trans.get("x", 0) - new_translation_x) <= TOL, f"translation.x: {trans}"
            rot = t0.get("rotation") or {}
            assert abs(rot.get("qw", 0) - new_rotation_qw) <= TOL, f"rotation.qw: {rot}"

        intrinsics = rebuilt_dict.get("intrinsics") or []
        if intrinsics:
            cam = intrinsics[0]
            ph = cam.get("pinhole") or cam.get("fisheye") or {}
            assert ph.get("width") == new_pinhole_width, f"pinhole.width: {ph.get('width')}"
            intr = ph.get("intrinsic") or []
            if intr:
                assert abs(float(intr[0]) - new_intrinsic_0) <= TOL, f"intrinsic[0]: {intr[0]}"


def test_long_decimal_within_tolerance_after_roundtrip():
    """S3: 长小数往返后在浮点容差内一致，不要求末位一字不差。"""
    test_dir = Path(__file__).resolve().parent
    pb_path = test_dir / "vehicle_config.pb.txt"
    if not pb_path.exists():
        pytest.skip(f"测试输入不存在: {pb_path}")

    with tempfile.TemporaryDirectory(prefix="pb_cal_roundtrip_") as tmpdir:
        parse(str(pb_path), tmpdir)
        rebuilt_pb_path = Path(tmpdir) / "rebuilt.pb.txt"
        build(tmpdir, str(rebuilt_pb_path))

        with open(pb_path, "r", encoding="utf-8") as f:
            orig_msg = parse_pb_text(f.read())
        with open(rebuilt_pb_path, "r", encoding="utf-8") as f:
            rebuilt_msg = parse_pb_text(f.read())

        orig_d = message_to_dict(orig_msg)
        rebuilt_d = message_to_dict(rebuilt_msg)

        vp_orig = orig_d.get("vehicle_param") or {}
        vp_rebuilt = rebuilt_d.get("vehicle_param") or {}
        for key in ("steer_ratio_a4", "steer_ratio_a5", "steer_ratio_a6", "min_steer_angle_rate"):
            a, b = vp_orig.get(key), vp_rebuilt.get(key)
            if a is not None and b is not None:
                assert abs(float(a) - float(b)) <= TOL, f"vehicle_param.{key}: {a!r} vs {b!r}"

        ext_orig = orig_d.get("extrinsics") or []
        ext_rebuilt = rebuilt_d.get("extrinsics") or []
        if ext_orig and ext_rebuilt:
            qw0 = ext_orig[0].get("rotation") or {}
            qw1 = ext_rebuilt[0].get("rotation") or {}
            a, b = qw0.get("qw"), qw1.get("qw")
            if a is not None and b is not None:
                assert abs(float(a) - float(b)) <= TOL, f"extrinsics[0].rotation.qw: {a!r} vs {b!r}"


if __name__ == "__main__":
    test_roundtrip_pb_to_yaml_to_pb()
    print("OK: 往返测试通过。")
