"""Flashing: image validation on the real official image, plan building,
and the runner against a fake esptool (no hardware)."""

import hashlib
import sys
import threading
from pathlib import Path

import pytest

from kota_tech.flash.image import ImageError, inspect_image, official_images
from kota_tech.flash.plan import Mode, build_plan, check_partition_table
from kota_tech.flash.runner import FlashRunner

FW = Path(__file__).resolve().parents[1] / "firmware" / "new_system.bin"


@pytest.fixture(scope="module")
def image():
    imgs = official_images()
    assert len(imgs) == 1
    return imgs[0]


# ------------------------------------------------------------------ image

def test_official_image_is_the_deployed_v1_build(image):
    assert image.app.project == "kota-sensorboard"
    assert image.app.elf_sha.startswith("45fa9d5b9")       # same as the live board
    assert image.flash_size == "8MB"
    assert image.partition("ota_0").offset == 0x40000
    assert image.partition("otadata").offset == 0x31000
    assert image.partition("nvs").offset == 0x10000


def test_app_bytes_are_a_complete_app_image(image):
    app = image.app_bytes()
    assert app[0] == 0xE9
    assert len(app) == len(image.data) - 0x40000              # ends where the file ends
    assert len(image.blank_otadata()) == 0x2000 and set(image.blank_otadata()) == {0xFF}


def test_modified_image_is_refused(tmp_path):
    data = bytearray(FW.read_bytes())
    data[-1] ^= 0xFF
    bad = tmp_path / "new_system.bin"
    bad.write_bytes(data)
    manifest = tmp_path / "firmware.toml"
    manifest.write_text(f'[[image]]\nfile = "new_system.bin"\nsha256 = "{hashlib.sha256(FW.read_bytes()).hexdigest()}"\n')
    with pytest.raises(ImageError, match="SHA-256"):
        official_images(manifest)


@pytest.mark.parametrize("patch,msg", [
    ((0, 0x00), "not an ESP32 firmware"),
    ((12, 0x05), "not ESP32-S3"),
])
def test_wrong_images_are_rejected(tmp_path, patch, msg):
    data = bytearray(FW.read_bytes())
    data[patch[0]] = patch[1]
    p = tmp_path / "x.bin"
    p.write_bytes(data)
    with pytest.raises(ImageError, match=msg):
        inspect_image(p)


# ------------------------------------------------------------------- plan

def test_app_only_plan_keeps_nvs(image, tmp_path):
    plan = build_plan(image, Mode.APP_ONLY, "COM5", tmp_path)
    kinds = [s.kind for s in plan.steps]
    assert kinds == ["read_pt", "write"]
    write = plan.steps[-1].args
    offsets = [a for a in write if a.startswith("0x")]
    assert offsets == ["0x31000", "0x40000"]          # otadata + ota_0 only
    assert "0x10000" not in write and "0x0" not in write   # NVS / bootloader untouched
    assert write[write.index("--after") + 1] == "hard_reset"
    assert (tmp_path / "app_ota0.bin").read_bytes() == image.app_bytes()


def test_full_plan_with_backup(image, tmp_path):
    backup = tmp_path / "bk.bin"
    plan = build_plan(image, Mode.FULL, "COM5", tmp_path, backup_to=backup)
    assert [s.kind for s in plan.steps] == ["backup", "write"]
    assert plan.steps[0].args[-3:] == ["0", "0x800000", str(backup)]
    assert plan.steps[0].args[plan.steps[0].args.index("--after") + 1] == "no_reset"
    assert plan.steps[1].args[-2:] == ["0x0", str(image.path)]
    assert "Back up" in plan.summary()[0]


def test_partition_table_check(image):
    assert check_partition_table(image.partition_table, image) is None
    assert "blank" in check_partition_table(b"\xff" * 0xC00, image)
    other = bytearray(image.partition_table)
    other[40] ^= 1
    assert "different partition layout" in check_partition_table(bytes(other), image)


# ----------------------------------------------------------------- runner

FAKE = """
import sys, pathlib
args = sys.argv[1:]
if "read_flash" in args:
    out = pathlib.Path(args[-1])
    src = pathlib.Path(r"{fw}").read_bytes()
    off, size = int(args[-3], 0), int(args[-2], 0)
    {pt_line}
    print("Read ... (50 %)", end="\\r"); print("Read ... (100 %)")
    sys.exit(0)
if "write_flash" in args:
    for p in (10, 60, 100):
        print(f"Writing at 0x00040000... ({{p}} %)", end="\\r", flush=True)
    print(); print("Hash of data verified.")
    sys.exit({write_code})
"""


def _fake(tmp_path, write_code=0, pt="same", fail_text=""):
    pt_line = {"same": "out.write_bytes(src[off:off+size])",
               "blank": "out.write_bytes(b'\\xff' * size)"}[pt]
    script = FAKE.format(fw=FW, pt_line=pt_line, write_code=write_code)
    if fail_text:
        script = script.replace('print("Hash of data verified.")', f'print("{fail_text}")')
    f = tmp_path / "fake_esptool.py"
    f.write_text(script)
    return [sys.executable, str(f)]


