#!/usr/bin/env python3
"""Build and self-check fixtures for TC157-TC160."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import shutil
import stat
import struct
import subprocess

from review_fixtures import (
    EROFS_INODE_COMPRESSED_FULL,
    ErofsImage,
    FixtureError,
    SUPER,
    align,
    normalize_times,
    run_mkfs,
    sha256,
)


FEATURE_INCOMPAT_48BIT = 0x80
EROFS_INODE_FLAT_PLAIN = 0
EROFS_I_DOT_OMITTED = 1 << 4
Z_EROFS_ADVISE_EXTENTS = 0x1
OFF_MAX = (1 << 63) - 1


def map_header_offset(inode) -> int:
    return align(inode.offset + inode.inode_size + inode.xattr_size, 8)


def record_size(advise: int) -> int:
    return 4 << ((advise >> 1) & 3)


def binary_search_indices(lstarts: list[int], logical: int) -> list[int]:
    visited: list[int] = []
    left = 0
    right = len(lstarts)
    while left < right:
        middle = left + (right - left) // 2
        visited.append(middle)
        if lstarts[middle] > logical:
            right = middle
        else:
            left = middle + 1
            if lstarts[middle] == logical:
                right = min(left + 1, right)
    return visited


def convert_extent_table(
    base: ErofsImage,
    inode,
    recsz: int,
    lstarts: list[int],
    payload: int,
) -> tuple[ErofsImage, int, int]:
    image = base.clone()
    header = map_header_offset(inode)
    table = align(header + 8, recsz)
    table_end = table + recsz * len(lstarts)
    if table_end > len(image.data):
        raise FixtureError("explicit extent table exceeds the base image")
    root = image.inode(image.root_nid)
    if not (
        table_end <= root.offset
        or root.offset + root.inode_size <= header
    ):
        raise FixtureError("explicit extent conversion overlaps the root inode")

    advise = Z_EROFS_ADVISE_EXTENTS | ({16: 2, 32: 3}[recsz] << 1)
    struct.pack_into("<IHH", image.data, header, len(lstarts), advise, 0)
    for index, lstart in enumerate(lstarts):
        offset = table + index * recsz
        image.data[offset : offset + recsz] = b"\0" * recsz
        struct.pack_into(
            "<IIII",
            image.data,
            offset,
            image.block_size,
            payload & 0xFFFFFFFF,
            payload >> 32,
            lstart & 0xFFFFFFFF,
        )
        if recsz == 32:
            struct.pack_into("<I", image.data, offset + 16, lstart >> 32)
    image.update_checksum()
    return image, header, payload


def verify_extent_fixture(
    image: ErofsImage,
    inode,
    recsz: int,
    expected: list[int],
    kind: str,
    logical_probe: int,
) -> tuple[int, list[int]]:
    header = map_header_offset(inode)
    count = image.u32(header) | (image.u16(header + 6) << 32)
    advise = image.u16(header + 4)
    if count != len(expected) or record_size(advise) != recsz:
        raise FixtureError("explicit extent header self-check failed")
    table = align(header + 8, recsz)
    observed: list[int] = []
    for index in range(count):
        offset = table + index * recsz
        value = image.u32(offset + 12)
        if recsz == 32:
            value |= image.u32(offset + 16) << 32
        observed.append(value)
    if observed != expected or any(value >= inode.size for value in observed):
        raise FixtureError("explicit extent lstart self-check failed")
    violations = [
        index
        for index in range(1, len(observed))
        if observed[index] <= observed[index - 1]
    ]
    if len(violations) != 1:
        raise FixtureError("fixture must contain exactly one ordering violation")
    violation = violations[0]
    if kind == "duplicate" and observed[violation] != observed[violation - 1]:
        raise FixtureError("duplicate fixture is not a duplicate")
    if kind != "duplicate" and observed[violation] >= observed[violation - 1]:
        raise FixtureError("descending fixture is not descending")
    visited = binary_search_indices(observed, logical_probe)
    if kind == "cross-branch" and violation - 1 in visited:
        raise FixtureError("cross-branch violation is visible to the old search")
    return table, visited


def make_extent_fixtures(output: Path, source: Path) -> tuple[list[str], dict]:
    target_source = source / "extent.bin"
    target_source.write_bytes(b"extent-ordering-payload\n" * 65536)
    normalize_times(source)
    base_path = output / ".extent-base.erofs"
    run_mkfs(
        source,
        base_path,
        "77777777-1111-4222-8333-000000000157",
        "-E",
        "legacy-compress,force-inode-extended",
        "-z",
        "lz4",
        "-C4096",
    )
    base = ErofsImage.load(base_path)
    inode = base.resolve_root_entry("extent.bin")
    if inode.inode_size != 64 or inode.layout != EROFS_INODE_COMPRESSED_FULL:
        raise FixtureError("extent target is not extended COMPRESSED_FULL")
    header = map_header_offset(inode)
    payload = base.u32(header + 20) << base.block_bits
    if payload == 0 or payload + base.block_size > header:
        raise FixtureError("original compressed payload is not isolated from metadata")

    cases = {
        "descending": ([0, 4096, 12288, 8192], 8192),
        "duplicate": ([0, 4096, 4096, 12288], 4096),
        "cross-branch": ([0, 8192, 4096, 12288], 4096),
    }
    evidence: list[str] = []
    manifest_cases: dict[str, object] = {}
    for recsz in (16, 32):
        for kind, (lstarts, logical_probe) in cases.items():
            image, header, payload = convert_extent_table(
                base, inode, recsz, lstarts, payload
            )
            table, visited = verify_extent_fixture(
                image, inode, recsz, lstarts, kind, logical_probe
            )
            name = f"extent-{recsz}-{kind}.erofs"
            path = output / name
            image.save(path)
            evidence.append(
                "extent-order "
                f"image={name} nid={inode.nid} recsize={recsz} "
                f"header={header} table={table} lstarts={lstarts} "
                f"probe={logical_probe} old_search_indices={visited} "
                f"payload_offset={payload} violation={kind}"
            )
            manifest_cases[name] = {
                "path": "/extent.bin",
                "nid": inode.nid,
                "record_size": recsz,
                "header_offset": header,
                "table_offset": table,
                "lstarts": lstarts,
                "logical_probe": logical_probe,
                "old_search_indices": visited,
                "payload_offset": payload,
                "payload_length": image.block_size,
                "expected_errno": 97,
            }
    base_path.unlink()
    return evidence, {"cases": manifest_cases}


def make_blocks_fixture(output: Path, source: Path) -> tuple[str, dict]:
    target_source = source / "compressed-blocks.bin"
    target_source.write_bytes(bytes(range(256)) * 4096)
    normalize_times(source)
    base_path = output / ".blocks-base.erofs"
    run_mkfs(
        source,
        base_path,
        "77777777-1111-4222-8333-000000000158",
        "-E",
        "legacy-compress,force-inode-extended",
        "-z",
        "lz4",
        "-C4096",
    )
    image = ErofsImage.load(base_path)
    inode = image.resolve_root_entry("compressed-blocks.bin")
    if inode.inode_size != 64 or inode.layout != EROFS_INODE_COMPRESSED_FULL:
        raise FixtureError("compressed block target has the wrong inode layout")
    root_nid = image.root_nid
    inode_blocks_lo = image.u32(inode.offset + 16)
    if inode_blocks_lo == 0:
        raise FixtureError("compressed target has zero blocks_lo")
    super_blocks_lo = image.u32(SUPER + 36)
    image.put_u64(SUPER + 112, root_nid)
    image.put_u32(
        SUPER + 80, image.u32(SUPER + 80) | FEATURE_INCOMPAT_48BIT
    )
    image.put_u16(SUPER + 14, 1)
    image.put_u16(inode.offset + 6, 1)
    image.update_checksum()
    path = output / "compressed-blocks-hi.erofs"
    image.save(path)
    base_path.unlink()

    data_blocks = (1 << 32) | inode_blocks_lo
    va_bytes = data_blocks << image.block_bits
    st_blocks = va_bytes // 512
    provider_blocks = (1 << 32) | super_blocks_lo
    provider_size = provider_blocks << image.block_bits
    if image.root_nid != root_nid or image.u16(inode.offset + 6) != 1:
        raise FixtureError("48-bit inode/superblock self-check failed")
    if provider_size <= 16 * 1024**4:
        raise FixtureError("sparse provider is not larger than 16 TiB")
    evidence = (
        "compressed-blocks-hi "
        f"nid={inode.nid} inode_offset={inode.offset} blocks_lo={inode_blocks_lo} "
        f"blocks_hi=1 data_blocks={data_blocks} block_size={image.block_size} "
        f"va_bytes={va_bytes} st_blocks={st_blocks} "
        f"provider_blocks={provider_blocks} provider_size={provider_size}"
    )
    manifest = {
        "image": path.name,
        "path": "/compressed-blocks.bin",
        "nid": inode.nid,
        "inode_offset": inode.offset,
        "blocks_lo": inode_blocks_lo,
        "blocks_hi": 1,
        "data_blocks": data_blocks,
        "block_size": image.block_size,
        "va_bytes": va_bytes,
        "st_blocks": st_blocks,
        "provider_size": provider_size,
        "seed_size": len(image.data),
    }
    return evidence, manifest


def make_special_fixture(output: Path, source: Path) -> tuple[str, dict]:
    fifo = source / "special.fifo"
    os.mkfifo(fifo, 0o640)
    normalize_times(source)
    path = output / "special-setattr.erofs"
    run_mkfs(
        source,
        path,
        "77777777-1111-4222-8333-000000000159",
        "-E",
        "force-inode-extended",
    )
    image = ErofsImage.load(path)
    inode = image.resolve_root_entry("special.fifo")
    mode = image.u16(inode.offset + 4)
    if not stat.S_ISFIFO(mode):
        raise FixtureError("special setattr target is not a FIFO")
    evidence = (
        f"special-setattr nid={inode.nid} inode_offset={inode.offset} "
        f"mode={mode:#o} type=fifo"
    )
    return evidence, {
        "image": path.name,
        "path": "/special.fifo",
        "nid": inode.nid,
        "mode": mode,
    }


def make_offmax_fixture(output: Path, source: Path) -> tuple[str, dict]:
    directory = source / "offmax-dir"
    directory.mkdir()
    for index in range(384):
        (directory / f"entry-{index:03d}").write_text(
            f"entry {index}\n", encoding="ascii"
        )
    normalize_times(source)
    base_path = output / ".offmax-base.erofs"
    run_mkfs(
        source,
        base_path,
        "77777777-1111-4222-8333-000000000160",
        "-E",
        "force-inode-extended",
    )
    image = ErofsImage.load(base_path)
    inode = image.resolve_root_entry("offmax-dir")
    mode = image.u16(inode.offset + 4)
    if inode.inode_size != 64 or not stat.S_ISDIR(mode) or inode.start_block == 0:
        raise FixtureError("OFF_MAX target is not an extended backed directory")
    inode_format = image.u16(inode.offset)
    inode_format &= ~0x0E
    inode_format |= EROFS_I_DOT_OMITTED
    image.put_u16(inode.offset, inode_format)
    image.put_u64(inode.offset + 8, OFF_MAX)
    image.update_checksum()
    path = output / "dot-omitted-offmax.erofs"
    image.save(path)
    base_path.unlink()
    mutated = image.resolve_root_entry("offmax-dir")
    if (
        mutated.inode_size != 64
        or mutated.layout != EROFS_INODE_FLAT_PLAIN
        or mutated.size != OFF_MAX
        or not image.u16(mutated.offset) & EROFS_I_DOT_OMITTED
    ):
        raise FixtureError("OFF_MAX directory field self-check failed")
    evidence = (
        f"dot-omitted-offmax nid={mutated.nid} inode_offset={mutated.offset} "
        f"format={image.u16(mutated.offset):#x} size={mutated.size} "
        f"start_block={mutated.start_block} expected_errno=97"
    )
    return evidence, {
        "image": path.name,
        "path": "/offmax-dir",
        "nid": mutated.nid,
        "inode_offset": mutated.offset,
        "size": mutated.size,
        "expected_errno": 97,
    }


def make_fixtures(output: Path) -> None:
    if shutil.which("mkfs.erofs") is None:
        raise FixtureError("mkfs.erofs is required")
    if output.exists() and any(output.iterdir()):
        raise FixtureError(f"output directory is not empty: {output}")
    output.mkdir(parents=True, exist_ok=True)
    sources = output / ".sources"
    extent_source = sources / "extent"
    blocks_source = sources / "blocks"
    special_source = sources / "special"
    offmax_source = sources / "offmax"
    for source in (extent_source, blocks_source, special_source, offmax_source):
        source.mkdir(parents=True)

    evidence, extent_manifest = make_extent_fixtures(output, extent_source)
    blocks_evidence, blocks_manifest = make_blocks_fixture(output, blocks_source)
    special_evidence, special_manifest = make_special_fixture(output, special_source)
    offmax_evidence, offmax_manifest = make_offmax_fixture(output, offmax_source)
    evidence.extend([blocks_evidence, special_evidence, offmax_evidence])
    manifest = {
        "extent_order": extent_manifest,
        "compressed_blocks_hi": blocks_manifest,
        "special_setattr": special_manifest,
        "dot_omitted_offmax": offmax_manifest,
    }
    shutil.rmtree(sources)
    (output / "fixture-evidence.txt").write_text(
        "\n".join(evidence) + "\n", encoding="ascii"
    )
    (output / "fixture-manifest.json").write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="ascii"
    )
    with (output / "SHA256SUMS").open("w", encoding="ascii") as sums:
        for path in sorted(output.glob("*.erofs")):
            sums.write(f"{sha256(path)}  {path.name}\n")
    print((output / "fixture-evidence.txt").read_text(encoding="ascii"), end="")
    print((output / "SHA256SUMS").read_text(encoding="ascii"), end="")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    make_parser = subparsers.add_parser("make")
    make_parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.command == "make":
        make_fixtures(args.output)


if __name__ == "__main__":
    try:
        main()
    except (FixtureError, OSError, subprocess.CalledProcessError) as error:
        raise SystemExit(f"final_review_fixtures.py: {error}") from error