#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Win32 显示器枚举 / 分辨率切换 / 显示拓扑(纯 ctypes, 无第三方依赖)
================================================================
get_displays()      枚举活动显示器(设备名/显卡/监视器/当前模式/矩形/EDID)
enum_modes()        显卡按该屏 EDID 提供的模式列表
apply_mode()        临时切换分辨率(不写注册表, 重启还原)
ensure_display_extend()  智能拓扑: 仅转接板(EDID 厂商 XLS)接入但未进桌面时强制扩展
"""

import ctypes
import ctypes.wintypes as wt
import winreg

# ============================================================
#  Win32 显示 API
# ============================================================
user32 = ctypes.windll.user32

ENUM_CURRENT_SETTINGS = 0xFFFFFFFF  # (DWORD)-1


class DISPLAY_DEVICEW(ctypes.Structure):
    _fields_ = [
        ('cb', wt.DWORD),
        ('DeviceName', ctypes.c_wchar * 32),
        ('DeviceString', ctypes.c_wchar * 128),
        ('StateFlags', wt.DWORD),
        ('DeviceID', ctypes.c_wchar * 128),
        ('DeviceKey', ctypes.c_wchar * 128),
    ]


class DEVMODEW(ctypes.Structure):
    _fields_ = [
        ('dmDeviceName', ctypes.c_wchar * 32),
        ('dmSpecVersion', wt.WORD),
        ('dmDriverVersion', wt.WORD),
        ('dmSize', wt.WORD),
        ('dmDriverExtra', wt.WORD),
        ('dmFields', wt.DWORD),
        ('dmPositionX', wt.LONG),
        ('dmPositionY', wt.LONG),
        ('dmDisplayOrientation', wt.DWORD),
        ('dmDisplayFixedOutput', wt.DWORD),
        ('dmColor', wt.WORD),
        ('dmDuplex', wt.WORD),
        ('dmYResolution', wt.WORD),
        ('dmTTOption', wt.WORD),
        ('dmCollate', wt.WORD),
        ('dmFormName', ctypes.c_wchar * 32),
        ('dmLogPixels', wt.WORD),
        ('dmBitsPerPel', wt.DWORD),
        ('dmPelsWidth', wt.DWORD),
        ('dmPelsHeight', wt.DWORD),
        ('dmDisplayFlags', wt.DWORD),
        ('dmDisplayFrequency', wt.DWORD),
        ('dmICMMethod', wt.DWORD),
        ('dmICMIntent', wt.DWORD),
        ('dmMediaType', wt.DWORD),
        ('dmDitherType', wt.DWORD),
        ('dmReserved1', wt.DWORD),
        ('dmReserved2', wt.DWORD),
        ('dmPanningWidth', wt.DWORD),
        ('dmPanningHeight', wt.DWORD),
    ]


user32.EnumDisplayDevicesW.argtypes = [ctypes.c_wchar_p, wt.DWORD,
                                       ctypes.POINTER(DISPLAY_DEVICEW), wt.DWORD]
user32.EnumDisplayDevicesW.restype = wt.BOOL
user32.EnumDisplaySettingsW.argtypes = [ctypes.c_wchar_p, wt.DWORD,
                                        ctypes.POINTER(DEVMODEW)]
user32.EnumDisplaySettingsW.restype = wt.BOOL
user32.ChangeDisplaySettingsExW.argtypes = [ctypes.c_wchar_p, ctypes.POINTER(DEVMODEW),
                                            ctypes.c_void_p, wt.DWORD, ctypes.c_void_p]
user32.ChangeDisplaySettingsExW.restype = ctypes.c_long

DM_BITSPERPEL = 0x00040000
DM_PELSWIDTH = 0x00080000
DM_PELSHEIGHT = 0x00100000
DM_DISPLAYFREQUENCY = 0x00400000


def read_edid_from_registry(device_id):
    """monitor DeviceID 'MONITOR\\<pnp>\\<inst>' → 注册表 EDID 原始字节

    注意: 间接显示器/虚拟显示驱动返回的 DeviceID 实例段是 '{类GUID}\\NNNN'
    格式, 注册表中并不存在该路径; 因此精确匹配失败后按 PNP ID 枚举其下
    全部实例, 取第一份 EDID(同一物理屏的重复实例 EDID 相同)。
    """
    if not device_id or not device_id.startswith('MONITOR\\'):
        return None
    parts = device_id.split('\\')
    if len(parts) < 3:
        return None
    pnp = parts[1]
    inst = device_id.split('\\', 2)[2]

    def _read(path):
        try:
            key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, path)
            try:
                val, _typ = winreg.QueryValueEx(key, 'EDID')
                if isinstance(val, (bytes, bytearray)):
                    return bytes(val)
                if isinstance(val, list):          # 某些系统返回 int 列表
                    return bytes(val)
            finally:
                winreg.CloseKey(key)
        except OSError:
            pass
        return None

    # 1) 精确实例路径(常规真实显示器)
    raw = _read(rf'SYSTEM\CurrentControlSet\Enum\DISPLAY\{pnp}\{inst}\Device Parameters')
    if raw:
        return raw

    # 2) 按 PNP ID 遍历全部实例(间接显示器/实例名不一致)
    try:
        pnp_root = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
                                  rf'SYSTEM\CurrentControlSet\Enum\DISPLAY\{pnp}')
    except OSError:
        return None
    try:
        for i in range(winreg.QueryInfoKey(pnp_root)[0]):
            sub = winreg.EnumKey(pnp_root, i)
            raw = _read(rf'SYSTEM\CurrentControlSet\Enum\DISPLAY\{pnp}\{sub}\Device Parameters')
            if raw:
                return raw
    finally:
        winreg.CloseKey(pnp_root)
    return None


def get_displays():
    """枚举活动显示器: 设备名/显卡/监视器名/当前模式/屏幕矩形/EDID 字节"""
    out = []
    i = 0
    while True:
        dd = DISPLAY_DEVICEW()
        dd.cb = ctypes.sizeof(dd)
        if not user32.EnumDisplayDevicesW(None, i, ctypes.byref(dd), 0):
            break
        if dd.StateFlags & 0x1:                 # ATTACHED_TO_DESKTOP
            # 找该输出口上活动的 monitor 条目
            mon = None
            j = 0
            while True:
                m = DISPLAY_DEVICEW()
                m.cb = ctypes.sizeof(m)
                if not user32.EnumDisplayDevicesW(dd.DeviceName, j, ctypes.byref(m), 0):
                    break
                if m.StateFlags & 0x1:
                    mon = m
                    break
                j += 1

            cur = ''
            rect = None
            dm = DEVMODEW()
            dm.dmSize = ctypes.sizeof(dm)
            if user32.EnumDisplaySettingsW(dd.DeviceName, ENUM_CURRENT_SETTINGS,
                                           ctypes.byref(dm)):
                cur = f'{dm.dmPelsWidth}x{dm.dmPelsHeight}@{dm.dmDisplayFrequency}'
                rect = (dm.dmPositionX, dm.dmPositionY, dm.dmPelsWidth, dm.dmPelsHeight)

            out.append({
                'device': dd.DeviceName,         # \\.\DISPLAY1
                'adapter': dd.DeviceString,      # 显卡名
                'mon_name': mon.DeviceString if mon else '',
                'edid': read_edid_from_registry(mon.DeviceID) if mon else None,
                'cur_mode': cur,
                'rect': rect,                    # (x, y, w, h) 物理像素
            })
        i += 1
    return out


def enum_modes(device):
    """枚举 device(\\\\.\\DISPLAYn) 可用模式 —— 显卡按该屏 EDID 提供,
    返回 (['WxH@Hz', ...] 按像素数降序, {'WxH@Hz': DEVMODEW})"""
    seen = {}
    i = 0
    while True:
        dm = DEVMODEW()
        dm.dmSize = ctypes.sizeof(dm)
        if not user32.EnumDisplaySettingsW(device, i, ctypes.byref(dm)):
            break
        i += 1
        if dm.dmBitsPerPel != 32:
            continue
        key = f'{dm.dmPelsWidth}x{dm.dmPelsHeight}@{dm.dmDisplayFrequency}'
        seen.setdefault(key, dm)
    keys = sorted(seen, key=lambda k: tuple(int(v) for v in k.replace('@', 'x').split('x')),
                  reverse=True)
    return keys, seen


def apply_mode(device, dm):
    """把 device 切到 dm 模式(临时切换, 不写注册表, 重启还原), 返回 (成功?, 返回码)"""
    dm.dmFields = DM_BITSPERPEL | DM_PELSWIDTH | DM_PELSHEIGHT | DM_DISPLAYFREQUENCY
    r = user32.ChangeDisplaySettingsExW(device, ctypes.byref(dm), None, 0, None)
    return r == 0, r          # DISP_CHANGE_SUCCESSFUL == 0


def ensure_display_extend():
    """仅当转接板显示器"已接入但未进入桌面"时, 才强制扩展拓扑。

    场景: 板卡热插拔后, 若 Windows 停留在"仅电脑屏幕"拓扑(会记住上次
    Win+P 的选择), 板卡的 EDID 虽已安装、GPU 也在按首选时序输出, 但桌面
    并未扩展过去, get_displays() 就枚举不到这块"显示器" —— 此时才需要
    强制扩展把它拉回桌面。

    用户主动选择的拓扑(家里的"仅第二屏幕"、手动关闭的显示器等)一律不动:
    枚举"接入但不在桌面"的显示器, 只有其 EDID 厂商是转接板签名 "XLS"
    (字节 0x61 0x93, 上游 lt8619c_edid.c) 时才扩展; 笔记本内屏等其它
    厂商视为用户有意关闭, 不干预。

    返回 SetDisplayConfig 的返回码(0=成功, -1=未触发, 其它=异常); 静默容错。
    """
    try:
        board_stuck = False
        i = 0
        while True:
            dd = DISPLAY_DEVICEW()
            dd.cb = ctypes.sizeof(dd)
            if not user32.EnumDisplayDevicesW(None, i, ctypes.byref(dd), 0):
                break
            i += 1
            j = 0
            while True:
                m = DISPLAY_DEVICEW()
                m.cb = ctypes.sizeof(m)
                if not user32.EnumDisplayDevicesW(dd.DeviceName, j, ctypes.byref(m), 0):
                    break
                j += 1
                if m.StateFlags & 0x1:            # ATTACHED_TO_DESKTOP: 已在桌面
                    continue
                raw = read_edid_from_registry(m.DeviceID)
                if len(raw or ()) >= 10 and raw[8:10] == b'\x61\x93':
                    board_stuck = True            # 转接板没进桌面 -> 需要拉回
        if not board_stuck:
            return -1
        SDC_TOPOLOGY_EXTEND = 0x00000004
        SDC_APPLY = 0x00000080
        return ctypes.windll.user32.SetDisplayConfig(
            0, None, 0, None, SDC_TOPOLOGY_EXTEND | SDC_APPLY)
    except Exception:
        return -1