#!/usr/bin/env python3
"""
Boot diagnosis — reads the live device and answers "why won't it boot?".

Goes beyond partition-match: inspects the actual boot chain on the eMMC and
reports the likely blocker. Checks:
  1. GPT (USER@0)        — present? partition table valid? offsets == scatter?
  2. preloader (BOOT1)   — present (not erased)?
  3. lk_a / boot_a       — bootloader + kernel images present & valid magic?
  4. vbmeta_a            — AVB magic + is verification disabled?
  5. seccfg              — lock state (erased == unlocked == good)?
  6. metadata            — erased (so userdata re-encrypts) ?
  7. para (A/B control)  — which slot is active / is slot A bootable?

Run with device in BROM mode:
    sudo python3 diagnose_boot.py
"""
import os, sys, time, 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")
SCATTER = ("/Users/home/Downloads/Stock rom F21 pro v1.0.8/"
           "image/MT6761_Android_scatter.txt")
SECTOR = 512


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


def erased(data):
    """True if the block looks erased (all 0x00 or all 0xFF)."""
    if not data:
        return True
    s = set(data[:4096])
    return s <= {0x00} or s <= {0xFF}


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))
            return True
        except Exception as e:
            log("warn", f"connect: {e}")
            mtk.disconnect(); time.sleep(0.5)
    return False


def parse_gpt(data):
    """Parse an EFI GPT from a USER@0 dump. Returns list of (name, byte_off)."""
    if data[SECTOR:SECTOR + 8] != b"EFI PART":
        return None
    hdr = data[SECTOR:SECTOR * 2]
    part_lba = struct.unpack_from("<Q", hdr, 72)[0]
    num = struct.unpack_from("<I", hdr, 80)[0]
    esize = struct.unpack_from("<I", hdr, 84)[0]
    out = []
    base = part_lba * SECTOR
    for i in range(min(num, 256)):
        off = base + i * esize
        if off + 128 > len(data):
            break
        ent = data[off:off + esize]
        type_guid = ent[0:16]
        if type_guid == b"\x00" * 16:
            continue
        first_lba = struct.unpack_from("<Q", ent, 32)[0]
        name = ent[56:128].decode("utf-16-le", "ignore").split("\x00")[0]
        out.append((name, first_lba * SECTOR))
    return out


