#!/usr/bin/env python3
"""
Check every boot-relevant partition on the device: present or ERASED, and
whether we even have a file for it. Finds the missing piece after the
bootloader — e.g. a firmware partition (vendor_boot/gz/...) that an aggressive
format erased and we never restored (no file), which stalls the kernel -> black
screen. Run with device in BROM mode:
    sudo python3 check_firmware.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")

# Partitions that must be NON-empty for the kernel/userspace to come up
# (slot A + shared). If any of these is erased AND we have no file, that's it.
BOOT_CRITICAL = {
    "preloader", "lk_a", "tee_a", "gz_a", "sspm_a", "scp_a", "spmfw_a",
    "md1img_a", "boot_a", "vendor_boot_a", "dtbo_a", "vbmeta_a",
    "vbmeta_system_a", "vbmeta_vendor_a", "super", "logo",
}


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


def erased(d):
    if not d:
        return True
    s = set(d[: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 main():
    print("=" * 74)
    print("  FIRMWARE INTEGRITY — which boot-critical partitions are EMPTY?")
    print("=" * 74)
    sf = parse_scatter(SCATTER)
    auto_load_roms(sf)
    have_file = {p.name.lower() for p in sf.partitions if p.file_path}

    mtk = MTKDevice(log_cb=lambda l, m: None)
    if not connect(mtk):
        log("FAIL", "could not connect"); return 2
    log("ok", "DA ready\n")
    print(f"  {'partition':<18}{'addr':>12}  {'on-device':<10}{'have file?':<11}note")
    print("  " + "-" * 68)

    missing = []
    for p in sf.partitions:
        name = p.name.lower()
        if name not in BOOT_CRITICAL:
            continue
        if p.start_addr >= 0xFFFF0000:
            continue
        parttype = 1 if p.region == EmmcRegion.EMMC_BOOT1 else 8
        d = mtk.read_partition(p.start_addr, 0x1000, parttype=parttype)
        empty = erased(d)
        hv = name in have_file
        state = "ERASED" if empty else "present"
        note = ""
        if empty and not hv:
            note = "<-- EMPTY + no file to restore = BOOT BLOCKER"
            missing.append(p.name)
        elif empty and hv:
            note = "(empty but we have a file — reflash it)"
        print(f"  {p.name:<18}{('%#x' % p.start_addr):>12}  "
              f"{state:<10}{('yes' if hv else 'NO'):<11}{note}")

    print("  " + "-" * 68)
    print("=" * 74)
    if missing:
        log("FAIL", f"BOOT BLOCKER FOUND — these are erased with no file to "
            f"restore: {missing}.")
        log("info", "These were wiped by a format and the stock package has no "
            "image for them. The kernel/userspace can't come up -> black screen. "
            "SP Flash Tool never erases them. Fix: don't format these, and/or "
            "obtain their images. (vendor_boot is required on GKI devices.)")
    else:
        log("ok", "All boot-critical partitions are present. The blocker is not "
            "an empty partition — look at vbmeta/AVB or seccfg next.")
    return 0


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