#!/usr/bin/env python3
"""Generate a virtual-display EDID with 4K CTA VICs (trimmed from EnriqueWood POC)."""

import math
import sys


def cvt_rb(w, h, refresh):
    hblank, hfront, hsync = 160, 48, 32
    htotal = w + hblank
    vfront, vsync = 3, 8
    rb_min_vblank = 460e-6
    vblank = math.ceil(rb_min_vblank * refresh * h / (1.0 - rb_min_vblank * refresh))
    vblank = max(vblank, vfront + vsync + 6)
    vtotal = h + vblank
    pixclk_khz = int(round(htotal * vtotal * refresh / 10000.0)) * 10
    return dict(
        w=w,
        h=h,
        refresh=refresh,
        htotal=htotal,
        vtotal=vtotal,
        hblank=hblank,
        hfront=hfront,
        hsync=hsync,
        vblank=vblank,
        vfront=vfront,
        vsync=vsync,
        pixclk_khz=pixclk_khz,
    )


def make_dtd(t, hmm=600, vmm=340):
    w, h = t["w"], t["h"]
    hblank, vblank = t["hblank"], t["vblank"]
    hfront, hsync = t["hfront"], t["hsync"]
    vfront, vsync = t["vfront"], t["vsync"]
    pc = t["pixclk_khz"] // 10
    d = bytearray(18)
    d[0] = pc & 0xFF
    d[1] = (pc >> 8) & 0xFF
    d[2] = w & 0xFF
    d[3] = hblank & 0xFF
    d[4] = ((w >> 4) & 0xF0) | ((hblank >> 8) & 0x0F)
    d[5] = h & 0xFF
    d[6] = vblank & 0xFF
    d[7] = ((h >> 4) & 0xF0) | ((vblank >> 8) & 0x0F)
    d[8] = hfront & 0xFF
    d[9] = hsync & 0xFF
    d[10] = ((vfront & 0x0F) << 4) | (vsync & 0x0F)
    d[11] = (
        (((hfront >> 8) & 3) << 6)
        | (((hsync >> 8) & 3) << 4)
        | (((vfront >> 4) & 3) << 2)
        | ((vsync >> 4) & 3)
    )
    d[12] = hmm & 0xFF
    d[13] = vmm & 0xFF
    d[14] = ((hmm >> 4) & 0xF0) | ((vmm >> 8) & 0x0F)
    d[17] = 0x1E
    return bytes(d)


def make_name(name):
    d = bytearray(18)
    d[3] = 0xFC
    nb = name.encode("ascii")[:13]
    d[5 : 5 + len(nb)] = nb
    pos = 5 + len(nb)
    if pos < 18:
        d[pos] = 0x0A
        pos += 1
    for i in range(pos, 18):
        d[i] = 0x20
    return bytes(d)


def make_range(min_v, max_v, min_h, max_h, max_pixclk_mhz):
    d = bytearray(18)
    d[3] = 0xFD
    flags = 0
    mv, mh = max_v, max_h
    if mh > 255:
        flags |= 0x08
        mh -= 255
    if mv > 255:
        flags |= 0x02
        mv -= 255
    d[4] = flags
    d[5] = min_v
    d[6] = mv
    d[7] = min_h
    d[8] = mh
    d[9] = max_pixclk_mhz // 10
    d[10] = 0x01
    for i in range(11, 18):
        d[i] = 0x0A if i == 11 else 0x20
    return bytes(d)


def checksum(block):
    return (256 - (sum(block[:-1]) % 256)) % 256


VICS = [97, 118, 96, 95, 63, 16, 4]


def build_base(base_dtds, name, max_pixclk_mhz):
    b = bytearray(128)
    b[0:8] = b"\x00\xFF\xFF\xFF\xFF\xFF\xFF\x00"
    b[8:10] = b"\x50\x74"
    b[10:12] = b"\x70\x02"
    b[16] = 22
    b[17] = 34
    b[18] = 1
    b[19] = 4
    b[20] = 0xA5
    b[21] = 60
    b[22] = 34
    b[23] = 120
    b[24] = 0x2E
    b[25:35] = bytes([0x35, 0x85, 0xA6, 0x56, 0x48, 0x9A, 0x24, 0x12, 0x50, 0x54])
    for i in range(8):
        b[38 + i * 2] = 0x01
        b[39 + i * 2] = 0x01
    b[54:72] = make_dtd(base_dtds[0])
    b[72:90] = make_dtd(base_dtds[1])
    b[90:108] = make_range(24, 240, 30, 510, max_pixclk_mhz)
    b[108:126] = make_name(name)
    b[126] = 1
    b[127] = checksum(b)
    return b


def build_cta_primary(dtds):
    b = bytearray(128)
    b[0] = 0x02
    b[1] = 0x03
    pos = 4
    b[pos] = (2 << 5) | len(VICS)
    for i, v in enumerate(VICS):
        b[pos + 1 + i] = v
    pos += 1 + len(VICS)
    b[pos] = (1 << 5) | 3
    b[pos + 1], b[pos + 2], b[pos + 3] = 0x17, 0x7F, 0x07
    pos += 4
    b[pos] = (4 << 5) | 3
    b[pos + 1] = 0x4F
    pos += 4
    b[2] = pos
    b[3] = 0x40
    used, p = [], pos
    for t in dtds:
        if p + 18 > 127:
            break
        b[p : p + 18] = make_dtd(t)
        used.append(t)
        p += 18
    b[127] = checksum(b)
    return b, used


def build(custom, name, max_pixclk_mhz=1200):
    if len(custom) < 2:
        custom = custom + custom[:1] * (2 - len(custom))
    base = build_base(custom[0:2], name, max_pixclk_mhz)
    rest = custom[2:]
    cta1, _used = build_cta_primary(rest)
    base[126] = 1
    base[127] = checksum(base)
    return bytes(base + cta1)


def main():
    output = sys.argv[1]
    custom = [
        cvt_rb(3840, 2160, 60),
        cvt_rb(2560, 1440, 60),
        cvt_rb(1920, 1080, 60),
    ]
    edid = build(custom, "SUN-4K")
    with open(output, "wb") as f:
        f.write(edid)
    assert len(edid) % 128 == 0
    for i in range(len(edid) // 128):
        block = edid[i * 128 : (i + 1) * 128]
        assert sum(block) % 256 == 0


if __name__ == "__main__":
    main()