#!/usr/bin/env python3
"""
Extended Display Identification Data Parser (mostly incomplete)
by Mike Crute (mike@crute.us)
licensed under the terms of the MIT license

This is basically just a minimum viable implementation of an EDID parser. There
is a weird subset of features supported because I got excited about parsing
this wacky binary structure. But there's still a lot left un-implemented
because I ran out of time and had already satisfied my use-case.

If you use this and need or want to implement the rest I'll gladly take pull
requests, see the TODOs below.

To use it just call `get_connected_edid` and consider the `EDID` classes
returned.

Based heavily on the information from the following two websites:
* https://en.wikipedia.org/wiki/Extended_Display_Identification_Data
* https://blog.fpmurphy.com/2011/07/retrieve-edid-information-via-sysfs.html
"""
import os
import glob
import struct
from collections import defaultdict


ESTABLISHED_TIMING = {
    0x800000: ( 720,  400, 70),
    0x400000: ( 720,  400, 88),
    0x200000: ( 640,  480, 60),
    0x100000: ( 640,  480, 67),
    0x80000:  ( 640,  480, 72),
    0x40000:  ( 640,  480, 75),
    0x20000:  ( 800,  600, 56),
    0x10000:  ( 800,  600, 60),
    0x8000:   ( 800,  600, 72),
    0x4000:   ( 800,  600, 75),
    0x2000:   ( 832,  624, 75),
    0x1000:   (1024,  768, 87),
    0x800:    (1024,  768, 60),
    0x400:    (1024,  768, 70),
    0x200:    (1024,  768, 75),
    0x100:    (1280, 1024, 75),
    0x80:     (1152,  870, 75),
}

DISPLAY_DESCRIPTOR_ASCII_TYPES = {
    0xFF: "Serial Number",
    0xFE: "Unspecified Text",
    0xFC: "Display Name",
}

DISPLAY_DESCRIPTOR_TYPES = {
    # 6 or 13-byte (with additional timing) binary descriptor
    0xFD: "Display Range Limits",
    # 2x 5-byte descriptors, padded with 0A 20 20.
    0xFB: "Additional White Point Data",
    # 6x 2-byte descriptors, padded with 0A.
    0xFA: "Additional Standard Timing Identifiers",
    0xF9: "Display Color Management",
    0xF8: "CVT 3-Byte Timing Codes",
    0xF7: "Additional Standard Timing 3",
    0x10: "Dummy Identifier",
    0x00: "Manufacturer Reserved Descriptor",
    0x01: "Manufacturer Reserved Descriptor",
    0x02: "Manufacturer Reserved Descriptor",
    0x03: "Manufacturer Reserved Descriptor",
    0x04: "Manufacturer Reserved Descriptor",
    0x05: "Manufacturer Reserved Descriptor",
    0x06: "Manufacturer Reserved Descriptor",
    0x07: "Manufacturer Reserved Descriptor",
    0x08: "Manufacturer Reserved Descriptor",
    0x09: "Manufacturer Reserved Descriptor",
    0x0A: "Manufacturer Reserved Descriptor",
    0x0B: "Manufacturer Reserved Descriptor",
    0x0C: "Manufacturer Reserved Descriptor",
    0x0D: "Manufacturer Reserved Descriptor",
    0x0E: "Manufacturer Reserved Descriptor",
    0x0F: "Manufacturer Reserved Descriptor",
}

BIT_DEPTHS = {
    0: "undefined",
    1: 6,
    2: 8,
    3: 10,
    4: 12,
    5: 14,
    6: 16,
    7: "reserved",
}

DIGITAL_INTERFACES = {
    0: "undefined",
    2: "HMDIa",
    3: "HDMIb",
    4: "MDDI",
    5: "DisplayPort",
}

IMAGE_ASPECT_RATIOS = {
    0: "16:10",
    1: "4:3",
    2: "5:4",
    3: "16:9",
}

