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

"""
Command-line interface for SPI flash image stripe / unstripe.

Implements functionality similar to the flash_stripe utility in the
QEMU source tree (util/flash-stripe.c).
"""

import sys
import os
import argparse

try:
    from spi_interleave import stripe, unstripe
except ModuleNotFoundError:
    # Relaunch the helper and look for nativepython3 in the same directory
    nativepython3 = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'nativepython3')
    args = [ nativepython3 ] + sys.argv
    env = os.environ
    os.execve(nativepython3, args, env)
    sys.exit(1)

def create_progress_callback(total_size: int, label: str = "Progress"):
    """Create a callback that displays current offset / total file size."""
    def progress_callback(offset: int) -> None:
        percentage = 100.0 * offset / total_size if total_size > 0 else 0
        print(f"\r{label}: {offset} / {total_size} bytes ({percentage:.1f}%)", 
              end='', file=sys.stderr)
        if offset >= total_size:
            print(file=sys.stderr)  # Newline at end
    return progress_callback


def main():
    parser = argparse.ArgumentParser(
        description='Stripe or unstripe flash images across multiple files'
    )
    parser.add_argument('--unstripe', action='store_true',
                        help='Unstripe mode (combine multiple files into one)')
    parser.add_argument('--be', action='store_true',
                        help='Use big-endian bit ordering')
    parser.add_argument('--byte-wise', action='store_true',
                        help='Byte-wise mode (no bit-level transformation)')
    parser.add_argument('--no-sparse', action='store_true',
                        help='Disable sparse file optimization (write all zeros)')
    parser.add_argument('files', nargs='+',
                        help='Files to process (first is input/output, rest are output/input)')
    
    # The name 'unstrip' is a common integration typo, so we check for both.
    #
    # Auto-enable unstripe if invoked name contains "_unstrip" or "_unstripe"
    invoked_name = os.path.basename(sys.argv[0])
    if '_unstrip' in invoked_name or '_unstripe' in invoked_name:
        parser.set_defaults(unstripe=True)

    # Auto-enable byte-wise mode if invoked name contains "_bw"
    if '_bw' in invoked_name:
        parser.set_defaults(byte_wise=True)

    args = parser.parse_args()
    
    if len(args.files) < 2:
        parser.error('At least 2 files are required')
    
    sparse = not args.no_sparse
    
    if args.unstripe:
        # First N-1 files are inputs, last is output
        input_files = args.files[:-1]
        output_file = args.files[-1]
        total_size = os.path.getsize(input_files[0]) if os.path.exists(input_files[0]) else 0
        callback = create_progress_callback(total_size, "Unstriping")
        result = unstripe(input_files, output_file, args.be, args.byte_wise, callback, sparse)
    else:
        # First file is input, rest are outputs
        input_file = args.files[0]
        output_files = args.files[1:]
        total_size = os.path.getsize(input_file) if os.path.exists(input_file) else 0
        callback = create_progress_callback(total_size, "Striping")
        result = stripe(input_file, output_files, args.be, args.byte_wise, callback, sparse)
    
    sys.exit(0 if result else 1)


if __name__ == '__main__':
    main()
