#!/usr/bin/env python3
"""
HDR10 信号输出验证脚本（UCD-323 HDMI Source + PG + sdpg DRM）。

验证 PG（BT.2020 / 10bit / Full RGB）与 sdpg 自定义 DRM(0x87) InfoFrame 能否让 DUT 进入 HDR 模式。

用法（在项目根目录）::

    python test.py
    python test.py --device 0
    python test.py --device ABCD1234 --max-cll 1000 --max-fall 400 --hold 60
    python test.py --save-bin drm_hdr10.bin

验收：脚本发出 100% PQ 白场后，用 CA-410 测 DUT 峰值亮度。
      若仍约 100–200 nits（SDR 水平），说明信令或 DUT 设置未生效。
"""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import UniTAP

# CTA-861 VIC 16 ≈ 3840×2160 @ 60Hz（设备需支持该 timing）
DEFAULT_VIC = 16
DEFAULT_MAX_CLL = 1000
DEFAULT_MAX_FALL = 400
DEFAULT_HOLD_SECONDS = 60
# 10bit PQ 白场码值（与 work_with_solid_color.py 一致）
PQ_WHITE_10BIT = 1023
PACKET_PAD_BYTES = 36


def xy_to_be16(value: float) -> tuple[int, int]:
    """CIE xy → 大端 uint16（单位 0.00002）。"""
    u = int(round(value / 0.00002))
    u = max(0, min(u, 0xFFFF))
    return (u >> 8) & 0xFF, u & 0xFF


def build_drm_hdr10_packet(
    max_cll: int = DEFAULT_MAX_CLL,
    max_fall: int = DEFAULT_MAX_FALL,
    max_master: int = DEFAULT_MAX_CLL,
    min_master_nits: float = 0.005,
    pad_to: int = PACKET_PAD_BYTES,
) -> bytearray:
    """构造 HDR10 静态 DRM InfoFrame（Type 0x87, Static Metadata Type 1）。"""
    body = bytearray(26)
    body[0] = 0x04  # SMPTE ST 2084 (PQ)
    body[1] = 0x01  # Static Metadata Type 1

    primaries = [
        (0.708, 0.292),    # R  BT.2020
        (0.170, 0.797),    # G
        (0.131, 0.046),    # B
        (0.3127, 0.3290),  # W  D65
    ]
    for i, (x, y) in enumerate(primaries):
        o = 2 + i * 4
        body[o], body[o + 1] = xy_to_be16(x)
        body[o + 2], body[o + 3] = xy_to_be16(y)

    body[18:20] = [(max_master >> 8) & 0xFF, max_master & 0xFF]
    min_code = int(round(min_master_nits / 0.0001))
    body[20:22] = [(min_code >> 8) & 0xFF, min_code & 0xFF]
    body[22:24] = [(max_cll >> 8) & 0xFF, max_cll & 0xFF]
    body[24:26] = [(max_fall >> 8) & 0xFF, max_fall & 0xFF]

    header = bytearray([0x87, 0x01, 0x1A, 0x00])
    header[3] = (256 - (sum(header[:3]) + sum(body)) % 256) & 0xFF

    packet = header + body
    if len(packet) < pad_to:
        packet.extend(b"\x00" * (pad_to - len(packet)))
    return packet


def verify_drm_checksum(packet: bytearray) -> bool:
    """校验 InfoFrame checksum（type + version + length + payload + checksum ≡ 0 mod 256）。"""
    if len(packet) < 30 or packet[0] != 0x87:
        return False
    length = packet[2]
    end = 4 + length
    if end > len(packet):
        return False
    return sum(packet[0:3] + packet[4:end] + packet[3:4]) % 256 == 0


def format_packet_hex(packet: bytes, bytes_per_line: int = 16) -> str:
    lines = []
    for i in range(0, len(packet), bytes_per_line):
        chunk = packet[i : i + bytes_per_line]
        hex_part = " ".join(f"{b:02X}" for b in chunk)
        lines.append(f"  {i:04X}: {hex_part}")
    return "\n".join(lines)


def build_hdr10_video_mode(role, vic: int) -> UniTAP.VideoMode:
    timing_mgr = role.hdtx.pg.timing_manager
    timing = timing_mgr.get_cta(vic)
    if timing is None:
        available = timing_mgr.get_all()
        raise RuntimeError(
            f"设备不支持 CTA VIC={vic}。"
            f"可用 timing 数量: {len(available)}，请尝试 --vic 其它值。"
        )

    color = UniTAP.ColorInfo()
    color.color_format = UniTAP.ColorInfo.ColorFormat.CF_RGB
    color.colorimetry = UniTAP.ColorInfo.Colorimetry.CM_ITUR_BT2020_RGB
    color.dynamic_range = UniTAP.ColorInfo.DynamicRange.DR_VESA  # Full
    color.bpc = 10

    return UniTAP.VideoMode(timing=timing, color_info=color)


def apply_pg_white(role, video_mode: UniTAP.VideoMode, white_code: int) -> None:
    pg = role.hdtx.pg
    pg.set_pattern(pattern=UniTAP.VideoPattern.SolidColor)
    pg.set_vm(vm=video_mode)
    if not pg.apply():
        raise RuntimeError("Pattern Generator apply() 失败")
    pg.set_pattern_params(
        pattern_params=UniTAP.SolidColorParams(white_code, white_code, white_code)
    )


