#!/usr/bin/env python3
"""Generate a disk-2.0 LittleFS customer_fs image on macOS."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

from littlefs import LittleFS, UserContext


IMAGE_SIZE = 589824
BLOCK_SIZE = 4096
BLOCK_COUNT = 144
DISK_VERSION = 0x00020000


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


def collect_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 main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()

    files = collect_source_files(args.source)
    context = UserContext(buffsize=IMAGE_SIZE)
    fs = LittleFS(
        context=context,
        mount=False,
        block_size=BLOCK_SIZE,
        block_count=BLOCK_COUNT,
        read_size=16,
        prog_size=16,
        cache_size=256,
        lookahead_size=32,
        block_cycles=500,
        disk_version=DISK_VERSION,
    )
    fs.format()
    fs.mount()
    try:
        for source_file in files:
            with source_file.open("rb") as source_handle:
                with fs.open("/" + source_file.name, "wb") as image_handle:
                    while chunk := source_handle.read(64 * 1024):
                        image_handle.write(chunk)
    finally:
        fs.unmount()

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_bytes(context.buffer)
    print(f"generated {args.output} from {len(files)} files")
    return 0


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