"""Tests for the pure-Python firmware flasher (no hardware)."""
import asyncio
import struct
import pytest
from unittest.mock import AsyncMock

from pm3py.core.protocol import (
    Cmd,
    OLD_FRAME_SIZE, START_FLASH_MAGIC,
    FLASH_START, BOOTLOADER_END, FLASH_END_512K, BL_VERSION_1_0_0,
    DEVICE_INFO_FLAG_BOOTROM_PRESENT, DEVICE_INFO_FLAG_OSIMAGE_PRESENT,
    DEVICE_INFO_FLAG_CURRENT_MODE_BOOTROM, DEVICE_INFO_FLAG_CURRENT_MODE_OS,
)
from pm3py.core.transport import PM3Response, encode_old_frame, decode_old_frame
from pm3py.core import flash as flashmod
from pm3py.core.flash import (
    Flasher, FlashError, parse_elf_segments, build_blocks, resolve_image,
    find_firmware_src, build_firmware,
)


def _proc(returncode):
    return type("P", (), {"returncode": returncode})()


def _fake_src(tmp_path):
    (tmp_path / "Makefile").write_text("x")
    (tmp_path / "armsrc").mkdir()
    (tmp_path / "armsrc" / "Makefile").write_text("x")
    return tmp_path

OS_FLAGS = (DEVICE_INFO_FLAG_BOOTROM_PRESENT | DEVICE_INFO_FLAG_OSIMAGE_PRESENT
            | DEVICE_INFO_FLAG_CURRENT_MODE_OS)
BOOTROM_FLAGS = (DEVICE_INFO_FLAG_BOOTROM_PRESENT | DEVICE_INFO_FLAG_OSIMAGE_PRESENT
                 | DEVICE_INFO_FLAG_CURRENT_MODE_BOOTROM)


def run(coro):
    # Match the rest of the suite's idiom — asyncio.run() would close and null the
    # shared global loop, breaking later tests that use get_event_loop().
    return asyncio.get_event_loop().run_until_complete(coro)


def old_resp(cmd, arg0=0, arg1=0, arg2=0):
    return PM3Response(cmd=int(cmd), status=0, reason=0, ng=False,
                       data=b"\x00" * 512, oldarg=[arg0, arg1, arg2])


def make_elf(paddr, data, *, machine=40, etype=2, cls=1, endian=1):
    """Build a minimal single-PT_LOAD ELF32 image for tests."""
    ident = b"\x7fELF" + bytes([cls, endian, 1, 0]) + bytes(8)
    e_phoff = 52
    e_phentsize = 32
    header = ident + struct.pack(
        "<HHIIIIIHHHHHH",
        etype, machine, 1, paddr, e_phoff, 0, 0,
        52, e_phentsize, 1, 0, 0, 0,
    )
    data_off = e_phoff + e_phentsize
    phdr = struct.pack("<IIIIIIII",
                       1,           # PT_LOAD
                       data_off,    # p_offset
                       paddr, paddr,
                       len(data), len(data),
                       5, 0x10000)
    return header + phdr + data


# --- OLD frame encoding ---

def test_old_frame_size_and_roundtrip():
    frame = encode_old_frame(Cmd.FINISH_WRITE, 0x102000, 1, 2, b"\xAB\xCD")
    assert len(frame) == OLD_FRAME_SIZE == 544
    d = decode_old_frame(frame)
    assert d["cmd"] == Cmd.FINISH_WRITE
    assert d["oldarg"] == [0x102000, 1, 2]
    assert d["data"][:2] == b"\xAB\xCD"
    assert d["data"][2:] == b"\x00" * 510  # zero-padded to 512


def test_old_frame_rejects_oversized_data():
    with pytest.raises(ValueError):
        encode_old_frame(Cmd.FINISH_WRITE, data=b"\x00" * 513)


# --- ELF parsing ---

def test_parse_elf_single_segment():
    segs = parse_elf_segments(make_elf(0x102000, b"hello world"))
    assert segs == [(0x102000, b"hello world")]


def test_parse_elf_rejects_bad_magic():
    with pytest.raises(FlashError):
        parse_elf_segments(b"NOTELF" + bytes(100))


