import argparse
import hashlib
import os
import sys

def sha256(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()

def read(path: str) -> bytes:
    with open(path, "rb") as f:
        return f.read()

def parse_hex(s: str) -> int:
    s = s.strip()
    return int(s, 16) if s.lower().startswith("0x") else int(s)

def parse_rcd(path: str):
    text = read(path).decode("utf-8", "replace").splitlines()
    expected_rom = None
    expected_disks = {}
    records = []
    # (optional) hold meta if you ever want to surface it later
    meta = {}

    for ln in text:
        s = ln.strip()
        if not s or s.startswith("#"):
            continue
        if s == "RCDv1":
            continue
        if s.startswith("ROM_SHA256"):
            # format: ROM_SHA256 <hash>
            parts = s.split()
            if len(parts) >= 2:
                expected_rom = parts[1]
            continue
        if s.startswith("DISK_SHA256"):
            # format: DISK_SHA256 <idx> <hash>
            parts = s.split()
            if len(parts) >= 3:
                try:
                    expected_disks[int(parts[1])] = parts[2]
                except ValueError:
                    pass
            continue
        if s.startswith("META"):
            # format: META key=value key=value ...
            # (ignored by the applier; parsed for future use)
            try:
                for kv in s.split()[1:]:
                    if "=" in kv:
                        k, v = kv.split("=", 1)
                        meta[k] = v
            except Exception:
                pass
            continue
        if s.startswith("COPY"):
            # format: COPY rom_off=0x... d=<n> disk_off=0x... len=0x...
            parts = dict(kv.split("=", 1) for kv in s.split()[1:])
            records.append((
                "COPY",
                parse_hex(parts["rom_off"]),
                int(parts["d"]),
                parse_hex(parts["disk_off"]),
                parse_hex(parts["len"]),
            ))
            continue
        if s.startswith("LITERAL"):
            # format: LITERAL rom_off=0x... data=<HEX>
            parts = dict(kv.split("=", 1) for kv in s.split()[1:])
            records.append((
                "LITERAL",
                parse_hex(parts["rom_off"]),
                bytes.fromhex(parts["data"]),
            ))
            continue
        if s.startswith("RLE"):
            # format: RLE rom_off=0x... byte=0x.. len=0x...
            parts = dict(kv.split("=", 1) for kv in s.split()[1:])
            records.append((
                "RLE",
                parse_hex(parts["rom_off"]),
                parse_hex(parts["byte"]),
                parse_hex(parts["len"]),
            ))
            continue
        # Be tolerant of future annotations: ignore unknown non-critical lines
        # (keeps backward/forward compatibility)
        # If you prefer strictness, raise here instead.
        # raise ValueError(f"Unknown line: {s}")
        continue

    return expected_rom, expected_disks, records  # meta is intentionally ignored

def apply_rcd(rcd_path, disk_paths, out_rom, verbose=False, strict=False, delete_on_mismatch=False):
    expected_rom, expected_disks, records = parse_rcd(rcd_path)
    if not records:
        print("ERROR: No records found in RCD; nothing to build.")
        sys.exit(1)

    disks = [read(p) for p in disk_paths]

    # Verify disk hashes if present; print side-by-side
    for i, h in expected_disks.items():
        actual = sha256(disks[i-1])
        print(f"Disk {i} SHA256 | expected: {h}")
        print(f"                 actual  : {actual}")
        if strict and actual != h:
            print(f"ERROR: Disk {i} SHA256 mismatch (strict mode).", file=sys.stderr)
            sys.exit(2)

    # Compute output size
    rom_size = 0
    for rec in records:
        if rec[0] == "COPY":
            _, roff, _, _, L = rec
            rom_size = max(rom_size, roff + L)
        elif rec[0] == "LITERAL":
            _, roff, data = rec
            rom_size = max(rom_size, roff + len(data))
        else:  # RLE
            _, roff, _, L = rec
            rom_size = max(rom_size, roff + L)

    if rom_size == 0:
        print("ERROR: Computed ROM size is 0; invalid RCD?", file=sys.stderr)
        sys.exit(1)

    out = bytearray(b"\x00" * rom_size)

    # Apply
    for rec in records:
        if rec[0] == "COPY":
            _, roff, di, doff, L = rec
            out[roff:roff+L] = disks[di-1][doff:doff+L]
            if verbose:
                print(f"COPY   @{roff:#08x} <= disk{di}@{doff:#08x} len={L:#x}")
        elif rec[0] == "LITERAL":
            _, roff, data = rec
            out[roff:roff+len(data)] = data
            if verbose:
                print(f"LITERAL@{roff:#08x} len={len(data):#x}")
        else:  # RLE
            _, roff, byte_val, L = rec
            out[roff:roff+L] = bytes([byte_val & 0xFF]) * L
            if verbose:
                print(f"RLE    @{roff:#08x} byte={byte_val & 0xFF:#04x} len={L:#x}")

    # Ensure output directory exists
    out_dir = os.path.dirname(os.path.abspath(out_rom))
    if out_dir and not os.path.exists(out_dir):
        os.makedirs(out_dir, exist_ok=True)

    # Write first, then verify
    with open(out_rom, "wb") as f:
        f.write(out)
    print(f"Wrote {out_rom} ({rom_size:,} bytes).")

    out_hash = sha256(out)
    print(f"ROM SHA256     | expected: {expected_rom if expected_rom else '(not in file)'}")
    print(f"                 actual  : {out_hash}")

    if expected_rom and out_hash != expected_rom:
        msg = "WARNING: ROM SHA256 mismatch."
        if strict:
            msg = "ERROR: ROM SHA256 mismatch (strict mode)."
        print(msg, file=sys.stderr)
        if delete_on_mismatch:
            try:
                os.remove(out_rom)
                print("Output file deleted due to mismatch (--delete-on-mismatch).", file=sys.stderr)
            except OSError:
                pass
        if strict:
            sys.exit(3)

def main():
    print("RCD Apply tool v2.3 - (c) Max Iwamoto SEP.06.2025")
    ap = argparse.ArgumentParser(description="Apply RCD to rebuild ROM from 8 patched DSKs (RLE-aware)")
    ap.add_argument("rcd", help="build.rcd")
    ap.add_argument("disks", nargs=8, help="disk1.dsk ... disk8.dsk")
    ap.add_argument("-o", "--out", default="reconstructed.rom", help="output ROM file")
    ap.add_argument("-v", "--verbose", action="store_true", help="print each record applied")
    ap.add_argument("--strict", action="store_true", help="treat any SHA mismatch as an error exit")
    ap.add_argument("--delete-on-mismatch", action="store_true", help="delete output if ROM SHA mismatches")
    args = ap.parse_args()
    apply_rcd(args.rcd, args.disks, args.out, verbose=args.verbose,
              strict=args.strict, delete_on_mismatch=args.delete_on_mismatch)

if __name__ == "__main__":
    main()
