#!/usr/bin/env python3
"""
Shift all data on a block device one chunk toward the end,
freeing the first chunk for an MBR + partition table.

Works backward from the end to avoid overwriting unread data.
"""

import argparse
import os
import sys

CHUNK = 1024 * 1024  # 1 MiB


def get_device_size(fd):
    """Get block device or file size via seek."""
    pos = os.lseek(fd, 0, os.SEEK_END)
    os.lseek(fd, 0, os.SEEK_SET)
    return pos


def main():
    ap = argparse.ArgumentParser(
        description="Shift block device contents by 1 MiB toward the end"
    )
    ap.add_argument("device", help="Block device or raw image file")
    ap.add_argument(
        "--chunk-size", type=int, default=CHUNK,
        help="Chunk size in bytes (default: 1 MiB)"
    )
    ap.add_argument(
        "--dry-run", action="store_true",
        help="Show what would happen without writing"
    )
    args = ap.parse_args()

    chunk = args.chunk_size

    fd = os.open(args.device, os.O_RDONLY if args.dry_run else os.O_RDWR | os.O_SYNC)
    dev_size = get_device_size(fd)

    if dev_size % chunk != 0:
        print(
            f"WARNING: device size {dev_size} is not evenly divisible by "
            f"chunk size {chunk}.  {dev_size % chunk} trailing bytes will "
            f"NOT be moved.",
            file=sys.stderr,
        )

    n = dev_size // chunk  # total number of whole chunks on the device

    if n < 2:
        print("Device too small (need at least 2 chunks).", file=sys.stderr)
        os.close(fd)
        sys.exit(1)

    # Data currently occupies chunks 0 .. N-2.
    # Chunk N-1 is the headroom we added when we created the larger device.
    # We shift every chunk forward by one position:
    #   chunk[N-2] -> chunk[N-1]   (into the free headroom)
    #   chunk[N-3] -> chunk[N-2]
    #   ...
    #   chunk[0]   -> chunk[1]
    # Then zero chunk[0] so we can later write an MBR + GRUB there.

    last = n - 2  # highest chunk index that contains real data

    print(f"Device : {args.device}")
    print(f"Size   : {dev_size:,} bytes  ({n} x {chunk // 1024} KiB chunks)")
    print(f"Shifting chunks 0..{last} → 1..{last + 1}")
    if args.dry_run:
        print("(dry run — no writes)")
    print()

    for i in range(last, -1, -1):            # N-2, N-3, … 1, 0
        src_off = i * chunk
        dst_off = (i + 1) * chunk

        if not args.dry_run:
            os.lseek(fd, src_off, os.SEEK_SET)
            data = os.read(fd, chunk)
            if len(data) != chunk:
                print(f"Short read at chunk {i} (got {len(data)})", file=sys.stderr)
                os.close(fd)
                sys.exit(1)

            os.lseek(fd, dst_off, os.SEEK_SET)
            written = os.write(fd, data)
            if written != chunk:
                print(f"Short write at chunk {i+1} (wrote {written})", file=sys.stderr)
                os.close(fd)
                sys.exit(1)

        # Progress every 256 chunks (~256 MiB) or on first/last
        if i == last or i == 0 or i % 256 == 0:
            done = last - i + 1
            pct = done * 100.0 / (last + 1)
            print(f"\r  [{pct:5.1f}%]  chunk {i:>8} → {i+1:<8}", end="", flush=True)

    print()

    # Zero out chunk 0
    print("Zeroing chunk 0 …")
    if not args.dry_run:
        os.lseek(fd, 0, os.SEEK_SET)
        os.write(fd, b"\x00" * chunk)

    os.close(fd)
    print("Done.")


if __name__ == "__main__":
    main()


##  # Dry run first — shows the plan, writes nothing
##  python3 block_shifter.py /dev/xen-vg/vm-root-new --dry-run

##  # Real run
##  python3 block_shifter.py /dev/xen-vg/vm-root-new

## Process
## if a 40GB/40960M ext4 raw disk:
#  1. fsck.ext4 -f /dev/vdb
#  2. resize2fs -p /dev/vdb  39G  or 40950M
#  3. python3 block_shifter.py /dev/vdb
#