def test_parse_elf_rejects_non_arm():
    with pytest.raises(FlashError):
        parse_elf_segments(make_elf(0x102000, b"x", machine=3))  # EM_386


def test_parse_elf_rejects_non_exec():
    with pytest.raises(FlashError):
        parse_elf_segments(make_elf(0x102000, b"x", etype=1))  # ET_REL


# --- block building ---

def test_build_blocks_aligns_and_pads():
    # 700 bytes at an aligned address -> 2 blocks, second 0xFF-padded
    blocks = build_blocks([(0x102000, b"\x11" * 700)])
    assert len(blocks) == 2
    assert blocks[0][0] == 0x102000
    assert blocks[1][0] == 0x102200
    assert all(len(b) == 512 for _, b in blocks)
    # tail padding is erased-flash 0xFF
    assert blocks[1][1][-100:] == b"\xFF" * 100


def test_build_blocks_merges_shared_page():
    # Two adjacent segments sharing the flash page at 0x102200: seg1 ends at
    # 0x1022AC, seg2 starts there. The shared page must keep seg1's real tail —
    # writing it twice (seg2's 0xFF head last) would wipe it (the on-hardware bug).
    seg1 = (0x102000, b"\x11" * 0x2AC)   # ends 0x1022AC (unaligned)
    seg2 = (0x1022AC, b"\x22" * 0x100)   # ends 0x1023AC
    blocks = dict(build_blocks([seg1, seg2]))
    shared = blocks[0x102200]            # 0x102200..0x102400
    assert shared[0x00:0xAC] == b"\x11" * 0xAC   # seg1 tail preserved, NOT 0xFF
    assert shared[0xAC:0x100] == b"\x22" * 0x54  # seg2 head
    # no duplicate block addresses (would mean a double-write of a page)
    addrs = [a for a, _ in build_blocks([seg1, seg2])]
    assert len(addrs) == len(set(addrs))


def test_build_blocks_separate_when_far_apart():
    # A real gap between segments -> separate groups, no giant fill between them.
    blocks = build_blocks([(0x102000, b"\x11" * 16), (0x120000, b"\x22" * 16)])
    assert {a for a, _ in blocks} == {0x102000, 0x120000}


def test_build_blocks_pads_low_side():
    # unaligned start -> leading 0xFF padding down to the block boundary
    blocks = build_blocks([(0x102010, b"\x22" * 16)])
    assert len(blocks) == 1
    addr, block = blocks[0]
    assert addr == 0x102000
    assert block[:0x10] == b"\xFF" * 0x10
    assert block[0x10:0x20] == b"\x22" * 16


# --- resolve_image ---

def test_resolve_image_bytes_passthrough():
    assert resolve_image(b"abc") == b"abc"


def test_resolve_image_path(tmp_path):
    p = tmp_path / "fw.elf"
    p.write_bytes(b"\x7fELFdata")
    assert resolve_image(str(p)) == b"\x7fELFdata"


def test_resolve_image_missing_path_errors(tmp_path):
    with pytest.raises(FlashError):
        resolve_image(str(tmp_path / "nope.elf"))


def test_resolve_image_none_no_build(tmp_path, monkeypatch):
    monkeypatch.delenv("PM3PY_FIRMWARE", raising=False)
    monkeypatch.chdir(tmp_path)  # no firmware/ build output here
    with pytest.raises(FlashError):
        resolve_image(None)


# --- build_firmware / find_firmware_src ---

def test_find_firmware_src_env(tmp_path, monkeypatch):
    src = _fake_src(tmp_path)
    monkeypatch.setenv("PM3PY_FIRMWARE_SRC", str(src))
    assert find_firmware_src() == src


def test_build_firmware_invalid_platform():
    with pytest.raises(FlashError, match="PLATFORM"):
        build_firmware("PM3; rm -rf /")


