import argparse
import struct

FLAG_ORDER = 'RWLDErweP_'

TWEAK_OFFSET = -0x4800 # Why? I don't know

# The ADFS docs say that the sector size is either 256 or 1024 bytes.
# But the Domesday images have addresses given in bytes, not sectors.

class DirectoryEntry:
    def __init__(self, f):

        self.name = ''
        self.flags = ''

        for i in range(10):
            c = f.read(1)[0]

            if c & 0x7F:
                self.name += chr(c & 0x7F)

            if c & 0x80:
                self.flags += FLAG_ORDER[i]

        self.name = self.name.strip()
        if ' ' in self.name:
            raise ValueError("Unlikely to be a directory entry")

        self.load_address = struct.unpack('<I', f.read(4))[0]
        self.exec_address = struct.unpack('<I', f.read(4))[0]
        self.length = struct.unpack('<I', f.read(4))[0]
        self.sector = struct.unpack('<I', b'\x00'+f.read(3))[0]
        self.sequence = f.read(1)[0]

    def __str__(self):
        return '%10s %8s %16x %16x %3d' % (
                self.name,
                self.flags,
                self.length,
                self.sector,
                self.sequence,
                )

class Directory:
    def __init__(self, f):

        self.pos = f.tell()
        self.sequence = f.read(1)
        header = f.read(4)
        if header != b'Hugo':
            raise ValueError("this is not a Hugo")

        self.entries = []
        for i in range(47):
            entry = DirectoryEntry(f)
            if entry.name:
                self.entries.append(entry)

        dummy = f.read(1)
        if dummy!=b'\x00':
            raise ValueError("Expecting a null byte here")

        self.name = self._get_name(f, 10)

        self.parent = struct.unpack('<I', b'\x00'+f.read(3))[0]

        self.title = self._get_name(f, 19)

        dummy = f.read(14)
        if dummy!=b'\x00'*14:
            raise ValueError(f"Expecting null bytes here but found {dummy}")

        seq_check = f.read(1)
        if self.sequence!=seq_check:
            raise ValueError(f"Sequence number mismatch: "
            f"{self.sequence} vs {seq_check}")

        trailer = f.read(4)
        if trailer != b'Hugo':
            raise ValueError("trailer mismatch")

    def _get_name(self, f, length):
        result = ''
        for i in range(length):
            c = f.read(1)
            if c!=b'\x00':
                result += chr(ord(c))

        return result

    def __str__(self):

        note = ''

        if self.name=='$':
            note += ' (which is the root)'

        result = (
                f"== directory at efm.dat offset {hex(self.pos)} is "
                f"{self.name}{note}:\n"
                f"   its title is {self.title}\n"
                f"   its parent is at {hex(self.parent)} in this image\n"
         )

        result += '%10s %8s %16s %16s %3s\n' % (
            'name', 'flags', 'length', 'sector', 'seq')

        for entry in self.entries:
            result += str(entry)+'\n'

        return result

def check_hugo(f, pos, extract):

    pos -= 1

    f.seek(pos, 0)

    try:
        directory = Directory(f)
    except ValueError:
        print(f"entry at {hex(pos)} is probably not a directory; skipping")
        return

    image_start = None

    if directory.name=="$":
        # for the root, directory.parent is itself.
        # So we know where on the disk "pos" is.
        image_start = pos-directory.parent

    print(directory)

    print("Sorted into sector order:")
    for entry in sorted(directory.entries,
            key=lambda e: e.sector):
        print(entry)

    """
    GAZETTEER *actually* begins 0x4800 (18432) bytes
    before we think it does. HELP and HELPTEXT do as well.

    But !BOOT is definitely at the right place
    """

    if extract:
        for entry in directory.entries:
            file_start = image_start+entry.sector + TWEAK_OFFSET
            f.seek(file_start)
            contents = f.read(entry.length)
            output_filename = "contents/"+entry.name.lower()
            with open(output_filename, 'wb') as out:
                print('Saving', entry.name,
                        'from efm.dat offset', hex(file_start),
                        'image offset', hex(entry.sector),
                        'length', entry.length,
                        'to', output_filename)
                out.write(contents)

def find_hugo(f):

    result = []

    while True:
        tell = f.tell()
        n = f.read(1024*1024)
        if not n:
            break

        start = 0
        while True:
            try:
                addr = n.index(b'Hugo', start)
                result.append(tell+addr)
                start = addr+4
            except ValueError:
                break

    return result

def main():
    parser = argparse.ArgumentParser(
            description='check ADFS data in Domesday output')
    parser.add_argument(
            'efm_dat', type=argparse.FileType('rb'),
            help='the Domesday .efm.dat file')
    parser.add_argument(
            '--extract', action='store_true',
            help='extract the files to contents/')
    args = parser.parse_args()

    hugos = find_hugo(args.efm_dat)

    check_hugo(args.efm_dat, hugos[0],
            extract=args.extract)

    args.efm_dat.close()

if __name__=='__main__':
    main()