def main():
    print("=" * 70)
    print("  BOOT DIAGNOSIS — what's on the device vs the stock")
    print("=" * 70)
    sf = parse_scatter(SCATTER)
    auto_load_roms(sf)
    smap = {p.name.lower(): p for p in sf.partitions}

    mtk = MTKDevice(log_cb=lambda l, m: None)
    if not connect(mtk):
        log("FAIL", "could not connect"); return 2
    log("ok", "DA ready\n")

    problems = []

    def rd(addr, size=0x10000, parttype=8):
        try:
            return mtk.read_partition(addr, size, parttype=parttype)
        except Exception as e:
            log("warn", f"read @ {addr:#x} failed: {e}")
            return b""

    # ---- 1. GPT ----
    print("-- 1. PARTITION TABLE (GPT) --")
    gpt_raw = rd(0x0, 0x10000, 8)
    gpt = parse_gpt(gpt_raw)
    if not gpt:
        problems.append("NO valid GPT at USER@0 — lk cannot find any partition. "
                        "Device cannot boot. (pgpt was never written / was erased)")
        log("FAIL", "no 'EFI PART' signature — partition table MISSING")
    else:
        log("ok", f"GPT valid, {len(gpt)} partitions on device")
        # compare a few key partitions to the scatter
        gmap = {n.lower(): off for n, off in gpt}
        mism = []
        for key in ("boot_a", "lk_a", "super", "vbmeta_a", "tee_a", "metadata", "seccfg"):
            if key in smap and key in gmap:
                exp = smap[key].start_addr
                got = gmap[key]
                tag = "ok" if exp == got else "MISMATCH"
                if exp != got:
                    mism.append(key)
                print(f"     {key:<12} device={got:#x}  scatter={exp:#x}  [{tag}]")
        if mism:
            problems.append(f"GPT offsets differ from scatter for {mism} — the "
                            f"on-device partition table points elsewhere than "
                            f"where we wrote. lk reads wrong data → no boot.")

    # ---- 2. preloader (BOOT1 + BOOT2) ----
    print("\n-- 2. PRELOADER (BOOT1 & BOOT2) --")
    pre1 = rd(0x0, 0x10000, 1)   # BOOT1
    pre2 = rd(0x0, 0x10000, 2)   # BOOT2
    if erased(pre1):
        problems.append("preloader (BOOT1) is ERASED — BootROM has nothing to "
                        "run. Totally dead device.")
        log("FAIL", "BOOT1 preloader erased/empty")
    else:
        magic1 = b"EMMC_BOOT" in pre1[:0x200] or b"MMM" in pre1[:0x40]
        log("ok", f"BOOT1 preloader present ({'MTK header' if magic1 else 'no header?'})")
    if erased(pre2):
        problems.append("preloader (BOOT2) is ERASED/empty — scatter requires "
                        "EMMC_BOOT1_BOOT2 (both). If this eMMC boots from BOOT2, "
                        "the device is dead. The tool's BOOT2 write is disabled — "
                        "ENABLE it (write preloader to BOTH boot partitions).")
        log("FAIL", "BOOT2 preloader ERASED — scatter wants BOTH boot parts")
    elif pre2[:0x40] == pre1[:0x40]:
        log("ok", "BOOT2 preloader present and matches BOOT1")
    else:
        log("warn", "BOOT2 has data but differs from BOOT1 (old preloader?)")

    # ---- 3. lk + boot ----
    print("\n-- 3. BOOTLOADER (lk_a) + KERNEL (boot_a) --")
    for key, magic, label in (("lk_a", None, "lk"),
                              ("boot_a", b"ANDROID!", "boot")):
        if key in smap:
            d = rd(smap[key].start_addr, 0x10000, 8)
            if erased(d):
                problems.append(f"{key} is ERASED — missing {label}. No boot.")
                log("FAIL", f"{key} erased")
            elif magic and d[:len(magic)] != magic:
                problems.append(f"{key} has no '{magic.decode()}' magic — bad image.")
                log("FAIL", f"{key} present but wrong magic: {d[:8]!r}")
            else:
                log("ok", f"{key} present"
                    + (f" ('{magic.decode()}' ok)" if magic else ""))

    # ---- 4. vbmeta ----
    print("\n-- 4. AVB (vbmeta_a) --")
    if "vbmeta_a" in smap:
        vb = rd(smap["vbmeta_a"].start_addr, 0x1000, 8)
        if erased(vb):
            log("warn", "vbmeta_a erased (AVB will treat as unlocked — usually OK)")
        elif vb[:4] == b"AVB0":
            # flags at offset 123 (avb header): 0=enabled,1=verify-disabled,2=hashtree-disabled
            flags = vb[123] if len(vb) > 123 else 0
            log("ok", f"vbmeta AVB0 present, flags={flags}")
            if flags == 0:
                log("info", "  AVB verification ENABLED — must match seccfg lock state")
        else:
            log("warn", f"vbmeta present but no AVB0 magic: {vb[:8]!r}")

    # ---- 5. seccfg (lock state) ----
    print("\n-- 5. seccfg (lock state) --")
    if "seccfg" in smap:
        sc = rd(smap["seccfg"].start_addr, 0x1000, 8)
        if erased(sc):
            log("ok", "seccfg ERASED → device treated as UNLOCKED (good for boot)")
        else:
            log("warn", f"seccfg has data: {sc[:16].hex()} — may enforce lock/AVB. "
                "If LOCKED + vbmeta mismatch → lk refuses boot.")
            problems.append("seccfg still has data after format — if it says LOCKED "
                            "and AVB doesn't match, lk blocks boot.")

    # ---- 6. metadata ----
    print("\n-- 6. metadata (userdata encryption) --")
    if "metadata" in smap:
        md = rd(smap["metadata"].start_addr, 0x1000, 8)
        if erased(md):
            log("ok", "metadata ERASED → userdata will re-encrypt fresh (good)")
        else:
            log("warn", "metadata still has data — stale encryption may block "
                "userdata mount (bootloop after logo).")
            problems.append("metadata not erased — stale userdata encryption.")

    # ---- 7. A/B slot control (para) ----
    print("\n-- 7. A/B ACTIVE SLOT (para / bootloader_control) --")
    if "para" in smap:
        pa = rd(smap["para"].start_addr, 0x10000, 8)
        idx = pa.find(struct.pack("<I", 0x42414342))  # 'BCAB' bootctrl magic
        if idx == -1:
            log("warn", "no bootloader_control magic in para — slot metadata "
                "missing/erased. lk may default to slot A (usually OK) or refuse.")
            problems.append("A/B slot metadata (bootloader_control) not found — "
                            "if lk needs an explicit active slot it won't boot. "
                            "May need to set active slot = A.")
        else:
            slot_suffix = pa[idx-4:idx]
            nb = pa[idx+5] if len(pa) > idx+5 else 0
            log("ok", f"bootloader_control found @ {idx:#x}, suffix={slot_suffix!r}, "
                f"slots={nb}")
            # slot_info array right after the 8-byte header
            for s in range(min(nb, 2)):
                b = pa[idx+8+s] if len(pa) > idx+8+s else 0
                pri = b & 0x0F
                successful = (b >> 6) & 1
                log("info", f"  slot {chr(65+s)}: priority={pri} successful={successful}")

    # ---- VERDICT ----
    print("\n" + "=" * 70)
    if not problems:
        print("  VERDICT: no boot-blocker found on device. The written firmware,")
        print("  GPT, preloader and lock state all look bootable. If it still")
        print("  won't power on, suspect: charge level, the preloader↔hardware")
        print("  DRAM(EMI) match, or a hardware fault.")
    else:
        print(f"  VERDICT: {len(problems)} likely boot-blocker(s):\n")
        for i, p in enumerate(problems, 1):
            print(f"   {i}. {p}\n")
    print("=" * 70)
    return 0


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