#!/usr/bin/env python3
"""
Comprehensive post-flash verification for a full stock ROM.

For every partition in the scatter that has a ROM file, it reads the region
back FROM THE DEVICE and compares it (first 64 KB, or the whole file if
smaller) to the file we flashed. This proves, per-partition, whether the
flash actually landed — including the preloader in BOOT1.

Output: a table of partition / addr / region / MATCH|MISMATCH|SKIP, and a
verdict. A MISMATCH names the partition and first differing offset.

Run with device in BROM mode:
    sudo python3 verify_stock.py
    sudo python3 verify_stock.py "/path/to/MT6761_Android_scatter.txt"
"""
import os, sys, time, hashlib, struct
from mtk_device import MTKDevice, MTKError
from da_parser import find_da_for_chip, load_da_regions
from scatter_parser import parse_scatter, auto_load_roms, EmmcRegion

HERE = os.path.dirname(os.path.abspath(__file__))
DA_PATH = os.path.join(HERE, "DA_PL.bin")
DEFAULT_SCATTER = ("/Users/home/Downloads/Stock rom F21 pro v1.0.8/"
                   "image/MT6761_Android_scatter.txt")
CHUNK = 0x10000  # 64 KB read window


def log(level, msg):
    print(f"[{level:7}] {msg}")


def parttype_for(region):
    # preloader lives in BOOT1; everything else in the USER area
    return 1 if region == EmmcRegion.EMMC_BOOT1 else 8


def connect(mtk):
    deadline = time.time() + 120
    while time.time() < deadline:
        try:
            mtk.connect(timeout_sec=10)
            mtk.brom_handshake()
            mtk.get_hw_code(); mtk.get_hw_sw_ver()
            try: mtk.get_target_config()
            except Exception: pass
            hw = mtk.hw_code
            da = find_da_for_chip(DA_PATH, hw) or \
                find_da_for_chip(DA_PATH, {0x0717:0x6761,0x0507:0x6761}.get(hw,hw))
            if not da:
                raise MTKError(f"No DA for {hw:#06x}")
            mtk.full_init(DA_PATH, load_da_regions(DA_PATH, da))
            log("success", "DA ready")
            return True
        except Exception as e:
            log("warn", f"connect: {e}")
            mtk.disconnect(); time.sleep(0.5)
    return False


def main():
    scatter_path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_SCATTER
    print("=" * 72)
    print("  STOCK ROM VERIFY — device readback vs flashed files")
    print(f"  scatter: {scatter_path}")
    print("=" * 72)
    if not os.path.isfile(scatter_path):
        log("error", "scatter not found"); return 2

    sf = parse_scatter(scatter_path)
    auto_load_roms(sf)
    targets = [p for p in sf.partitions
               if p.enabled and p.file_path and os.path.isfile(p.file_path)]
    if not targets:
        log("error", "no partitions with ROM files found"); return 2
    log("info", f"{len(targets)} partitions have files; verifying each...")

    mtk = MTKDevice(log_cb=lambda l, m: None)  # quiet DA chatter
    if not connect(mtk):
        log("error", "could not connect"); return 2

    rows = []
    for p in targets:
        pt = parttype_for(p.region)
        region = "BOOT1" if pt == 1 else "USER"
        fsize = os.path.getsize(p.file_path)
        n = min(fsize, CHUNK)
        with open(p.file_path, "rb") as f:
            expected = f.read(n)
        try:
            got = mtk.read_partition(p.start_addr, n, parttype=pt)
        except Exception as e:
            rows.append((p.name, p.start_addr, region, fsize, "ERR", str(e)[:24]))
            continue
        m = min(len(got), len(expected))
        if m == 0:
            rows.append((p.name, p.start_addr, region, fsize, "NOREAD", ""))
        elif got[:m] == expected[:m]:
            rows.append((p.name, p.start_addr, region, fsize, "MATCH", ""))
        else:
            off = next((i for i in range(0, m, 16)
                        if got[i:i+16] != expected[i:i+16]), 0)
            rows.append((p.name, p.start_addr, region, fsize, "MISMATCH",
                         f"@{off:#x}"))

    # ---- report ----
    print("=" * 72)
    print(f"  {'partition':<16}{'addr':>14} {'reg':<6}{'size':>12}  {'result':<9}note")
    print("  " + "-" * 68)
    ok = bad = 0
    for name, addr, region, size, result, note in rows:
        mark = "✓" if result == "MATCH" else "✗"
        if result == "MATCH": ok += 1
        else: bad += 1
        print(f"  {name:<16}{('%#x' % addr):>14} {region:<6}{size:>12}  "
              f"{mark} {result:<7}{note}")
    print("  " + "-" * 68)
    print("=" * 72)
    if bad == 0:
        log("success", f"ALL {ok} PARTITIONS MATCH — every flashed partition is "
                       f"byte-identical on the device. The flash fully succeeded.")
        log("info", "If it still won't boot, the cause is outside the flash data "
                    "(A/B active slot, eMMC boot-config, or hardware).")
    else:
        log("error", f"{bad} partition(s) FAILED ({ok} OK). The ones marked "
                     f"MISMATCH/NOREAD did not land correctly — reflash them.")
    return 0


if __name__ == "__main__":
    sys.exit(main())
