#!/bin/sh
set -eu

script_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
fixture_dir=${FIXTURE_DIR:-"$script_dir/fixture"}
artifact_dir=${ARTIFACT_DIR:-"$script_dir/artifacts"}

for tool in mkfs.erofs python3 sha256sum; do
	command -v "$tool" >/dev/null 2>&1 || {
		echo "missing tool: $tool" >&2
		exit 1
	}
done
test ! -e "$fixture_dir" || {
	echo "fixture directory already exists: $fixture_dir" >&2
	exit 1
}
test ! -e "$artifact_dir" || {
	echo "artifact directory already exists: $artifact_dir" >&2
	exit 1
}

mkdir -p "$fixture_dir/fallback" "$fixture_dir/oversize" \
	"$fixture_dir/extent" "$fixture_dir/large/huge" "$artifact_dir"

FIXTURE_DIR="$fixture_dir" python3 <<'PY'
from pathlib import Path
import os

root = Path(os.environ["FIXTURE_DIR"])
(root / "fallback" / "root.txt").write_text(
    "48-bit fallback root\n", encoding="ascii")
(root / "oversize" / "big.dat").write_bytes(b"x")
(root / "extent" / "hole.dat").write_bytes(b"\0" * (1024 * 1024))
(root / "large" / "huge" / "anchor.txt").write_text(
    "large directory anchor\n", encoding="ascii")
for index in range(400):
    (root / "large" / "huge" /
     f"entry-{index:03d}-abcdefghijklmnopqrstuvwxyz.txt").write_text(
        f"entry {index:03d}\n", encoding="ascii")
PY
find "$fixture_dir" -exec touch -h -t 197001010000.00 {} +

build_image()
{
	uuid=$1
	image=$2
	source=$3
	shift 3
	mkfs.erofs -d0 -x-1 -T0 --all-time --all-root --workers=1 \
		-U "$uuid" "$@" "$artifact_dir/$image" "$fixture_dir/$source"
}

build_image 22222222-3333-4444-5555-666666666650 \
	fallback-base.erofs fallback -E force-inode-compact
build_image 22222222-3333-4444-5555-666666666651 \
	oversize-base.erofs oversize -E force-inode-extended
build_image 22222222-3333-4444-5555-666666666652 \
	extent-base.erofs extent -E legacy-compress,force-inode-extended -z lz4
build_image 22222222-3333-4444-5555-666666666653 \
	large-dir-base.erofs large -E force-inode-extended

ARTIFACT_DIR="$artifact_dir" python3 <<'PY'
from pathlib import Path
import hashlib
import os
import struct

artifact_dir = Path(os.environ["ARTIFACT_DIR"])
SUPER = 1024
FEATURE_INCOMPAT = SUPER + 80
ROOTNID_2B = SUPER + 14
BLOCKS_LO = SUPER + 36
META_BLKADDR = SUPER + 40
ROOTNID_8B = SUPER + 112
EROFS_FEATURE_INCOMPAT_48BIT = 0x80
EROFS_INODE_COMPRESSED_FULL = 1
Z_EROFS_ADVISE_EXTENTS = 0x1
Z_EROFS_EXTENT_RECSZ_16 = 0x4
HOLE_SIZE = 5 * 1024 * 1024 * 1024 + 4096
HUGE_DIR_BLOCKS = (1 << 31) + 1


def u16(image, offset):
    return struct.unpack_from("<H", image, offset)[0]


def u32(image, offset):
    return struct.unpack_from("<I", image, offset)[0]


def u64(image, offset):
    return struct.unpack_from("<Q", image, offset)[0]


def put_u16(image, offset, value):
    struct.pack_into("<H", image, offset, value)


def put_u32(image, offset, value):
    struct.pack_into("<I", image, offset, value)


def put_u64(image, offset, value):
    struct.pack_into("<Q", image, offset, value)


def update_superblock_checksum(image):
    polynomial = 0x82F63B78
    checksum = 0xFFFFFFFF
    block_size = 1 << image[SUPER + 12]
    put_u32(image, SUPER + 4, 0)
    for byte in image[SUPER:block_size]:
        checksum ^= byte
        for _ in range(8):
            checksum = (checksum >> 1) ^ (
                polynomial if checksum & 1 else 0)
    put_u32(image, SUPER + 4, checksum & 0xFFFFFFFF)


def inode_offset(image, nid):
    block_bits = image[SUPER + 12]
    return (u32(image, META_BLKADDR) << block_bits) + (nid << 5)


def inode_info(image, nid):
    offset = inode_offset(image, nid)
    inode_format = u16(image, offset)
    inode_size = 64 if inode_format & 1 else 32
    size = u64(image, offset + 8) if inode_size == 64 else u32(
        image, offset + 8)
    xattr_count = u16(image, offset + 2)
    xattr_size = 0 if xattr_count == 0 else 12 + 4 * (xattr_count - 1)
    layout = (inode_format >> 1) & 7
    return offset, inode_format, inode_size, xattr_size, layout, size


def directory_entries(image, nid):
    offset, _, inode_size, xattr_size, layout, size = inode_info(image, nid)
    block_size = 1 << image[SUPER + 12]
    assert 0 < size <= block_size
    if layout == 2:
        data = offset + inode_size + xattr_size
    elif layout == 0:
        data = u32(image, offset + 16) << image[SUPER + 12]
    else:
        raise AssertionError(f"unsupported directory layout {layout}")
    first_nameoff = u16(image, data + 8)
    assert first_nameoff >= 12 and first_nameoff % 12 == 0
    count = first_nameoff // 12
    entries = []
    for index in range(count):
        entry = data + index * 12
        child_nid = u64(image, entry)
        nameoff = u16(image, entry + 8)
        endoff = (u16(image, entry + 20)
                  if index + 1 < count else size)
        name = bytes(image[data + nameoff:data + endoff]).split(b"\0", 1)[0]
        assert name
        entries.append((name, child_nid))
    return entries