class EDID:

    manufacturer_id = None
    product_code = None
    serial_number = None
    mfg_week = None
    mfg_year = None
    edid_version = None
    edid_revision = None
    interface_is_digital = None
    iterface_type = None
    bit_depth = None
    gamma = None
    resolutions = None
    x_resolution = None
    vertical_freqency = None
    image_aspect_ratio = None
    timing_descriptors = None
    display_descriptors = None
    number_of_extensions = None
    checksum = None

    # These are just here for convenience
    display_name = None
    mfg_serial_number = None
    other_mfg_text = None

    # DELL XPS 13 laptops have broken EDID so this tries to still give the
    # display a reasonable name. This might need to be more clever.
    @property
    def reasonable_name(self):
        if self.display_name:
            return self.display_name
        else:
            return "{}_{}".format(self.manufacturer_id, self.product_code)

    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            if hasattr(self, k):
                setattr(self, k, v)

    def __repr__(self):
        attrs = []
        for k, v in self.__dict__.items():
            attrs.append("{}={!r}".format(k, v))

        return "{}({})".format(self.__class__.__name__, ", ".join(attrs))


class DisplayDescriptors(defaultdict):

    def get_first(self, key, default=None):
        return self.get(key, [default])[0]

    def __repr__(self):
        return repr(dict(self))


