# Copyright (C) 2026, Advanced Micro Devices, Inc.  All rights reserved.
# Portions of this file consist of AI-generated content.
# SPDX-License-Identifier: MIT

"""
Stripe and unstripe flash images for QEMU SPI emulation.

Implements functionality similar to the flash_stripe utility in the
QEMU source tree (util/flash-stripe.c), which performs the same
bit/byte interleaving that SPI controllers apply when driving
parallel flash chips.

SPI controllers that drive parallel flash chips interleave data across
the chip-selects at the bit level (or byte level in some modes).  This
module splits a single flash image into per-chip files ("stripe") or
recombines them ("unstripe"), matching that hardware behaviour.

The core operation is a bit-level matrix transpose of N-byte groups:
treat N bytes as an 8-row x N-column bit matrix, then transpose so
column-wise bits become row-wise.  For performance the transpose is
encoded into 256-byte lookup tables and applied with bytes.translate().
"""

import os
import sys
from typing import Callable, List, Optional

# Bytes per chip-select file per iteration.
_CHUNK = 1 << 20  # 1 MiB

# Fill byte for padding short / unaligned inputs.
# 0xFF matches erased NOR flash state.
_FILL = 0xFF


# -------------------------------------------------------------------
# Bit-transpose lookup tables
# -------------------------------------------------------------------

def _make_tables(n: int, inverse: bool, big_end: bool) -> List[List[bytes]]:
    """Build 256-byte translation tables for the N-way bit transpose.

    Returns ``tbl`` where ``tbl[src_col][dst_col]`` is a 256-byte
    ``bytes`` object.  For input byte value *v* in source column
    *src_col*, ``tbl[src_col][dst_col][v]`` gives the bits that land
    in destination column *dst_col*.

    The mapping is derived from first principles: linearise the 8*N
    bits, then redistribute them via row/column transposition.
    """
    total = 8 * n

    # Walk every bit position through the transpose to build a
    # (src_col, dst_col) -> [(src_bit, dst_bit), ...] mapping.
    routes: dict = {}
    for pos in range(total):
        src_col = pos // 8
        src_bit = pos % 8
        dst_col = pos % n
        dst_bit = pos // n

        if big_end:
            src_bit = 7 - src_bit
            dst_bit = 7 - dst_bit

        if inverse:
            src_col, dst_col = dst_col, src_col
            src_bit, dst_bit = dst_bit, src_bit

        routes.setdefault((src_col, dst_col), []).append((src_bit, dst_bit))

    tbl = [[bytes(256) for _ in range(n)] for _ in range(n)]
    for (sc, dc), bits in routes.items():
        lut = bytearray(256)
        for v in range(256):
            acc = 0
            for sb, db in bits:
                if v & (1 << sb):
                    acc |= 1 << db
            lut[v] = acc
        tbl[sc][dc] = bytes(lut)
    return tbl


# -------------------------------------------------------------------
# Vectorised column transform
# -------------------------------------------------------------------

def _apply_transpose(columns: List[bytes],
                     tbl: List[List[bytes]]) -> List[bytes]:
    """Transpose *columns* using precomputed lookup tables.

    Each element of *columns* is a ``bytes`` of equal length (one per
    chip-select).  Returns a new list of the same shape.
    """
    n = len(columns)
    length = len(columns[0])
    out = []
    for dc in range(n):
        acc = int.from_bytes(columns[0].translate(tbl[0][dc]), 'little')
        for sc in range(1, n):
            acc |= int.from_bytes(
                columns[sc].translate(tbl[sc][dc]), 'little')
        out.append(acc.to_bytes(length, 'little'))
    return out


# -------------------------------------------------------------------
# Helpers
# -------------------------------------------------------------------

def _merge_columns(cols: List[bytes]) -> bytes:
    """Interleave equal-length byte columns into one buffer."""
    n = len(cols)
    length = len(cols[0])
    buf = bytearray(length * n)
    for i in range(n):
        buf[i::n] = cols[i]
    return bytes(buf)


def _write_maybe_sparse(fobj, data: bytes, sparse: bool) -> None:
    """Write *data*, or seek past it if *sparse* and all zero."""
    if sparse and not any(data):
        fobj.seek(len(data), os.SEEK_CUR)
    else:
        fobj.write(data)


def _safe_open_all(paths: List[str], mode: str) -> List:
    """Open every path in *paths*.  Close all on failure."""
    opened = []
    try:
        for p in paths:
            opened.append(open(p, mode))
    except Exception:
        for f in opened:
            f.close()
        raise
    return opened


# -------------------------------------------------------------------
# Public API
# -------------------------------------------------------------------

