#!/bin/sh
set -eu

script_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
fixture_dir="$script_dir/fixture"
artifact_dir="$script_dir/artifacts"

mkdir -p \
	"$fixture_dir/alpha/bravo/charlie" \
	"$fixture_dir/alpha/sibling" \
	"$fixture_dir/wide" \
	"$artifact_dir"

printf '%s\n' 'cold nested lookup payload' > \
	"$fixture_dir/alpha/bravo/charlie/payload.txt"
printf '%s\n' 'repeat lookup payload' > \
	"$fixture_dir/alpha/bravo/repeat.txt"
printf '%s\n' 'sibling marker' > \
	"$fixture_dir/alpha/sibling/marker.txt"

index=0
while [ "$index" -lt 320 ]; do
	name=$(printf 'entry-%03d-abcdefghijklmnopqrstuvwxyz.txt' "$index")
	printf 'wide entry %03d\n' "$index" > "$fixture_dir/wide/$name"
	index=$((index + 1))
done

find "$fixture_dir" -exec touch -h -t 197001010000.00 {} +

build_image()
{
	image=$1
	mkfs.erofs -d0 -x-1 -T0 --all-time --all-root --workers=1 \
		-U 11111111-2222-3333-4444-555555555555 \
		"$image" "$fixture_dir"
}

build_image "$artifact_dir/nested.erofs"
build_image "$artifact_dir/nested-repeat.erofs"
cmp "$artifact_dir/nested.erofs" "$artifact_dir/nested-repeat.erofs"

ARTIFACT_DIR="$artifact_dir" python3 <<'PY'
import hashlib
import os
import shutil
import struct

artifact_dir = os.environ["ARTIFACT_DIR"]
base_path = os.path.join(artifact_dir, "nested.erofs")


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 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 update_superblock_checksum(image):
    polynomial = 0x82F63B78
    checksum = 0xFFFFFFFF
    put_u32(image, 1024 + 4, 0)
    for byte in image[1024:4096]:
        checksum ^= byte
        for _ in range(8):
            checksum = (checksum >> 1) ^ (polynomial if checksum & 1 else 0)
    put_u32(image, 1024 + 4, checksum & 0xFFFFFFFF)


def inode_offset(nid):
    return nid << 5


def compact_inline_dir(image, nid):
    offset = inode_offset(nid)
    inode_format = u16(image, offset)
    layout = (inode_format >> 1) & 0x7
    assert (inode_format & 0x1) == 0
    assert layout == 2
    assert u16(image, offset + 2) == 0
    return offset, offset + 32, u32(image, offset + 8)


def dir_entries(image, data_offset, size):
    first_nameoff = u16(image, data_offset + 8)
    assert first_nameoff >= 12
    assert first_nameoff % 12 == 0
    assert first_nameoff < size
    count = first_nameoff // 12
    entries = []
    for index in range(count):
        entry_offset = data_offset + index * 12
        nid, nameoff, file_type, reserved = struct.unpack_from(
            "<QHBB", image, entry_offset
        )
        endoff = (
            u16(image, data_offset + (index + 1) * 12 + 8)
            if index + 1 < count
            else size
        )
        assert first_nameoff <= nameoff < endoff <= size
        slot = bytes(image[data_offset + nameoff : data_offset + endoff])
        name = slot.split(b"\0", 1)[0]
        assert name
        entries.append((nid, nameoff, file_type, reserved, name))
    return entries


with open(base_path, "rb") as source:
    base = bytearray(source.read())

assert u32(base, 1024) == 0xE0F5E1E2
assert base[1024 + 12] == 12
assert u16(base, 1024 + 14) == 36
root_nid = 36
root_inode, root_data, root_size = compact_inline_dir(base, root_nid)
root_entries = dir_entries(base, root_data, root_size)
assert [entry[4] for entry in root_entries] == [b".", b"..", b"alpha", b"wide"]

dot_omitted = bytearray(base)
removed = 13
new_size = root_size - removed
new_entries = root_entries[1:]
new_first_nameoff = len(new_entries) * 12
for index, entry in enumerate(new_entries):
    nid, nameoff, file_type, reserved, _ = entry
    struct.pack_into(
        "<QHBB",
        dot_omitted,
        root_data + index * 12,
        nid,
        nameoff - removed,
        file_type,
        reserved,
    )
names = b"".join(entry[4] for entry in new_entries)
dot_omitted[root_data + new_first_nameoff : root_data + new_size] = names
dot_omitted[root_data + new_size : root_data + root_size] = b"\0" * (
    root_size - new_size
)
put_u16(dot_omitted, root_inode, u16(dot_omitted, root_inode) | (1 << 4))
put_u32(dot_omitted, root_inode + 8, new_size)
assert [entry[4] for entry in dir_entries(dot_omitted, root_data, new_size)] == [
    b"..",
    b"alpha",
    b"wide",
]

short_block = bytearray(base)
put_u32(short_block, root_inode + 8, 8)

nonmonotonic = bytearray(base)
first_nameoff = u16(nonmonotonic, root_data + 8)
put_u16(nonmonotonic, root_data + 12 + 8, first_nameoff)

wide_nid = next(entry[0] for entry in root_entries if entry[4] == b"wide")
wide_inode = inode_offset(wide_nid)
assert ((u16(base, wide_inode) >> 1) & 0x7) == 2
wide_startblk = u32(base, wide_inode + 16)
wide_block = wide_startblk << 12
wide_first_nameoff = u16(base, wide_block + 8)
wide_count = wide_first_nameoff // 12
wide_last_nameoff = u16(base, wide_block + (wide_count - 1) * 12 + 8)
padding_nul = base.index(0, wide_block + wide_last_nameoff, wide_block + 4096)
assert padding_nul + 1 < wide_block + 4096
assert base[padding_nul + 1] == 0
bad_padding = bytearray(base)
bad_padding[padding_nul + 1] = ord("X")

outputs = {
    "dot-omitted.erofs": dot_omitted,
    "corrupt-short-block.erofs": short_block,
    "corrupt-nameoff.erofs": nonmonotonic,
    "corrupt-padding.erofs": bad_padding,
}
for filename, image in outputs.items():
    update_superblock_checksum(image)
    path = os.path.join(artifact_dir, filename)
    with open(path, "wb") as output:
        output.write(image)

with open(os.path.join(artifact_dir, "SHA256SUMS"), "w", encoding="ascii") as sums:
    for filename in sorted(["nested.erofs", "nested-repeat.erofs", *outputs]):
        path = os.path.join(artifact_dir, filename)
        with open(path, "rb") as image_file:
            digest = hashlib.sha256(image_file.read()).hexdigest()
        sums.write(f"{digest}  {filename}\n")
PY

cat "$artifact_dir/SHA256SUMS"