"""A filesystem small enough to read in one sitting: an inode table, a block
bitmap, and an array of fixed-size data blocks, all inside ONE flat file.

    +--------+--------------+---------+-------------------------+
    | super  | inode table  | bitmap  | data blocks 0..N-1      |
    +--------+--------------+---------+-------------------------+

A name resolves to an inode; an inode holds a size and a list of block
numbers; a block is `block_size` bytes at a computable offset. No
directories, no permissions, no journal. What the toy is *about* is what
`rm` does and what it doesn't:

1. `rm` clears the inode and flips the file's bits in the bitmap back to 0.
   It does not touch the data blocks -- that would cost one disk write per
   block, on an operation the user expects to be instant.
2. The next `write` first-fits the bitmap and gets those same blocks back.
   A short file only overwrites the *front* of its first block.
3. Reads are capped at `inode.size`, so the leftover tail is unreachable...
   until a file's size legitimately grows. `truncate()` is the crack: POSIX
   says growing a file reads back as zeros, and this toy doesn't do that.

Mount with `zero_on_free` or `zero_on_alloc` to close the crack, and watch
what each costs in `self.block_writes`.

Deterministic: no clock, no randomness, first-fit allocation, and
`format()` always starts from an all-zero image.
"""

import os
import struct
import sys

MAGIC = b"TINYFS\0\0"
SUPER_SIZE = 32          # magic(8) + block_size, num_blocks, num_inodes, direct
SUPER_FMT = "<8sIIII"
NAME_LEN = 24


class Inode:
    """One fixed-size record: is it in use, how many bytes long is the file,
    what is it called, and which blocks hold it. Fixed-size is the point --
    inode number `i` is at a computable offset, so there is no index to
    search and no allocator to run when you want inode 5.
    """

    def __init__(self, used=0, size=0, name=b"", blocks=()):
        self.used = used
        self.size = size
        self.name = name
        self.blocks = list(blocks)

    def pack(self, direct):
        blocks = (self.blocks + [0] * direct)[:direct]
        return struct.pack("<BI%ds%dI" % (NAME_LEN, direct), self.used,
                           self.size, self.name.ljust(NAME_LEN, b"\0"), *blocks)

    @classmethod
    def unpack(cls, raw, direct):
        used, size, name, *blocks = struct.unpack(
            "<BI%ds%dI" % (NAME_LEN, direct), raw)
        return cls(used, size, name.rstrip(b"\0"), blocks)