def stripe(input_file: str, output_files: List[str],
           be: bool = False, byte_wise: bool = False,
           status_callback: Optional[Callable[[int], None]] = None,
           sparse: bool = True) -> bool:
    """Split a single flash image across *len(output_files)* chip-select files.

    Args:
        input_file:      Path to the combined flash image.
        output_files:    One output path per chip-select.
        be:              Use big-endian (MSB-first) bit ordering.
        byte_wise:       Byte-level interleave only (skip bit transpose).
        status_callback: Called with current input offset after each chunk.
        sparse:          Skip writing all-zero regions (create sparse files).

    Returns:
        ``True`` on success, ``False`` on error.
    """
    n = len(output_files)
    if n < 2:
        print("Error: stripe needs at least 2 output files", file=sys.stderr)
        return False
    if not os.path.exists(input_file):
        print(f"Error: '{input_file}' not found", file=sys.stderr)
        return False

    total = os.path.getsize(input_file)
    tbl = None if byte_wise else _make_tables(n, False, be)
    read_sz = _CHUNK * n
    fill = bytes([_FILL])
    created: List[str] = []

    try:
        with open(input_file, 'rb') as inp:
            created = list(output_files)
            outs = _safe_open_all(output_files, 'wb')
            try:
                offset = 0
                warned = False
                while True:
                    raw = inp.read(read_sz)
                    if not raw:
                        break
                    size = len(raw)

                    # Pad to a multiple of n
                    if size % n:
                        if not warned:
                            print(f"Warning: input not a multiple of {n} "
                                  f"bytes, padding with 0x{_FILL:02X}",
                                  file=sys.stderr)
                            warned = True
                        pad = n - size % n
                        raw += fill * pad
                        size += pad

                    # Split into per-chip columns and optionally transpose
                    cols = [raw[i::n] for i in range(n)]
                    if not byte_wise:
                        cols = _apply_transpose(cols, tbl)

                    for i in range(n):
                        fi = (n - 1 - i) if (byte_wise and be) else i
                        _write_maybe_sparse(outs[fi], cols[i], sparse)

                    offset += size
                    if status_callback:
                        status_callback(min(offset, total))
            finally:
                expected = (total + n - 1) // n
                for f in outs:
                    f.truncate(expected)
                    f.close()

        return True

    except Exception as exc:
        print(f"Error striping file: {exc}", file=sys.stderr)
        for p in created:
            try:
                os.unlink(p)
            except OSError:
                pass
        return False


def unstripe(input_files: List[str], output_file: str,
             be: bool = False, byte_wise: bool = False,
             status_callback: Optional[Callable[[int], None]] = None,
             sparse: bool = True) -> bool:
    """Recombine per-chip-select files into a single flash image.

    Args:
        input_files:     One input path per chip-select.
        output_file:     Path for the combined output image.
        be:              Use big-endian (MSB-first) bit ordering.
        byte_wise:       Byte-level de-interleave only (skip bit transpose).
        status_callback: Called with current per-file offset after each chunk.
        sparse:          Skip writing all-zero regions (create sparse output).

    Returns:
        ``True`` on success, ``False`` on error.
    """
    n = len(input_files)
    if n < 2:
        print("Error: unstripe needs at least 2 input files", file=sys.stderr)
        return False
    for p in input_files:
        if not os.path.exists(p):
            print(f"Error: '{p}' not found", file=sys.stderr)
            return False

    sizes = [os.path.getsize(p) for p in input_files]
    longest = max(sizes)
    tbl = None if byte_wise else _make_tables(n, True, be)
    fill = bytes([_FILL])

    try:
        ins = _safe_open_all(input_files, 'rb')
        try:
            with open(output_file, 'wb') as out:
                offset = 0
                short_warned: set = set()

                while offset < longest:
                    want = min(_CHUNK, longest - offset)
                    cols = []
                    for i in range(n):
                        fi = (n - 1 - i) if (byte_wise and be) else i
                        chunk = ins[fi].read(want)
                        if len(chunk) < want:
                            if fi not in short_warned:
                                print(f"Warning: '{input_files[fi]}' is "
                                      f"shorter, padding with "
                                      f"0x{_FILL:02X}", file=sys.stderr)
                            short_warned.add(fi)
                            chunk += fill * (want - len(chunk))
                        cols.append(chunk)

                    # All-zero across every chip means the output is
                    # also all-zero (transpose preserves the zero vector).
                    if sparse and all(not any(c) for c in cols):
                        out.seek(want * n, os.SEEK_CUR)
                    else:
                        out_cols = cols if byte_wise \
                            else _apply_transpose(cols, tbl)
                        _write_maybe_sparse(out,
                                            _merge_columns(out_cols), sparse)

                    offset += want
                    if status_callback:
                        status_callback(offset)

                out.truncate(longest * n)

        finally:
            for f in ins:
                f.close()

        return True

    except Exception as exc:
        print(f"Error unstriping files: {exc}", file=sys.stderr)
        try:
            os.unlink(output_file)
        except OSError:
            pass
        return False