def _run(plan, cmd):
    steps, progress, output, done = [], [], [], threading.Event()
    result = {}
    r = FlashRunner(plan, lambda i, t: steps.append(i), progress.append, output.append,
                    lambda res: (result.setdefault("r", res), done.set()), esptool_cmd=cmd)
    r.start()
    assert done.wait(30)
    return result["r"], steps, progress, output


def test_runner_app_only_success(image, tmp_path):
    plan = build_plan(image, Mode.APP_ONLY, "COM5", tmp_path)
    res, steps, progress, output = _run(plan, _fake(tmp_path))
    assert res.ok, res
    assert steps == [0, 1]
    assert 60 in progress and progress[-1] == 100
    assert any("Hash of data verified" in o for o in output)


def test_runner_refuses_app_only_on_blank_chip(image, tmp_path):
    plan = build_plan(image, Mode.APP_ONLY, "COM5", tmp_path)
    res, steps, *_ = _run(plan, _fake(tmp_path, pt="blank"))
    assert not res.ok and "Full mode" in res.message
    assert steps == [0]                       # never reached the write step


def test_runner_reports_connect_failure_with_hint(image, tmp_path):
    plan = build_plan(image, Mode.FULL, "COM5", tmp_path)
    cmd = _fake(tmp_path, write_code=2,
                fail_text="A fatal error occurred: Failed to connect to ESP32-S3: No serial data received.")
    res, *_ = _run(plan, cmd)
    assert not res.ok and res.failed_step == 0
    assert "BOOT" in res.hint


# ------------------------------------------------- developer builds (v0.5)

from kota_tech.flash.image import FirmwareLibrary, app_image_length  # noqa: E402


def _app_only_file(image, tmp_path, name="kota-sensorboard.bin"):
    p = tmp_path / name
    p.write_bytes(image.app_bytes())          # what a developer build/ folder contains
    return p


def test_integrity_check_on_factory_image(image):
    d = image.data
    assert app_image_length(d, 0) > 0                          # bootloader
    assert app_image_length(d, 0x40000) == len(d) - 0x40000    # app ends exactly at file end


@pytest.mark.parametrize("damage,msg", [
    (lambda b: b[:len(b) // 2], "truncated"),
    (lambda b: b[:0x40000 + 5000] + bytes([b[0x40000 + 5000] ^ 1]) + b[0x40000 + 5001:], "corrupted"),
])
def test_corrupted_or_truncated_images_refused(tmp_path, damage, msg):
    p = tmp_path / "broken.bin"
    p.write_bytes(damage(FW.read_bytes()))
    with pytest.raises(ImageError, match=msg):
        inspect_image(p, origin="developer")


def test_import_app_only_developer_build(image, tmp_path):
    lib = FirmwareLibrary(tmp_path / "lib")
    src = _app_only_file(image, tmp_path)
    img = lib.import_file(src, reference=image)
    assert img.is_developer and img.kind == "app"
    assert img.partition("ota_0").offset == 0x40000            # layout from the factory image
    assert img.app_bytes() == image.app_bytes()
    listed, problems = lib.images(image)
    assert [i.sha256 for i in listed] == [img.sha256] and problems == []
    assert lib.import_file(src, reference=image).sha256 == img.sha256   # re-import: no duplicate
    assert len(lib.images(image)[0]) == 1


def test_app_only_build_cannot_be_flashed_full(image, tmp_path):
    lib = FirmwareLibrary(tmp_path / "lib")
    img = lib.import_file(_app_only_file(image, tmp_path), reference=image)
    plan = build_plan(img, Mode.APP_ONLY, "COM5", tmp_path / "w")
    assert [s.kind for s in plan.steps] == ["read_pt", "write"]
    with pytest.raises(ValueError, match="app only"):
        build_plan(img, Mode.FULL, "COM5", tmp_path / "w")


def test_import_merged_developer_build(image, tmp_path):
    lib = FirmwareLibrary(tmp_path / "lib")
    src = tmp_path / "dev_full.bin"
    src.write_bytes(image.data)
    img = lib.import_file(src, reference=image)
    assert img.kind == "merged" and img.is_developer
    assert [s.kind for s in build_plan(img, Mode.FULL, "COM5", tmp_path / "w").steps] == ["write"]


def test_library_detects_file_changed_after_import(image, tmp_path):
    lib = FirmwareLibrary(tmp_path / "lib")
    img = lib.import_file(_app_only_file(image, tmp_path), reference=image)
    stored = tmp_path / "lib" / img.path.name
    data = bytearray(stored.read_bytes())
    data[-40] ^= 1
    stored.write_bytes(data)
    listed, problems = lib.images(image)
    assert listed == [] and problems


def test_library_remove(image, tmp_path):
    lib = FirmwareLibrary(tmp_path / "lib")
    img = lib.import_file(_app_only_file(image, tmp_path), reference=image)
    lib.remove(img.sha256)
    assert lib.images(image) == ([], []) and not img.path.exists()


def test_non_firmware_file_refused(tmp_path, image):
    p = tmp_path / "notes.bin"
    p.write_bytes(b"hello firmware team")
    with pytest.raises(ImageError, match="not an ESP32 firmware"):
        FirmwareLibrary(tmp_path / "lib").import_file(p, reference=image)