class FileSystem:
    """A mounted image. `format` lays one out; `mount` reads the geometry
    back off the superblock, which is the only reason a superblock exists.
    """

    def __init__(self, path, block_size, num_blocks, num_inodes, direct,
                 zero_on_free=False, zero_on_alloc=False):
        self.path = path
        self.block_size = block_size
        self.num_blocks = num_blocks
        self.num_inodes = num_inodes
        self.direct = direct
        # Mount options, not on-disk state: a filesystem's zeroing policy is
        # a property of the driver that mounted it, not of the image.
        self.zero_on_free = zero_on_free
        self.zero_on_alloc = zero_on_alloc
        self.block_writes = 0                      # the I/O bill, counted
        self.inode_size = struct.calcsize("<BI%ds%dI" % (NAME_LEN, direct))
        self.inode_start = SUPER_SIZE
        self.bitmap_start = self.inode_start + num_inodes * self.inode_size
        self.data_start = self.bitmap_start + num_blocks
        self.image_size = self.data_start + num_blocks * block_size

    @classmethod
    def format(cls, path, block_size=64, num_blocks=32, num_inodes=8,
               direct=4, **opts):
        fs = cls(path, block_size, num_blocks, num_inodes, direct, **opts)
        with open(path, "wb") as f:
            f.write(b"\0" * fs.image_size)
        fs._write(0, struct.pack(SUPER_FMT, MAGIC, block_size, num_blocks,
                                 num_inodes, direct))
        return fs

    @classmethod
    def mount(cls, path, **opts):
        with open(path, "rb") as f:
            raw = f.read(SUPER_SIZE)
        if raw[:len(MAGIC)] != MAGIC:
            raise OSError("not a tinyfs image: %s" % path)
        bs, nb, ni, direct = struct.unpack("<4I", raw[len(MAGIC):24])
        return cls(path, bs, nb, ni, direct, **opts)

    # -- the "device": every read and write goes through these two ---------
    def _read(self, off, n):
        with open(self.path, "rb") as f:
            f.seek(off)
            return f.read(n)

    def _write(self, off, data):
        with open(self.path, "r+b") as f:
            f.seek(off)
            f.write(data)

    # -- metadata ----------------------------------------------------------
    def inode(self, i):
        return Inode.unpack(
            self._read(self.inode_start + i * self.inode_size,
                       self.inode_size), self.direct)

    def put_inode(self, i, ino):
        self._write(self.inode_start + i * self.inode_size,
                    ino.pack(self.direct))

    def bitmap(self):
        return bytearray(self._read(self.bitmap_start, self.num_blocks))

    def put_bitmap(self, bm):
        self._write(self.bitmap_start, bytes(bm))

    def nblocks(self, size):
        return (size + self.block_size - 1) // self.block_size

    # -- data blocks -------------------------------------------------------
    def block(self, b):
        return self._read(self.data_start + b * self.block_size,
                          self.block_size)

    def put_block(self, b, data):
        self.block_writes += 1
        self._write(self.data_start + b * self.block_size,
                    data[:self.block_size])

    def alloc_block(self):
        """First fit: the lowest-numbered free block wins. Deterministic,
        two lines, and the reason a just-freed block comes straight back.
        """
        bm = self.bitmap()
        for b in range(self.num_blocks):
            if not bm[b]:
                bm[b] = 1
                self.put_bitmap(bm)
                if self.zero_on_alloc:
                    self.put_block(b, b"\0" * self.block_size)
                return b
        raise OSError("no free blocks")

    def free_block(self, b):
        bm = self.bitmap()
        bm[b] = 0
        self.put_bitmap(bm)
        if self.zero_on_free:
            self.put_block(b, b"\0" * self.block_size)

    # -- the file API ------------------------------------------------------
    def lookup(self, name):
        for i in range(self.num_inodes):
            ino = self.inode(i)
            if ino.used and ino.name == name.encode():
                return i, ino
        return None, None

    def ls(self):
        out = []
        for i in range(self.num_inodes):
            ino = self.inode(i)
            if ino.used:
                out.append((ino.name.decode(), ino.size,
                            ino.blocks[:self.nblocks(ino.size)]))
        return out

    def write(self, name, data):
        if self.lookup(name)[1]:
            self.rm(name)
        for i in range(self.num_inodes):
            if not self.inode(i).used:
                break
        else:
            raise OSError("no free inodes")
        if self.nblocks(len(data)) > self.direct:
            raise OSError("file too large: %d direct blocks" % self.direct)
        blocks = []
        for off in range(0, len(data), self.block_size):
            b = self.alloc_block()
            blocks.append(b)
            # Only len(data)-off bytes are written. Whatever sat in the rest
            # of this block is still sitting there. That is the whole toy.
            self.put_block(b, data[off:off + self.block_size])
        self.put_inode(i, Inode(1, len(data), name.encode(), blocks))
        return i

    def read(self, name):
        i, ino = self.lookup(name)
        if ino is None:
            raise FileNotFoundError(name)
        out = b"".join(self.block(ino.blocks[n])
                       for n in range(self.nblocks(ino.size)))
        return out[:ino.size]      # the cap that hides the tail

    def truncate(self, name, newsize):
        """POSIX `ftruncate`: any owner may resize their own file, and the
        bytes a grow exposes are required to read back as zeros. This
        implementation just moves the size field -- which is the bug.
        """
        i, ino = self.lookup(name)
        if ino is None:
            raise FileNotFoundError(name)
        need, have = self.nblocks(newsize), self.nblocks(ino.size)
        if need > self.direct:
            raise OSError("file too large: %d direct blocks" % self.direct)
        # Trim first: unpack always hands back `direct` entries, and the ones
        # past the file's length are whatever the last, longer file left.
        blocks = ino.blocks[:have]
        for _ in range(have, need):
            blocks.append(self.alloc_block())
        for b in blocks[need:]:
            self.free_block(b)
        ino.blocks = blocks[:need]
        ino.size = newsize
        self.put_inode(i, ino)

    def rm(self, name):
        """Free the blocks, clear the inode. Note what is missing: any write
        to the data region at all, unless zero_on_free is set.
        """
        i, ino = self.lookup(name)
        if ino is None:
            raise FileNotFoundError(name)
        for b in ino.blocks[:self.nblocks(ino.size)]:
            self.free_block(b)
        self.put_inode(i, Inode())


def main(argv):
    opts = {"zero_on_free": "--zero-on-free" in argv,
            "zero_on_alloc": "--zero-on-alloc" in argv}
    args = [a for a in argv if not a.startswith("--")]
    img = os.environ.get("TINYFS_IMG", "disk.img")
    cmd, rest = (args[0], args[1:]) if args else ("help", [])
    if cmd == "format":
        FileSystem.format(img)
        return print("formatted %s" % img)
    fs = FileSystem.mount(img, **opts)
    if cmd == "write":
        fs.write(rest[0], rest[1].encode() if len(rest) > 1
                 else sys.stdin.buffer.read())
    elif cmd == "read":
        sys.stdout.buffer.write(fs.read(rest[0]))
    elif cmd == "rm":
        fs.rm(rest[0])
    elif cmd == "truncate":
        fs.truncate(rest[0], int(rest[1]))
    elif cmd == "ls":
        for name, size, blocks in fs.ls():
            print("%-14s %4d bytes  blocks=%s" % (name, size, blocks))
    else:
        print(__doc__.strip().splitlines()[0])
        print("usage: tinyfs.py [--zero-on-free|--zero-on-alloc] "
              "format|ls|write NAME [DATA]|read NAME|rm NAME|truncate NAME N")


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