# Name: Maxim van den Berg
# ID: 11867000
# Course: Besturingssystemen
#
# Opdracht 3: Filesystem

import sys
import struct
import time

# Constants
BLOCK_SIZE = 1024
INDEX_SIZE = 16
INODE_SIZE = 32
DIR_ENTRY_SIZE = -1
MINIX_SUPER_MAGIC = 0x137F
S_IFREG = 0o0100000
S_IFDIR = 0o0040000
S_IRUSR = 0o00400
S_IWUSR = 0o00200
S_IXUSR = 0o00100

unpack_format = {1: "<B", 2: "<H", 4: "<I"}

# Parse the superblock. Convert data to dictionary.
def parse_superblock(sbdata):
    sbdict = {}

    # Tuple of (name, number of byte) for all items in the superblock.
    sbitems = [("#inodes", 2) , ("#zones", 2), ("imap blocks", 2),
               ("zmap blocks", 2), ("first data zone", 2),
               ("log zone size", 2), ("max file size", 4),
               ("minix magic", 2), ("mount state", 2)]

    idx = 0
    for name, length in sbitems:
        (sbdict[name],) = struct.unpack(unpack_format[length],
                 sbdata[idx:idx + length])
        idx += length

    return sbdict

# Convert inode number to byte index in filesystem.
def get_inode_index(sbdict, inode_n):
    return ((2 + sbdict["imap blocks"] + sbdict["zmap blocks"]) * BLOCK_SIZE
            + (inode_n - 1) * INODE_SIZE)

# Convert inode data to inode dictionary.
def parse_inode_entry(inodedata):
    inodedict = {}

    inodeitems = [("mode", 2) , ("uid", 2), ("size", 4),
                  ("time", 4), ("gid", 1), ("links", 1)]

    idx = 0
    for name, length in inodeitems:
        (inodedict[name],) = struct.unpack(unpack_format[length],
                 inodedata[idx:idx + length])
        idx += length


    # Parse the zones.
    inodedict["zones"] = list(struct.unpack("<9H", inodedata[idx:]))

    return inodedict

