#!/usr/bin/env python3
"""Validate a customer_fs LittleFS image for the EG800AKCN_11LG kernel."""

from __future__ import annotations

import argparse
import hashlib
import struct
import sys
from pathlib import Path


DEFAULT_SIZE = 589824
DEFAULT_DISK_VERSION = 0x00020000
DEFAULT_BLOCK_SIZE = 4096
DEFAULT_BLOCK_COUNT = 144
SUPERBLOCK_MAGIC_OFFSETS = (8, DEFAULT_BLOCK_SIZE + 8)


def fail(message: str) -> None:
    raise SystemExit(f"customer_fs verification failed: {message}")


def source_files(source: Path) -> list[Path]:
    if not source.is_dir():
        fail(f"source directory does not exist: {source}")

    entries = sorted(source.iterdir(), key=lambda path: path.name.encode("utf-8"))
    if not entries:
        fail(f"source directory is empty: {source}")

    invalid = [entry.name for entry in entries if not entry.is_file() or entry.name.startswith(".")]
    if invalid:
        fail(
            "source must contain only visible root-level files; invalid entries: "
            + ", ".join(invalid)
        )
    return entries


def verify_superblock(
    image: bytes,
    magic_offset: int,
    disk_version: int,
    block_size: int,
    block_count: int,
) -> None:
    if image[magic_offset : magic_offset + 8] != b"littlefs":
        fail(f"LittleFS magic is missing at 0x{magic_offset:X}")

    actual_version, actual_block_size, actual_block_count = struct.unpack_from(
        "<III", image, magic_offset + 12
    )
    if actual_version != disk_version:
        fail(
            f"disk version at 0x{magic_offset:X} is 0x{actual_version:08X}; "
            f"expected 0x{disk_version:08X}"
        )
    if actual_block_size != block_size:
        fail(
            f"block size at 0x{magic_offset:X} is {actual_block_size}; "
            f"expected {block_size}"
        )
    if actual_block_count != block_count:
        fail(
            f"block count at 0x{magic_offset:X} is {actual_block_count}; "
            f"expected {block_count}"
        )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", required=True, type=Path)
    parser.add_argument("--source", type=Path, help="also check root-level source filenames")
    parser.add_argument("--size", type=int, default=DEFAULT_SIZE)
    parser.add_argument(
        "--disk-version", type=lambda value: int(value, 0), default=DEFAULT_DISK_VERSION
    )
    parser.add_argument("--block-size", type=int, default=DEFAULT_BLOCK_SIZE)
    parser.add_argument("--block-count", type=int, default=DEFAULT_BLOCK_COUNT)
    args = parser.parse_args()

    if not args.image.is_file():
        fail(f"image does not exist: {args.image}")
    image = args.image.read_bytes()
    if len(image) != args.size:
        fail(f"image size is {len(image)}; expected {args.size}")
    if args.block_size * args.block_count != args.size:
        fail("block size multiplied by block count does not equal image size")

    for magic_offset in (8, args.block_size + 8):
        verify_superblock(
            image,
            magic_offset,
            args.disk_version,
            args.block_size,
            args.block_count,
        )

    files: list[Path] = []
    if args.source:
        files = source_files(args.source)
        missing = [path.name for path in files if path.name.encode("utf-8") not in image]
        if missing:
            fail("source filenames missing from image: " + ", ".join(missing))

    digest = hashlib.sha256(image).hexdigest()
    print("customer_fs verified")
    print(f"  image:        {args.image}")
    print(f"  size:         {len(image)}")
    print(f"  disk_version: 0x{args.disk_version:08X}")
    print(f"  block_size:   {args.block_size}")
    print(f"  block_count:  {args.block_count}")
    if files:
        print(f"  source_files: {len(files)}")
    print(f"  sha256:       {digest}")
    return 0


if __name__ == "__main__":
    sys.exit(main())