class EDIDParser:

    _EDID_MAGIC = (0x00, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x00)

    @staticmethod
    def _decode_manufacturer_id(id):
        out = ""

        packed = id[0] << 8 | id[1] # Big endian
        out += chr((packed >> 10 & 0x1f) + 64)
        out += chr((packed >> 5  & 0x1f) + 64)
        out += chr((packed       & 0x1f) + 64)
        chks = packed >> 11 & 0x1f # TODO: chks should be 0

        return out

    @staticmethod
    def _unpack_16bit_le_int(bytes):
        v = struct.unpack("<2B", bytes)
        return v[0] << 8 | v[1]

    @staticmethod
    def _parse_descriptors(bytes):
        timing_descriptors, display_descriptors = [], DisplayDescriptors(list)

        for i in range(54, 126, 18):
            d = bytes[i:i+18]

            # Probably a Display Descriptor
            if d[0:3] == b'\x00\x00\x00' and d[4] == 0x00:
                type, data = d[3], d[5:]

                if type in DISPLAY_DESCRIPTOR_ASCII_TYPES:
                    type_name = DISPLAY_DESCRIPTOR_ASCII_TYPES[type]
                    display_descriptors[type_name].append(
                        data.decode("us-ascii", "backslashreplace").strip())
                elif type in DISPLAY_DESCRIPTOR_TYPES:
                    type_name = DISPLAY_DESCRIPTOR_TYPES[type]
                    display_descriptors[type_name].append(data)
            else:
                # TODO: Parse descriptors
                timing_descriptors.append(d)

        return timing_descriptors, display_descriptors

    @staticmethod
    def _parse_standard_resolutions(bytes):
        resolutions = []

        stdres = (bytes[0] << 16) | (bytes[1] << 8) | bytes[2]
        for k, v in ESTABLISHED_TIMING.items():
            if stdres & k != 0:
                resolutions.append(v)

        return resolutions

    # Expects binary EDID blob, not hex encoded
    def parse(self, edid):
        if edid == b"":
            return

        magic = struct.unpack("8B", edid[:8])
        if magic != self._EDID_MAGIC:
            raise Exception("Invalid magic number")

        manufacturer_id = self._decode_manufacturer_id(
            struct.unpack(">2B", edid[8:10]))

        mfg_product_code = self._unpack_16bit_le_int(edid[10:12])
        serial_number = struct.unpack("<i", edid[12:16])[0]

        mfg_week = edid[16] # Numbers not standardized
        mfg_year = 1990 + edid[17]
        edid_version = edid[18] # 01 for 1.3 and 1.4
        edid_rev = edid[19] # 03 for 1.3 or 04 for 1.4

        is_digital = edid[20] & 0b10000000 == 128
        interface = None
        bit_depth = None

        if is_digital:
            interface = edid[20] & 0b00001111
            bit_depth = (edid[20] & 0b01110000) >> 3

        horizontal_screen_size = edid[21]
        vertical_screen_size = edid[22]
        display_interface_type = DIGITAL_INTERFACES[interface]
        bit_depth = BIT_DEPTHS[bit_depth]
        gamma = edid[23]

        # TODO: parse edid[24] Features bitmap
        # TODO: parse edid[25:35] Chromacity Coordinates

        resolutions = []
        resolutions.extend(self._parse_standard_resolutions(edid[35:38]))

        x_resolution = None
        if edid[38] != 0:
            x_resolution = (edid[38] + 31) * 8

        vertical_freq = (edid[39] & 0b00111111) + 60
        image_aspect_ratio = IMAGE_ASPECT_RATIOS[(edid[39] & 0b11000000) >> 6]

        timing_descriptors, display_descriptors = self._parse_descriptors(edid)

        number_of_extensions = edid[126]
        base_checksum = edid[127]
        remainder = edid[128:]

        # TODO: Parse extensions
        #self._parse_extensions(remainder)

        return EDID(
            manufacturer_id=manufacturer_id,
            product_code=mfg_product_code,
            serial_number=serial_number,
            mfg_week=mfg_week,
            mfg_year=mfg_year,
            edid_version=edid_version,
            edid_revision=edid_rev,
            interface_is_digital=is_digital,
            iterface_type=interface,
            bit_depth=bit_depth,
            gamma=gamma,
            display_name=display_descriptors.get_first("Display Name"),
            mfg_serial_number=display_descriptors.get_first("Serial Number"),
            other_mfg_text=display_descriptors.get_first("Unspecified Text"),
            resolutions=resolutions,
            x_resolution=x_resolution,
            vertical_freqency=vertical_freq,
            image_aspect_ratio=image_aspect_ratio,
            timing_descriptors=timing_descriptors,
            display_descriptors=display_descriptors,
            number_of_extensions=number_of_extensions,
            checksum=base_checksum,
        )

    @staticmethod
    def _parse_extensions(bytes):
        # TODO: Mostly incomplete
        # DTD = Detailed Timing Descriptors
        ext_tag = bytes[0]
        ext_rev = bytes[1]

        # We only sort of understand the "CEA EDID Timing Extension" v3
        if ext_tag != 2 or ext_rev != 3:
            raise Exception("Unknown extension")

        dtd_offset = bytes[2] # 04 means no DTDs, 00 means no data

        supports_underscan = bytes[3] & 0b10000000 != 0
        supports_basic_audio = bytes[3] & 0b01000000 != 0
        supports_ycbcr_444 = bytes[3] & 0b00100000 != 0
        supports_ycbcr_422 = bytes[3] & 0b00010000 != 0
        num_dtds_in_native_format = bytes[3] & 0b00001111

        raw_before_dtds = bytes[4:29]
        raw_dtds = bytes[29:127]
        checksum = bytes[127]

        types = {
            1: "audio",
            2: "video",
            3: "vendor",
            4: "speaker allocation",
            5: "VESA DTC",
        }

        i = 0
        dtds = bytes[4:29]
        while i < len(dtds):
            header = dtds[i]
            i += 1
            type = (header & 0b11100000) >> 5
            length = header & 0b00011111
            body = dtds[i:i+length+1]
            i += length
            print(type, types.get(type), length, body)

        # DTDs are 18-byte packets the same as the main EDID
        # They are padded to 128 bytes with \x00
        # Thus len(raw_dtds) % 18 is padding length

        return None


def get_status(dev_path):
    status_path = os.path.join(dev_path, "status")

    if os.path.exists(status_path):
        with open(status_path, "r") as fp:
            return fp.read().strip()

    return "unknown"


def read_edid_sysfs(dev_path):
    edid_path = os.path.join(dev_path, "edid")
    if not os.path.exists(edid_path):
        return None

    with open(edid_path, "rb") as fp:
        return EDIDParser().parse(fp.read())


def get_connected_edid():
    out = {}

    for g in glob.glob("/sys/class/drm/*"):
        try:
            card, port = os.path.basename(g).split("-", 1)
        except ValueError:
            continue

        if get_status(g) == "connected":
            out[(card, port)] = read_edid_sysfs(g)

    return out


if __name__ == "__main__":
    for (card, display), edid in get_connected_edid().items():
        print("{}:{}".format(display, edid.reasonable_name))