def child_nid(image, parent_nid, name):
    return next(nid for entry_name, nid in directory_entries(image, parent_nid)
                if entry_name == name)


def convert_inline_directory_to_plain(image, nid):
    offset, inode_format, inode_size, xattr_size, layout, size = inode_info(
        image, nid)
    assert layout == 2 and inode_size == 64 and size > 0
    block_size = 1 << image[SUPER + 12]
    tail_size = size % block_size
    assert tail_size > 0
    raw_block = u32(image, offset + 16)
    full_size = size - tail_size
    inline_offset = offset + inode_size + xattr_size
    directory_data = bytes(
        image[raw_block * block_size:raw_block * block_size + full_size] +
        image[inline_offset:inline_offset + tail_size])
    assert len(directory_data) == size

    new_raw_block = (len(image) + block_size - 1) // block_size
    image.extend(b"\0" * (new_raw_block * block_size - len(image)))
    image.extend(directory_data)
    image.extend(b"\0" * (-len(image) % block_size))
    put_u16(image, offset, inode_format & ~(7 << 1))
    put_u32(image, offset + 16, new_raw_block)
    put_u32(image, BLOCKS_LO, len(image) // block_size)
    assert inode_info(image, nid)[4] == 0
    return new_raw_block, len(image) // block_size


def write_image(name, image):
    update_superblock_checksum(image)
    (artifact_dir / name).write_bytes(image)


evidence = []

fallback = bytearray((artifact_dir / "fallback-base.erofs").read_bytes())
fallback_root = u16(fallback, ROOTNID_2B)
fallback_blocks = u32(fallback, BLOCKS_LO)
assert fallback_root != 0 and u64(fallback, ROOTNID_8B) == 0
put_u32(fallback, FEATURE_INCOMPAT,
        u32(fallback, FEATURE_INCOMPAT) | EROFS_FEATURE_INCOMPAT_48BIT)
put_u64(fallback, ROOTNID_8B, 0)
write_image("fallback-48bit-root2.erofs", fallback)
evidence.append(
    f"fallback rootnid_2b={fallback_root} rootnid_8b=0 "
    f"blocks_lo={fallback_blocks} union=0x{u16(fallback, ROOTNID_2B):04x}")

oversize = bytearray((artifact_dir / "oversize-base.erofs").read_bytes())
oversize_root = u16(oversize, ROOTNID_2B)
oversize_nid = child_nid(oversize, oversize_root, b"big.dat")
oversize_off, oversize_format, oversize_isize, _, _, _ = inode_info(
    oversize, oversize_nid)
assert oversize_format & 1 and oversize_isize == 64
put_u64(oversize, oversize_off + 8, 1 << 63)
write_image("extended-size-bit63.erofs", oversize)
evidence.append(
    f"oversize nid={oversize_nid} inode_off={oversize_off} "
    f"i_size=0x{u64(oversize, oversize_off + 8):016x}")

extent = bytearray((artifact_dir / "extent-base.erofs").read_bytes())
extent_root = u16(extent, ROOTNID_2B)
extent_nid = child_nid(extent, extent_root, b"hole.dat")
extent_off, extent_format, extent_isize, extent_xattr, extent_layout, _ = \
    inode_info(extent, extent_nid)
assert extent_format & 1 and extent_isize == 64
assert extent_layout == EROFS_INODE_COMPRESSED_FULL
header = (extent_off + extent_isize + extent_xattr + 7) & ~7
record = (header + 8 + 15) & ~15
root_off = inode_offset(extent, extent_root)
assert not (header < root_off + 64 and root_off < record + 16)
put_u64(extent, extent_off + 8, HOLE_SIZE)
struct.pack_into("<IHH", extent, header, 1,
                 Z_EROFS_ADVISE_EXTENTS | Z_EROFS_EXTENT_RECSZ_16, 0)
struct.pack_into("<IIII", extent, record, 0, 0, 0, 0)
write_image("extent-hole-5g.erofs", extent)
evidence.append(
    f"extent-hole nid={extent_nid} inode_off={extent_off} size={HOLE_SIZE} "
    f"header={header} record={record} extents=1 recsize=16 plen=0 lstart=0")

large = bytearray((artifact_dir / "large-dir-base.erofs").read_bytes())
large_root = u16(large, ROOTNID_2B)
large_nid = child_nid(large, large_root, b"huge")
large_off, large_format, large_isize, _, _, _ = inode_info(large, large_nid)
assert large_format & 1 and large_isize == 64
large_raw_block, large_image_blocks = convert_inline_directory_to_plain(
    large, large_nid)
large_size = HUGE_DIR_BLOCKS * (1 << large[SUPER + 12])
put_u64(large, large_off + 8, large_size)
write_image("large-dir-intmax.erofs", large)
evidence.append(
    f"large-dir nid={large_nid} inode_off={large_off} size={large_size} "
    f"blocks={HUGE_DIR_BLOCKS} last_block={HUGE_DIR_BLOCKS - 1} "
    f"layout=flat-plain raw_block={large_raw_block} "
    f"image_blocks={large_image_blocks}")

(artifact_dir / "fixture-evidence.txt").write_text(
    "\n".join(evidence) + "\n", encoding="ascii")
with (artifact_dir / "SHA256SUMS").open("w", encoding="ascii") as sums:
    for path in sorted(artifact_dir.glob("*.erofs")):
        sums.write(f"{hashlib.sha256(path.read_bytes()).hexdigest()}  "
                   f"{path.name}\n")
PY

cat "$artifact_dir/fixture-evidence.txt"
cat "$artifact_dir/SHA256SUMS"