#!/usr/bin/env python3
"""
Forces 2160x3840 portrait @ 30Hz as the ONLY Detailed Timing Descriptor
(DTD slot 1 / preferred timing) and strips the Video Data Block from the
CTA-861 extension so no other video mode is advertised in any way.

Usage:
  python3 force_4k30_portrait_only.py donor.bin fixed.bin
"""

import sys


def build_dtd(pixel_clock_khz, h_active, h_blank, h_sync_offset, h_sync_width,
              v_active, v_blank, v_sync_offset, v_sync_width,
              h_size_mm=0, v_size_mm=0, h_border=0, v_border=0,
              interlaced=False, hsync_pos=True, vsync_pos=True):
    pc = round(pixel_clock_khz / 10)  # EDID stores pixel clock in 10kHz units
    dtd = bytearray(18)
    dtd[0] = pc & 0xFF
    dtd[1] = (pc >> 8) & 0xFF

    dtd[2] = h_active & 0xFF
    dtd[3] = h_blank & 0xFF
    dtd[4] = ((h_active >> 8) << 4) | (h_blank >> 8)

    dtd[5] = v_active & 0xFF
    dtd[6] = v_blank & 0xFF
    dtd[7] = ((v_active >> 8) << 4) | (v_blank >> 8)

    dtd[8] = h_sync_offset & 0xFF
    dtd[9] = h_sync_width & 0xFF
    dtd[10] = ((v_sync_offset & 0xF) << 4) | (v_sync_width & 0xF)
    dtd[11] = (((h_sync_offset >> 8) & 0x3) << 6) | \
              (((h_sync_width >> 8) & 0x3) << 4) | \
              (((v_sync_offset >> 4) & 0x3) << 2) | \
              ((v_sync_width >> 4) & 0x3)

    dtd[12] = h_size_mm & 0xFF
    dtd[13] = v_size_mm & 0xFF
    dtd[14] = ((h_size_mm >> 8) << 4) | (v_size_mm >> 8)

    dtd[15] = h_border
    dtd[16] = v_border

    flags = 0x18  # digital separate sync
    if interlaced:
        flags |= 0x80
    if hsync_pos:
        flags |= 0x02
    if vsync_pos:
        flags |= 0x04
    dtd[17] = flags

    return dtd


TAG_VIDEO = 2


def strip_video_data_block(ext):
    """Remove all Video Data Blocks from a CTA-861 extension block."""
    dtd_offset = ext[2]
    pos = 4  # skip 4-byte CTA header
    blocks = []
    while pos < dtd_offset:
        tag = (ext[pos] >> 5) & 0x07
        length = ext[pos] & 0x1F
        block_data = bytes(ext[pos + 1:pos + 1 + length])
        blocks.append((tag, length, block_data))
        pos += 1 + length

    blocks = [(t, l, d) for t, l, d in blocks if t != TAG_VIDEO]

    new_db = bytearray()
    for tag, length, data in blocks:
        new_db.append((tag << 5) | length)
        new_db.extend(data)

    new_ext = bytearray(128)
    new_ext[0] = ext[0]  # CTA tag
    new_ext[1] = ext[1]  # revision
    new_ext[2] = 4 + len(new_db)  # new DTD offset
    new_ext[3] = ext[3]  # flags (audio, YCbCr, etc.)
    new_ext[4:4 + len(new_db)] = new_db

    new_ext[0x7F] = 0
    new_ext[0x7F] = (-sum(new_ext)) % 256
    return new_ext


if __name__ == "__main__":
    infile, outfile = sys.argv[1], sys.argv[2]
    data = bytearray(open(infile, "rb").read())

    # --- 2160x3840 portrait @ 30Hz ---
    # H blanking uses CVT-RB-style fixed values (front 48 / sync 32 / back 80)
    # V blanking computed to satisfy CVT-RB min vertical blank time (~460us)
    # Pixel clock ~= htotal * vtotal * refresh
    dtd_4k30 = build_dtd(
        pixel_clock_khz=271020,
        h_active=2160, h_blank=160, h_sync_offset=48, h_sync_width=32,
        v_active=3840, v_blank=54, v_sync_offset=3, v_sync_width=5,
        h_size_mm=685, v_size_mm=1218,  # swapped panel physical size for portrait
    )

    # Slot 1 = bytes 54-71 (was 1920x1080@60, the preferred timing)
    data[54:72] = dtd_4k30

    # Slot 2 = bytes 72-89 (was 1280x720@60) - blank out to a dummy
    # descriptor so no second timing is advertised at all.
    # Dummy Descriptor format: 00 00 00 10 00, then 13 zero-padding bytes.
    dummy = bytearray(18)
    dummy[3] = 0x10  # tag byte for "Dummy Descriptor"
    data[72:90] = dummy

    # Unknown Display Descriptor slot and Product Name descriptor are
    # left as-is.

    # Recompute base block checksum (byte 127)
    data[127] = (256 - sum(data[0:127]) % 256) % 256

    # Strip Video Data Block from CTA-861 extension (bytes 128-255)
    data[128:256] = strip_video_data_block(data[128:256])

    open(outfile, "wb").write(data)
    print(f"Wrote {outfile}")
    print("Run `edid-decode` on the output to verify before flashing.")