def test_build_firmware_runs_make(tmp_path, monkeypatch):
    src = _fake_src(tmp_path)
    obj = src / "armsrc" / "obj"
    obj.mkdir()
    (obj / "fullimage.elf").write_bytes(b"\x7fELF")
    seen = {}

    def fake_run(cmd, *a, **k):
        seen["cmd"] = cmd
        return _proc(0)

    monkeypatch.setattr(flashmod.subprocess, "run", fake_run)
    elf = build_firmware("PM3GENERIC", firmware_dir=str(src))
    assert elf == obj / "fullimage.elf"
    # explicit platform, non-shell argv, armsrc/all target
    assert seen["cmd"] == ["make", "-C", str(src), "PLATFORM=PM3GENERIC", "armsrc/all"]


def test_build_firmware_make_fails(tmp_path, monkeypatch):
    src = _fake_src(tmp_path)
    monkeypatch.setattr(flashmod.subprocess, "run", lambda *a, **k: _proc(2))
    with pytest.raises(FlashError, match="build failed"):
        build_firmware("PM3RDV4", firmware_dir=str(src))


def test_build_firmware_missing_source(monkeypatch):
    monkeypatch.setattr(flashmod, "find_firmware_src", lambda fd=None: None)
    with pytest.raises(FlashError, match="source not found"):
        build_firmware("PM3GENERIC")


def test_build_firmware_make_not_on_path(tmp_path, monkeypatch):
    # make missing/not-executable -> FlashError, not a raw OSError traceback
    src = _fake_src(tmp_path)

    def boom(*a, **k):
        raise FileNotFoundError(2, "No such file or directory", "make")

    monkeypatch.setattr(flashmod.subprocess, "run", boom)
    with pytest.raises(FlashError, match="could not run make"):
        build_firmware("PM3GENERIC", firmware_dir=str(src))


def test_find_firmware_src_ignores_cwd(tmp_path, monkeypatch):
    # A firmware/ in the CWD must NOT be picked up — build_firmware runs `make` on
    # the result, so a stray ./firmware would execute an untrusted Makefile.
    cwd_fw = tmp_path / "firmware"
    cwd_fw.mkdir()
    _fake_src(cwd_fw)
    monkeypatch.chdir(tmp_path)
    monkeypatch.delenv("PM3PY_FIRMWARE_SRC", raising=False)
    monkeypatch.setattr(flashmod, "_PACKAGE_ROOT", tmp_path / "nonexistent")
    assert find_firmware_src() is None


# --- Flasher: probing ---

def test_device_info_parses_flags():
    t = AsyncMock()
    t.send_old.return_value = old_resp(Cmd.DEVICE_INFO, arg0=BOOTROM_FLAGS)
    info = run(Flasher(t).device_info())
    assert info["in_bootrom"] is True
    assert info["in_os"] is False
    assert info["bootrom_present"] is True


# --- Flasher: enter_bootloader ---

def test_enter_bootloader_already_in_bootrom():
    t = AsyncMock()
    t.send_old.return_value = old_resp(Cmd.DEVICE_INFO, arg0=BOOTROM_FLAGS)
    info = run(Flasher(t).enter_bootloader())
    assert info["in_bootrom"] is True
    t.send_old_no_response.assert_not_called()   # no handover needed
    t.reopen.assert_not_called()


def test_enter_bootloader_from_os_handover():
    t = AsyncMock()
    t.send_old.side_effect = [
        old_resp(Cmd.DEVICE_INFO, arg0=OS_FLAGS),       # first probe: in OS
        old_resp(Cmd.DEVICE_INFO, arg0=BOOTROM_FLAGS),  # after handover: in bootrom
    ]
    run(Flasher(t).enter_bootloader())
    # handover = START_FLASH with no args, then reopen the re-enumerated port
    t.send_old_no_response.assert_awaited_once_with(Cmd.START_FLASH)
    t.reopen.assert_awaited_once()


def test_enter_bootloader_fails_if_not_entered():
    t = AsyncMock()
    t.send_old.side_effect = [
        old_resp(Cmd.DEVICE_INFO, arg0=OS_FLAGS),
        old_resp(Cmd.DEVICE_INFO, arg0=OS_FLAGS),   # still in OS after handover
    ]
    with pytest.raises(FlashError):
        run(Flasher(t).enter_bootloader())


