"""
fix_mbr.py

이미지 파일의 MBR 파티션 테이블에서 1번과 2번 파티션의 start LBA와 섹터 수를 채워 새로운 이미지로 저장합니다.

기본 동작:
  - 1번 파티션 start LBA = 63
  - 2번 파티션 start LBA = 4306239
  - 1번 섹터 수 = 2번 start - 1번 start
  - 2번 섹터 수 = 디스크 총 섹터 수 - 2번 start

사용법:
  python fix_mbr.py 260607_ence.001
  python fix_mbr.py input.001 -o output.001 --p1 63 --p2 4306239 --p1-type 0x83 --p2-type 0x83
"""

from __future__ import annotations

import argparse
import shutil
import struct
import sys
from pathlib import Path


def update_image_stream(inp: Path, outp: Path, new_first_sector: bytes) -> None:
    """첫 섹터(new_first_sector)를 교체하고 나머지는 스트리밍으로 복사합니다."""
    with inp.open('rb') as fr, outp.open('wb') as fw:
        fw.write(new_first_sector)
        # fr is at pos 0; advance to 512
        fr.seek(512)
        shutil.copyfileobj(fr, fw)


def update_partition_entry(mbr: bytearray, index: int, start_lba: int, num_sectors: int, part_type: int = 0x83, boot_flag: int = 0x00) -> None:
    table_off = 0x1BE
    off = table_off + index * 16
    # boot flag
    mbr[off] = boot_flag & 0xFF
    # CHS start (leave zero)
    mbr[off + 1:off + 4] = b"\x00\x00\x00"
    # type
    mbr[off + 4] = part_type & 0xFF
    # CHS end (leave zero)
    mbr[off + 5:off + 8] = b"\x00\x00\x00"
    # start LBA
    mbr[off + 8:off + 12] = struct.pack('<I', start_lba)
    # number of sectors
    mbr[off + 12:off + 16] = struct.pack('<I', num_sectors)


def main(argv=None):
    p = argparse.ArgumentParser(description='Fix MBR partition starts and sizes')
    p.add_argument('input', help='입력 이미지 파일')
    p.add_argument('-o', '--output', help='수정된 이미지 출력 파일 (기본: <input>.fixed.001)')
    p.add_argument('--p1', type=int, default=63, help='1번 파티션 start LBA (default=63)')
    p.add_argument('--p2', type=int, default=4306239, help='2번 파티션 start LBA (default=4306239)')
    p.add_argument('--p1-type', type=lambda x: int(x, 0), default=0x83, help='1번 파티션 타입 (hex/dec, default=0x83)')
    p.add_argument('--p2-type', type=lambda x: int(x, 0), default=0x83, help='2번 파티션 타입 (hex/dec, default=0x83)')
    args = p.parse_args(argv)

    inp = Path(args.input)
    if not inp.exists():
        print('Input file not found:', inp, file=sys.stderr)
        return 2

    outp = Path(args.output) if args.output else inp.with_name(inp.stem + '.fixed' + inp.suffix)
    bakp = inp.with_suffix(inp.suffix + '.bak')

    # backup
    shutil.copy2(inp, bakp)
    print(f'Backup created: {bakp}')

    size = inp.stat().st_size
    if size < 512:
        print('Input image too small (<512 bytes)', file=sys.stderr)
        return 3
    total_sectors = size // 512

    p1_start = args.p1
    p2_start = args.p2
    if not (0 <= p1_start < total_sectors and 0 <= p2_start < total_sectors):
        print('Start LBA values out of range for this image', file=sys.stderr)
        return 4

    if p2_start <= p1_start:
        print('p2 start must be greater than p1 start', file=sys.stderr)
        return 5

    p1_sectors = p2_start - p1_start
    p2_sectors = total_sectors - p2_start

    # Read only first sector, modify it, then stream-copy remainder
    with inp.open('rb') as fr:
        first = bytearray(fr.read(512))
    # Update partition entries 0 and 1 in the first sector
    update_partition_entry(first, 0, p1_start, p1_sectors, part_type=args.p1_type, boot_flag=0x00)
    update_partition_entry(first, 1, p2_start, p2_sectors, part_type=args.p2_type, boot_flag=0x00)
    # Ensure boot signature exists
    first[510:512] = b"\x55\xAA"

    # Stream out
    update_image_stream(inp, outp, bytes(first))
    print(f'Wrote fixed image: {outp}')
    print(f'Partition 1: start={p1_start}, sectors={p1_sectors}')
    print(f'Partition 2: start={p2_start}, sectors={p2_sectors}')
    return 0


if __name__ == '__main__':
    raise SystemExit(main())