def start_sdpg(role, drm_packet: bytearray):
    sdpg = role.hdtx.sdpg
    sdpg.reset()
    sdpg.add_packet(drm_packet)
    sdpg.start()
    return sdpg.status()


def parse_device_arg(value: str):
    if value.isdigit():
        return int(value)
    return value


def list_devices(tsi: UniTAP.TsiLib) -> None:
    try:
        devices = tsi.get_list_of_available_devices()
    except Exception as exc:
        print(f"无法枚举设备: {exc}")
        return
    if not devices:
        print("未发现 UCD 设备。")
        return
    print("已连接设备:")
    for line in devices:
        print(f"  {line}")


def main() -> int:
    parser = argparse.ArgumentParser(
        description="UCD-323 HDR10 信号输出验证（PG + sdpg DRM InfoFrame）"
    )
    parser.add_argument(
        "--device",
        default="0",
        help="设备序号（整数）或 8 位序列号，默认 0",
    )
    parser.add_argument("--vic", type=int, default=DEFAULT_VIC, help=f"CTA VIC，默认 {DEFAULT_VIC}")
    parser.add_argument("--max-cll", type=int, default=DEFAULT_MAX_CLL, help="DRM MaxCLL (nits)")
    parser.add_argument("--max-fall", type=int, default=DEFAULT_MAX_FALL, help="DRM MaxFALL (nits)")
    parser.add_argument(
        "--white-code",
        type=int,
        default=PQ_WHITE_10BIT,
        help=f"PQ 10bit 白场码值，默认 {PQ_WHITE_10BIT}",
    )
    parser.add_argument(
        "--hold",
        type=float,
        default=DEFAULT_HOLD_SECONDS,
        help=f"保持输出秒数，默认 {DEFAULT_HOLD_SECONDS}",
    )
    parser.add_argument(
        "--save-bin",
        metavar="PATH",
        help="将 DRM 包保存为 .bin 文件（便于 sdpg.load_packet 调试）",
    )
    parser.add_argument("--list-devices", action="store_true", help="列出设备后退出")
    args = parser.parse_args()

    tsi = UniTAP.TsiLib()
    if args.list_devices:
        list_devices(tsi)
        tsi.cleanup()
        return 0

    drm = build_drm_hdr10_packet(max_cll=args.max_cll, max_fall=args.max_fall)
    if not verify_drm_checksum(drm):
        print("错误: DRM 包 checksum 校验失败。", file=sys.stderr)
        return 1

    if args.save_bin:
        out = Path(args.save_bin)
        out.write_bytes(drm)
        print(f"已保存 DRM 包: {out.resolve()} ({len(drm)} bytes)")

    print("=== HDR10 信号输出测试 ===")
    print(f"DRM MaxCLL={args.max_cll}, MaxFALL={args.max_fall}")
    print(f"CTA VIC={args.vic}, bpc=10, BT.2020 RGB Full, PQ 白场码值={args.white_code}")
    print("DRM 包内容:")
    print(format_packet_hex(drm))
    print()

    dev = None
    role = None
    try:
        device_id = parse_device_arg(args.device)
        print(f"正在打开设备: {device_id!r} ...")
        dev = tsi.open(device_id)
        role = dev.select_role(UniTAP.dev.UCD323.HDMISource)

        print("配置 Pattern Generator ...")
        vm = build_hdr10_video_mode(role, args.vic)
        t = vm.timing
        refresh_hz = t.frame_rate / 1000.0 if t.frame_rate else 0.0
        print(f"  Timing: {t.hactive}x{t.vactive} @ {refresh_hz:g} Hz")
        print(
            f"  Color: {vm.color_info.colorimetry.name}, "
            f"bpc={vm.color_info.bpc}, range={vm.color_info.dynamic_range.name}"
        )

        apply_pg_white(role, vm, args.white_code)
        print("  PG apply() 成功，已输出 PQ 白场。")

        print("启动 sdpg (DRM InfoFrame) ...")
        status = start_sdpg(role, drm)
        print(f"  sdpg status: {status}")
        if status.error:
            print("  警告: sdpg 报告 error=True，请检查固件或包格式。", file=sys.stderr)
        if not status.active and not status.enabled:
            print("  警告: sdpg 未 active/enabled。", file=sys.stderr)

        print()
        print("--- 请用 CA-410 测量 DUT 100% 白场峰值亮度 ---")
        print("  HDR 生效: 通常显著高于 SDR（具体取决于电视峰值）")
        print("  仍为 ~100–200 nits: 检查 DUT HDR 开关 / HDMI 口 / 信令")
        print(f"--- 保持输出 {args.hold:.0f} 秒，Ctrl+C 可提前结束 ---")
        print()

        deadline = time.monotonic() + args.hold
        while time.monotonic() < deadline:
            time.sleep(1.0)

        print("测试完成。")
        return 0

    except KeyboardInterrupt:
        print("\n用户中断。")
        return 130
    except Exception as exc:
        print(f"错误: {exc}", file=sys.stderr)
        list_devices(tsi)
        return 1
    finally:
        if role is not None:
            try:
                role.hdtx.sdpg.stop()
            except Exception:
                pass
        if dev is not None:
            try:
                tsi.close(dev)
            except Exception:
                pass
        tsi.cleanup()


if __name__ == "__main__":
    sys.exit(main())