# --- Flasher: flash_image ---

def test_flash_image_full_sequence():
    t = AsyncMock()
    t.send_old.side_effect = [
        old_resp(Cmd.BL_VERSION, arg0=BL_VERSION_1_0_0),  # bl_version
        old_resp(Cmd.ACK),                                # START_FLASH ack
        old_resp(Cmd.ACK),                                # block 0
        old_resp(Cmd.ACK),                                # block 1
    ]
    image = make_elf(0x102000, b"\x11" * 700)  # 2 blocks
    result = run(Flasher(t).flash_image(image))

    assert result == {"success": True, "blocks": 2, "bytes": 1024,
                      "flash_end": FLASH_END_512K}

    calls = t.send_old.call_args_list
    # START_FLASH: fullimage-only window (low=BOOTLOADER_END, no magic)
    assert calls[1].args == (Cmd.START_FLASH, BOOTLOADER_END, FLASH_END_512K, 0)
    # FINISH_WRITE per block at the right addresses, 512-byte payloads
    assert calls[2].args[0] == Cmd.FINISH_WRITE
    assert calls[2].args[1] == 0x102000
    assert len(calls[2].args[4]) == 512
    assert calls[3].args[1] == 0x102200


def test_flash_image_refuses_bootrom_region():
    t = AsyncMock()
    image = make_elf(FLASH_START, b"\x00" * 16)  # 0x100000 = bootrom region
    with pytest.raises(FlashError, match="bootrom"):
        run(Flasher(t).flash_image(image))
    t.send_old.assert_not_called()


def test_flash_image_allow_bootrom_uses_magic():
    t = AsyncMock()
    t.send_old.side_effect = [
        old_resp(Cmd.BL_VERSION, arg0=BL_VERSION_1_0_0),
        old_resp(Cmd.ACK),   # START_FLASH
        old_resp(Cmd.ACK),   # 1 block
    ]
    image = make_elf(FLASH_START, b"\x00" * 16)
    run(Flasher(t).flash_image(image, allow_bootrom=True))
    # window opens at FLASH_START and carries the unlock magic
    assert t.send_old.call_args_list[1].args == (
        Cmd.START_FLASH, FLASH_START, FLASH_END_512K, START_FLASH_MAGIC)


def test_flash_image_nack_raises_with_status():
    t = AsyncMock()
    t.send_old.side_effect = [
        old_resp(Cmd.BL_VERSION, arg0=BL_VERSION_1_0_0),
        old_resp(Cmd.ACK),                 # START_FLASH
        old_resp(Cmd.NACK, arg0=0x8),      # write NACK, EFC PROGE
    ]
    image = make_elf(0x102000, b"\x11" * 16)
    with pytest.raises(FlashError, match="0x8"):
        run(Flasher(t).flash_image(image))


def test_flash_image_start_flash_not_acked():
    t = AsyncMock()
    t.send_old.side_effect = [
        old_resp(Cmd.BL_VERSION, arg0=BL_VERSION_1_0_0),
        old_resp(Cmd.NACK),   # window arm rejected
    ]
    image = make_elf(0x102000, b"\x11" * 16)
    with pytest.raises(FlashError, match="START_FLASH"):
        run(Flasher(t).flash_image(image))


# --- Flasher: end-to-end flash() ---

def test_flash_end_to_end_resets():
    t = AsyncMock()
    t.send_old.side_effect = [
        old_resp(Cmd.DEVICE_INFO, arg0=BOOTROM_FLAGS),   # enter_bootloader probe
        old_resp(Cmd.BL_VERSION, arg0=BL_VERSION_1_0_0),
        old_resp(Cmd.ACK),   # START_FLASH
        old_resp(Cmd.ACK),   # block 0
        old_resp(Cmd.ACK),   # block 1
    ]
    image = make_elf(0x102000, b"\x11" * 700)
    result = run(Flasher(t).flash(image))
    assert result["success"] is True
    # reboot into the freshly written OS
    t.send_old_no_response.assert_awaited_with(Cmd.HARDWARE_RESET)