# List all files in a given block.
def block_list_files(block_data):
    files = []

    # Loop over all directory entries in the block.
    for entry in range(BLOCK_SIZE // DIR_ENTRY_SIZE):
        loc = entry * DIR_ENTRY_SIZE
        inode = struct.unpack("<H", block_data[loc:loc + 2])[0]

        # Check if the directory entry is filled in.
        if not inode:
            continue

        # Get the file name.
        string = struct.unpack(f"{DIR_ENTRY_SIZE - 2}c",
                block_data[loc + 2:loc + DIR_ENTRY_SIZE])

        files.append((inode, b''.join(string).decode().rstrip('\x00')))

    return files

# List all files in the blocks linked to by the given block.
def block_list_files_indirect(block_data, f):
    files = []

    for entry in range(BLOCK_SIZE // INDEX_SIZE):
        loc = entry * INDEX_SIZE
        f.seek(struct.unpack("<H",
            block_double_indirect[loc:loc + INDEX_SIZE])[0])

        files.extend(block_list_files(f.read(BLOCK_SIZE)))

    return files

# Print all files in directory inode.
def dir_list_files(inode, f):
    files = []

    if not (inode["mode"] & S_IFDIR):
        print(f"Error: inode is not a directory",
                file=sys.stderr)
        return None

    for i, zone in enumerate(inode["zones"]):
        if not zone:
            continue

        f.seek(zone * BLOCK_SIZE)

        if i < 7:
            files.extend(block_list_files(f.read(BLOCK_SIZE)))
        elif i == 7:
            files.extend(block_list_files_indirect(f.read(BLOCK_SIZE), f))
        else:
            block_double_indirect = f.read(BLOCK_SIZE)

            for entry in range(BLOCK_SIZE / INDEX_SIZE):
                loc = entry * INDEX_SIZE
                f.seek(struct.unpack("<H",
                    block_double_indirect[loc:loc + INDEX_SIZE])[0])

                files.extend(block_list_files_indirect(f.read(BLOCK_SIZE), f))

    return files

# Run the ls command.
def cmd_ls(sbdict, rootdict, f):
    # List all files in the root directory
    files = dir_list_files(rootdict, f)
    for _, file in files:
        print(file)

# Run the cat command.
def cmd_cat(sbdict, root, f, filepath):
    inode = root;

    # Traverse the inode tree to get to the file.
    for p in filepath.split("/"):
        # Find the file or directory in the current directory.
        files = dir_list_files(inode, f)
        try:
            inode_n_next = files[[file for _, file in files].index(p)][0]
        except ValueError:
            print(f"Error: file not found", file=sys.stderr)
            return

        # Goto the next inode.
        f.seek(get_inode_index(sbdict, inode_n_next))
        inode = parse_inode_entry(f.read(INODE_SIZE))

    # The last inode should be the file we are looking for.
    # We check if it is a regular file as is required.
    if not inode["mode"] & S_IFREG:
        print(f"Error: invalid file type", file=sys.stderr)
        return

    # We print the file contents (with the assumption that
    # no indirect zones exists)
    for block_n in inode["zones"][0:7]:
        f.seek(block_n * BLOCK_SIZE)
        buffer = f.read(BLOCK_SIZE)
        # sys.stdout.buffer.write(buffer)

# Check if file/directory can be created.
def check_creation(name, files):
    if "/" in name:
        print(f"Error: touch only implemented for files in the root dir",
                file=sys.stderr)
        return False

    if files != None and name in [file for _, file in files]:
        print(f"Error: file already exists", file=sys.stderr)
        return False

    if len(name) > DIR_ENTRY_SIZE - 2:
        print(f"Error: filename too long", file=sys.stderr)
        return False

    return True


# Find the index of the the first bit that is equal to 0
def first_available(buffer):
    for j, b in enumerate(buffer):
        # Find byte with a open spot.
        value = b
        if value != 255:
            for shift in range(8):
                # Find the first bit which equals 0.
                if (value >> shift) & 1 == 0:
                    return shift + 8 * j, value


# Extra operations necessary for the creation of directories
def cmd_create_dir_extra(sbdict, f, zone_claim, zone_byte, inode_claim):
    # Create the directory zone and fill it with '.' (itself) and '..' (root).
    f.seek((sbdict["first data zone"] + zone_claim - 1) * BLOCK_SIZE)
    f.write(struct.pack(f"<Hc", inode_claim + 1, '.'.encode()))
    f.seek((sbdict["first data zone"] + zone_claim - 1) * BLOCK_SIZE + DIR_ENTRY_SIZE)
    f.write(struct.pack(f"<Hcc", 1, '.'.encode(), '.'.encode()))

    # Claim the zone if it is a directory.
    f.seek((2 + sbdict["imap blocks"]) * BLOCK_SIZE + zone_claim // 8)
    f.write(struct.pack("<B", zone_byte | (1 << (zone_claim % 8))))

    # Add the first zone reference to the inode.
    f.seek(get_inode_index(sbdict, inode_claim + 1) + 30)
    f.write(struct.pack("<H", sbdict["first data zone"] + zone_claim - 1))


# Run the touch or mkdir command.
def cmd_create(sbdict, root, f, diskimg, filename, is_dir):
    files = dir_list_files(root, f)

    if not check_creation(filename, files):
        return

    # Find the first available inode.
    f.seek(2 * BLOCK_SIZE)
    inode_claim, inode_byte = first_available(
            f.read(sbdict["imap blocks"] * BLOCK_SIZE))

    # Find the first available zone if it is a directory.
    if is_dir:
        f.seek((2 + sbdict["imap blocks"]) * BLOCK_SIZE)
        zone_claim, zone_byte = first_available(
                f.read(sbdict["zmap blocks"] * BLOCK_SIZE))

    # Close reading and start writing.
    f.close()
    f = open(diskimg, 'r+b')

    # Claim the inode.
    f.seek(2 * BLOCK_SIZE + inode_claim // 8)
    f.write(struct.pack("<B", inode_byte | (1 << (inode_claim % 8))))

    if is_dir:
        cmd_create_dir_extra(sbdict, f, zone_claim, zone_byte, inode_claim)

    # Create the inode.
    f.seek(get_inode_index(sbdict, inode_claim + 1))
    f.write(struct.pack("<HHIIBB", (S_IFREG if not is_dir else S_IFDIR
            | S_IRUSR | S_IWUSR | S_IXUSR), 0, 0, int(time.time()), 0, 1))

    # Add the entry to the first zone of the root directory.
    f.seek(root["zones"][0] * BLOCK_SIZE + len(files) * DIR_ENTRY_SIZE)
    f.write(struct.pack(f"<H", inode_claim + 1))
    f.write(struct.pack(f"{DIR_ENTRY_SIZE - 2}B",
        *filename.ljust(DIR_ENTRY_SIZE - 2, "\x00").encode()))

    # Update the root inode (due to our assumptions, only the size
    # and modification time are necessary).
    f.seek(get_inode_index(sbdict, 1) + 4)
    f.write(struct.pack("<I", root["size"] + DIR_ENTRY_SIZE))

    f.seek(get_inode_index(sbdict, 1) + 8)
    f.write(struct.pack("<I", int(time.time())))

    f.close()

# Run the mkdir command.
def cmd_mkdir(sbdict, root, f, dirname):
    files = dir_list_files(root, f)

    if not check_creation(filename, files):
        return

    # Find the first available block and claim it.

    # Find the first available inode and claim it.

    # Add the entry to the first zone of the root directory.
    f.seek(root["zones"][0] * BLOCK_SIZE + len(files) * DIR_ENTRY_SIZE)
    f.write()

    # Update the root inode

    # Update the superblock


# Check if the program argument count is correct (for some commands).
def check_argument_count(arguments):
    if len(arguments) != 4:
        print(f"Error: please specify a (single) file",
                file=sys.stderr)


# The main function.
if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: mfstool.py image command [params]")
        sys.exit(0)

    diskimg = sys.argv[1]
    cmd = sys.argv[2]

    f = open(diskimg, 'rb')

    # Read the super block.
    f.seek(BLOCK_SIZE)
    sbdict = parse_superblock(f.read(BLOCK_SIZE))
    DIR_ENTRY_SIZE = 16 if sbdict["minix magic"] == MINIX_SUPER_MAGIC else 32

    # Go to the inode table and parse the first inode (the root dir).
    f.seek(get_inode_index(sbdict, 1))
    rootdict = parse_inode_entry(f.read(INODE_SIZE))

    if cmd == "ls":
        cmd_ls(sbdict, rootdict, f)
    elif cmd == "cat":
        check_argument_count(sys.argv)
        cmd_cat(sbdict, rootdict, f, sys.argv[3])
    elif cmd == "touch":
        check_argument_count(sys.argv)
        cmd_create(sbdict, rootdict, f, diskimg, sys.argv[3], False)
    elif cmd == "mkdir":
        check_argument_count(sys.argv)
        cmd_create(sbdict, rootdict, f, diskimg, sys.argv[3], True)
    elif cmd == "append":
        ...