#!/usr/bin/python

# This is a hack to work around a dvdauthor bug which is #323140 in
# the Debian BTS.

# Given a root directory for a DVD tree, check that certain sector
# number fields in the .IFO files match their expected values, and if
# not, fix them up.  Don't bother fixing the .BUP files because
# mkisofs ignores them.

def pread(file, offset, length):
    file.seek(offset, 0)
    return file.read(length)

def pwrite(file, offset, data):
    file.seek(offset, 0)
    file.write(data)

def main(root_dir):

    import os, os.path, struct

    SECTOR_SIZE = 2048
    video_dir = os.path.join(root_dir, 'VIDEO_TS')

    vmgm_name = os.path.join(video_dir, 'VIDEO_TS.IFO')
    vmgm_file = open(vmgm_name, 'r+b')
    ts_count = struct.unpack('>H', pread(vmgm_file, 0x3E, 2))[0]
    ts_sector_diffs = []

    for ts_num in range(ts_count + 1):

        if ts_num == 0:
            ifo_name = vmgm_name
        else:
            ifo_name = os.path.join(video_dir, 'VTS_%02d_0.IFO' % ts_num)
        print 'Checking', ifo_name

        # Find actual sizes
        ifo_size = os.stat(ifo_name).st_size / SECTOR_SIZE
        total_size = ifo_size * 2
        if ts_num == 0:
            vob_name = os.path.join(video_dir, 'VIDEO_TS.VOB')
            if os.path.exists(vob_name):
                total_size = total_size + os.stat(vob_name).st_size / SECTOR_SIZE
        else:
            for vob_num in range(0, 10):
                vob_name = os.path.join(video_dir,
                                        'VTS_%02d_%d.VOB' % (ts_num, vob_num))
                if not os.path.exists(vob_name):
                    break
                total_size = total_size + os.stat(vob_name).st_size / SECTOR_SIZE

        # Fix broken fields
        ifo_file = open(ifo_name, 'r+b')
        def fix(offset, name, value):
            value_read = struct.unpack('>I', pread(ifo_file, offset, 4))[0]
            if value_read != value:
                print ('Changing %s field from 0x%x to 0x%x'
                       % (name, value_read, value))
                value_packed = struct.pack('>I', value)
                pwrite(ifo_file, offset, value_packed)
            return value - value_read
        # I don't know what the real names of the fields are, but these
        # are their apparent meanings.
        ts_sector_diffs.append(fix(0x0C, 'titleset last sector', total_size - 1))
        fix(0x1C, 'IFO last sector', ifo_size - 1)
        fix(0xC0, 'menus first sector', ifo_size) # FIXME: should be 0 if no menus
        ifo_file.close()

    # If any titleset other than the last had a wrong titleset last sector, the
    # TOC entries for the following titlesets will also be wrong.
    if [diff for diff in ts_sector_diffs[:-1] if diff != 0]:
        print 'Updating TOC in', vmgm_name
        title_count = struct.unpack('>H', pread(vmgm_file, SECTOR_SIZE, 2))[0]
        for title_num in range(1, title_count + 1):
            tt_entry = pread(vmgm_file, SECTOR_SIZE + 12 * title_num - 4, 12)
            ts_num = ord(tt_entry[6])
            first_sector = struct.unpack('>I', tt_entry[8:12])[0]
            total_diff = sum(ts_sector_diffs[0:ts_num])
            if total_diff != 0:
                print ('Changing title %d first sector from 0x%x to 0x%x'
                       % (title_num, first_sector, first_sector + total_diff))
                pwrite(vmgm_file, SECTOR_SIZE + 12 * title_num + 4,
                       struct.pack('>I', first_sector + total_diff))

if __name__ == '__main__':
    import sys
    main(sys.argv[1])
