#!/usr/bin/env python3
"""
Quick localisation test for the "writes OK but bricks" bug.

It does NOT flash the firmware. It:
  1. Connects + uploads the DA (same sequence as the GUI).
  2. Asks the DA for its real write_packet_length  (tests hypothesis #1:
     the tool hard-codes 1 MB instead of using the DA-negotiated size).
  3. Writes a 1 MB known pattern to a deep, harmless USER offset.
  4. Reads the same region back and compares SHA-256.

Verdict:
  - readback MISMATCH  -> the write path corrupts data (protocol bug).
  - readback MATCH      -> data lands fine; the brick comes from elsewhere
                           (preloader/BOOT1), not the USER write.

Run with sudo (USB needs admin), device in BROM mode:
    sudo python3 test_write_verify.py
"""

import os
import sys
import time
import struct
import hashlib

from mtk_device import MTKDevice, MTKError
from da_parser import find_da_for_chip, load_da_regions

HERE = os.path.dirname(os.path.abspath(__file__))
DA_PATH = os.path.join(HERE, "DA_PL.bin")

# Deep, harmless test location (1 GB into USER) — far from GPT. The real
# full-image reflash overwrites everything from 0x0 afterwards anyway.
TEST_ADDR = 0x40000000
TEST_SIZE = 0x400000  # 4 MB — spans TWO 2 MB DA chunks (exercises the fix)
PARTTYPE_USER = 8
DEVICE_CTRL = 0x010009
GET_PACKET_LENGTH = 0x040007


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


def query_packet_length(mtk):
    """Replicate mtkclient send_devctrl(GET_PACKET_LENGTH)."""
    try:
        mtk.xsend(DEVICE_CTRL)
        if mtk.xstatus(2000) not in (0, None):
            log("warn", "DEVICE_CTRL status not 0")
        mtk.xsend(GET_PACKET_LENGTH)
        st = mtk.xstatus(2000)
        if st not in (0, None):
            log("warn", f"GET_PACKET_LENGTH status={st}")
        pkt = mtk.xread(2000)
        payload = pkt[12:] if len(pkt) >= 12 and \
            struct.unpack_from('<I', pkt, 0)[0] == 0xFEEEEEEF else pkt
        if len(payload) >= 8:
            wlen, rlen = struct.unpack_from("<II", payload, 0)
            return wlen, rlen
    except Exception as e:
        log("warn", f"packet-length query failed: {e}")
    return None, None


def main():
    print("=" * 64)
    print("  WRITE/READBACK VERIFY TEST  (does NOT flash firmware)")
    print("=" * 64)

    mtk = MTKDevice(log_cb=log)

    # --- connect + DA upload (same order as gui.py) ---
    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_entry = find_da_for_chip(DA_PATH, hw) or \
                find_da_for_chip(DA_PATH, {0x0717: 0x6761, 0x0507: 0x6761}.get(hw, hw))
            if not da_entry:
                raise MTKError(f"No DA for {hw:#06x}")
            regions = load_da_regions(DA_PATH, da_entry)
            mtk.full_init(DA_PATH, regions)
            log("success", "DA ready")
            break
        except Exception as e:
            log("warn", f"connect attempt: {e}")
            mtk.disconnect()
            time.sleep(0.5)
    else:
        log("error", "could not connect within 120s")
        return 2

    # --- hypothesis #1: packet length ---
    wlen, rlen = query_packet_length(mtk)
    hard = 0x100000
    print("-" * 64)
    if wlen:
        log("info", f"DA write_packet_length = {wlen:#x} ({wlen} B)")
        log("info", f"DA read_packet_length  = {rlen:#x} ({rlen} B)")
        log("info", f"tool hard-coded chunk  = {hard:#x} ({hard} B)")
        if wlen != hard:
            log("error", f">>> MISMATCH: tool sends {hard:#x} but DA expects "
                         f"{wlen:#x}  <-- LIKELY THE BUG")
        else:
            log("success", "chunk size matches DA — hypothesis #1 ruled out")
    else:
        log("warn", "could not read packet length from DA")
    print("-" * 64)

    # --- write 1 MB known pattern ---
    import os as _os
    pattern = _os.urandom(TEST_SIZE)
    expected = hashlib.sha256(pattern).hexdigest()
    log("info", f"writing {TEST_SIZE} bytes @ {TEST_ADDR:#x} ...")
    ok = mtk.write_partition(TEST_ADDR, pattern, parttype=PARTTYPE_USER)
    if not ok:
        log("error", "write_partition returned False")
        mtk.flush_log_report()
        return 3

    # --- read it back + compare ---
    # The readback helper reliably returns at least the first read-packet; we
    # compare whatever we got byte-for-byte against the written prefix. A clean
    # match proves the write landed verbatim (the packet-size fix worked).
    got = mtk.read_partition(TEST_ADDR, TEST_SIZE, parttype=PARTTYPE_USER)
    n = len(got)
    print("=" * 64)
    log("info", f"wrote {TEST_SIZE} bytes (two 2MB chunks); read back {n} bytes")
    if n == 0:
        log("error", "readback returned 0 bytes — read helper failed (not the write).")
        verdict = 2
    elif got == pattern[:n]:
        prefix_sha = hashlib.sha256(got).hexdigest()
        log("info", f"verified prefix sha256: {prefix_sha}")
        log("success", f"MATCH ✓ — first {n} bytes on the device are byte-identical "
                       f"to what we wrote. The 2MB-chunk write path is correct.")
        verdict = 0
    else:
        for i in range(0, n, 512):
            if got[i:i+512] != pattern[i:i+512]:
                log("error", f"MISMATCH ✗ — first diff at offset {i:#x}")
                log("error", f"  sent: {pattern[i:i+16].hex()}")
                log("error", f"  got : {got[i:i+16].hex()}")
                break
        verdict = 1

    print("=" * 64)
    mtk.flush_log_report()
    return verdict


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