"""Safety regression tests; all device I/O is simulated, no serial port opens."""
import contextlib
import hashlib
import io
import json
from pathlib import Path
import struct
import subprocess
import tempfile
import unittest
from unittest.mock import patch

import device_flash as flash


MAC = "B4:3A:45:16:34:A8"


def table_bytes(layout=None):
    records = b"".join(struct.pack("<HBBII16sI", 0x50AA, *entry[:4], name.encode(), entry[4])
                       for name, entry in (layout or flash.LAYOUT).items())
    records += b"\xeb\xeb" + b"\xff" * 14 + hashlib.md5(records).digest()
    return records.ljust(0xC00, b"\xff")


def executable(application=False, chip=9):
    header = bytearray(24)
    header[:4] = bytes((0xE9, 1, 3, 0x4F))
    struct.pack_into("<H", header, 12, chip)
    header[23] = 1
    payload = bytearray(288 if application else 64)
    if application:
        struct.pack_into("<I", payload, 0, 0xABCD5432)
        payload[16:16 + 4] = b"test"
        payload[48:48 + 11] = b"VolcRTCDemo\0"
        payload[112:112 + 6] = b"v5.5.1"
        version = b"1.7.0-controls.28\0"
        payload[256:256 + len(version)] = version
    result = bytes(header) + struct.pack("<II", 0x3C000020, len(payload)) + bytes(payload)
    checksum = 0xEF
    for value in payload:
        checksum ^= value
    result = result.ljust((len(result) // 16 + 1) * 16 - 1, b"\0") + bytes((checksum,))
    return result + hashlib.sha256(result).digest()


def make_build(directory):
    directory.mkdir()
    manifest = {"extra_esptool_args": {"chip": "esp32s3"},
                "flash_settings": {"flash_mode": "dout", "flash_size": "detect", "flash_freq": "80m"},
                "flash_files": {}}
    for name, (offset, _) in flash.REGIONS.items():
        if name == "bootloader":
            data = executable()
        elif name == "app":
            data = executable(True)
        elif name == "partition-table":
            data = table_bytes()
        elif name == "otadata":
            data = b"\xff" * 0x2000
        else:
            data = b"asset contents\0"
        filename = name + ".bin"
        (directory / filename).write_bytes(data)
        manifest[name] = {"offset": hex(offset), "file": filename, "encrypted": "false"}
        manifest["flash_files"][hex(offset)] = filename
    (directory / "flasher_args.json").write_text(json.dumps(manifest))
    return manifest


def make_factory(directory):
    directory.mkdir()
    data = b"\xaa" * flash.FACTORY_SIZE
    (directory / "ble_factory.bin").write_bytes(data)
    manifest = {"schema_version": 1, "chip": "esp32s3", "flash_size": flash.FLASH_SIZE,
                "device_id": "KP01-123456789ABC", "expected_mac": MAC,
                "image": {"file": "ble_factory.bin", "offset": "0xe4e000",
                          "size": "0x6000", "sha256": flash.sha256(data)}}
    (directory / "manifest.json").write_text(json.dumps(manifest))
    return manifest


class SimulatedEsptool:
    def __init__(self, *, mac=MAC, size="16MB", chip="ESP32-S3", version="4.12.0",
                 security="Disabled", short_backup=False, blank=False, incompatible=False,
                 failure=None, duplicate_mac=None):
        self.mac, self.size, self.chip, self.version = mac, size, chip, version
        self.security, self.short_backup, self.blank = security, short_backup, blank
        self.incompatible, self.failure = incompatible, failure
        self.duplicate_mac = duplicate_mac
        self.calls = []
        self.commands = []

    def __call__(self, args, **kwargs):
        self.calls.append(args)
        commands = {"version", "read_mac", "flash_id", "get_security_info", "read_flash", "write_flash", "verify_flash"}
        command = next(arg.replace("-", "_") for arg in args if arg.replace("-", "_") in commands)
        self.commands.append(command)
        output = "OK"
        if command == self.failure:
            return subprocess.CompletedProcess(args, 2, "Simulated failure")
        if command == "version":
            output = f"esptool v{self.version}\n{self.version}\n"
        elif command == "read_mac":
            output = f"Chip is {self.chip}\nMAC: {self.mac.lower()}\n"
            if self.duplicate_mac:
                output += f"Uploading stub...\nRunning stub...\nStub running...\nMAC: {self.duplicate_mac.lower()}\n"
        elif command == "flash_id":
            output = f"Detected flash size: {self.size}\n"
        elif command == "get_security_info":
            output = f"Secure Boot: {self.security}\nFlash Encryption: Disabled\n"
        elif command == "read_flash":
            data = bytearray(b"\xff" * flash.FLASH_SIZE)
            if not self.blank:
                data[flash.TABLE_OFFSET:flash.TABLE_OFFSET + 0xC00] = table_bytes()
            if self.incompatible:
                data[flash.TABLE_OFFSET] = 0
            if self.short_backup:
                data = data[:64]
            Path(args[-1]).write_bytes(data)
        return subprocess.CompletedProcess(args, 0, output)


class FlashSafetyTests(unittest.TestCase):
    def setUp(self):
        # Keep test writes in the workspace on Windows sandboxed runners.
        test_parent = Path(__file__).resolve().parent
        self.temporary = tempfile.TemporaryDirectory(prefix=".flash-test-", dir=test_parent)
        self.addCleanup(self.temporary.cleanup)
        self.root = Path(self.temporary.name)
        self.assertTrue(self.root.resolve().is_relative_to(test_parent))
        # ACL creation is independently tested by device_prepare; do not run
        # Windows account/ACL commands for each mocked hardware case.
        private = patch.object(flash, "private_directory", side_effect=lambda p: p.mkdir(mode=0o700))
        self.private = private.start()
        self.addCleanup(private.stop)
        self.build = self.root / "build"
        self.manifest = make_build(self.build)
        self.factory = self.root / "factory"
        self.factory_manifest = make_factory(self.factory)

    def update_build_manifest(self):
        (self.build / "flasher_args.json").write_text(json.dumps(self.manifest))

    def plan(self, **kwargs):
        return flash.build_plan(self.build, **kwargs)

    def execute(self, plan=None, **kwargs):
        runner = SimulatedEsptool(**kwargs)
        with contextlib.redirect_stdout(io.StringIO()):
            result = flash.execute_plan(plan or self.plan(reset_runtime=True), python="python", port="COM6", expect_mac=MAC,
                                        backup_dir=self.root / "backup", runner=runner)
        return result, runner

    def blocked(self, *, plan=None, **kwargs):
        runner = SimulatedEsptool(**kwargs)
        with self.assertRaises(flash.FlashError), contextlib.redirect_stdout(io.StringIO()):
            flash.execute_plan(plan or self.plan(reset_runtime=True), python="python", port="COM6", expect_mac=MAC,
                               backup_dir=self.root / "backup", runner=runner)
        self.assertNotIn("write_flash", runner.commands)
        return runner

    def test_default_plan_never_opens_hardware(self):
        with patch.object(flash.subprocess, "run") as run, contextlib.redirect_stdout(io.StringIO()):
            self.assertEqual(flash.main(["--build-dir", str(self.build)]), 0)
        run.assert_not_called()

    def test_plan_has_required_display_and_model_assets(self):
        plan = self.plan()
        self.assertEqual({image.name for image in plan.images}, set(flash.REGIONS))
        self.assertEqual(plan.application["firmware_version"], "1.7.0-controls.28")
        self.assertFalse(plan.reset_runtime)
        self.assertTrue(plan.summary()["execute_requires_reset_runtime"])

    def test_full_flash_requires_runtime_reset_before_any_side_effect(self):
        runner = self.blocked(plan=self.plan())
        self.assertEqual(runner.calls, [])
        self.private.assert_not_called()
        self.assertFalse((self.root / "backup").exists())

    def test_factory_id_is_independent_of_hardware_mac(self):
        plan = self.plan(factory_dir=self.factory, expect_mac=MAC)
        self.assertEqual(plan.factory["device_id"], "KP01-123456789ABC")

    def test_identity_only_and_reset_are_exact_regions(self):
        plan = self.plan(factory_dir=self.factory, identity_only=True, reset_runtime=True)
        self.assertEqual([(x.offset, len(x.data)) for x in plan.images], [(0x9000, 0x6000), (0xE4E000, 0x6000)])
        self.assertEqual(plan.images[0].data, b"\xff" * 0x6000)
        self.assertIsNone(plan.summary()["application"]["firmware_version"])
        self.assertEqual(plan.summary()["reference_build"]["firmware_version"], "1.7.0-controls.28")
        self.assertFalse(plan.summary()["execute_requires_reset_runtime"])

    def test_identity_only_does_not_claim_reference_build_was_installed(self):
        result, runner = self.execute(plan=self.plan(factory_dir=self.factory, identity_only=True))
        self.assertIsNone(result["firmware_version"])
        self.assertEqual(result["firmware_action"], "preserved")
        self.assertIn(str(self.root / "backup" / "preserved-runtime.bin"), runner.calls[-1])

    def test_acl_failure_prevents_staging_or_device_access(self):
        self.private.side_effect = PermissionError("ACL failed")
        runner = SimulatedEsptool()
        with self.assertRaises(PermissionError):
            flash.execute_plan(self.plan(reset_runtime=True), python="python", port="COM6", expect_mac=MAC,
                               backup_dir=self.root / "backup", runner=runner)
        self.assertEqual(runner.commands, ["version"])
        self.assertFalse((self.root / "backup" / "ble_factory.bin").exists())

    def test_identity_only_requires_factory(self):
        with self.assertRaises(flash.FlashError):
            self.plan(identity_only=True)

    def test_factory_mac_must_match(self):
        with self.assertRaises(flash.FlashError):
            self.plan(factory_dir=self.factory, expect_mac="10:20:30:40:50:60")

    def test_factory_hash_tampering_is_rejected(self):
        (self.factory / "ble_factory.bin").write_bytes(b"\xbb" * flash.FACTORY_SIZE)
        with self.assertRaises(flash.FlashError):
            self.plan(factory_dir=self.factory)

    def test_factory_wrong_offset_rejected(self):
        self.factory_manifest["image"]["offset"] = "0x9000"
        (self.factory / "manifest.json").write_text(json.dumps(self.factory_manifest))
        with self.assertRaises(flash.FlashError):
            self.plan(factory_dir=self.factory)

    def test_missing_eaf_rejected(self):
        (self.build / "eaf.bin").unlink()
        with self.assertRaises(flash.FlashError):
            self.plan()

    def test_build_cannot_include_nvs(self):
        self.manifest["flash_files"]["0x9000"] = "runtime.bin"
        self.update_build_manifest()
        with self.assertRaises(flash.FlashError):
            self.plan()

    def test_escape_paths_rejected(self):
        for path in ("../outside.bin", "C:\\outside.bin", "/outside.bin", "..\\outside.bin"):
            with self.subTest(path=path), self.assertRaises(flash.FlashError):
                flash.contained_file(self.build, path)

    def test_image_address_mismatch_rejected(self):
        self.manifest["app"]["offset"] = "0x9000"
        self.update_build_manifest()
        with self.assertRaises(flash.FlashError):
            self.plan()

    def test_oversize_app_rejected(self):
        (self.build / "app.bin").write_bytes(b"\xff" * (0x400000 + 1))
        with self.assertRaises(flash.FlashError):
            self.plan()

    def test_wrong_chip_image_rejected(self):
        (self.build / "app.bin").write_bytes(executable(True, chip=2))
        with self.assertRaises(flash.FlashError):
            self.plan()

    def test_image_checksum_rejected(self):
        image = bytearray(executable(True))
        image[300] ^= 1
        with self.assertRaisesRegex(flash.FlashError, "checksum"):
            flash.esp_image(bytes(image), True)

    def test_image_sha_rejected(self):
        image = bytearray(executable(True))
        image[-1] ^= 1
        with self.assertRaisesRegex(flash.FlashError, "SHA256"):
            flash.esp_image(bytes(image), True)

    def test_bad_partition_md5_rejected(self):
        table = bytearray(table_bytes())
        table[16] ^= 1
        with self.assertRaises(flash.FlashError):
            flash.partition_table(bytes(table))

    def test_valid_md5_wrong_layout_rejected(self):
        layout = dict(flash.LAYOUT)
        layout["ota_0"] = (0, 0x10, 0x10000, 0x300000, 0)
        with self.assertRaises(flash.FlashError):
            flash.partition_table(table_bytes(layout))

    def test_mutated_input_types_fail_closed(self):
        for name in ("extra_esptool_args", "app", "flash_settings"):
            with self.subTest(name=name):
                old = self.manifest[name]
                self.manifest[name] = None
                self.update_build_manifest()
                with self.assertRaises(flash.FlashError):
                    self.plan()
                self.manifest[name] = old

    def test_missing_execute_gate_prevents_io(self):
        with patch.object(flash.subprocess, "run") as run, contextlib.redirect_stderr(io.StringIO()):
            self.assertEqual(flash.main(["flash", "--build-dir", str(self.build)]), 2)
        run.assert_not_called()

    def test_plan_execute_is_rejected(self):
        with contextlib.redirect_stderr(io.StringIO()):
            self.assertEqual(flash.main(["plan", "--execute"]), 2)

    def test_wrong_physical_mac_stops_before_backup(self):
        runner = self.blocked(mac="10:20:30:40:50:60")
        self.assertNotIn("read_flash", runner.commands)

    def test_duplicate_matching_mac_from_esptool_connection_is_accepted(self):
        result, _ = self.execute(duplicate_mac=MAC)
        self.assertEqual(result["status"], "written_and_verified")

    def test_differing_mac_lines_are_rejected(self):
        runner = self.blocked(duplicate_mac="10:20:30:40:50:60")
        self.assertNotIn("read_flash", runner.commands)

    def test_wrong_physical_chip_stops_before_backup(self):
        runner = self.blocked(chip="ESP32-C3")
        self.assertNotIn("read_flash", runner.commands)

    def test_wrong_flash_size_stops_before_backup(self):
        runner = self.blocked(size="8MB")
        self.assertNotIn("read_flash", runner.commands)

    def test_secure_boot_device_is_rejected(self):
        self.blocked(security="Enabled")

    def test_incomplete_backup_prevents_write(self):
        self.blocked(short_backup=True)

    def test_incompatible_live_partition_prevents_write(self):
        self.blocked(incompatible=True)

    def test_blank_board_without_factory_prevents_write(self):
        self.blocked(blank=True)

    def test_blank_board_rejects_identity_only(self):
        self.blocked(plan=self.plan(factory_dir=self.factory, identity_only=True), blank=True)

    def test_blank_board_full_factory_flash_succeeds(self):
        result, runner = self.execute(plan=self.plan(factory_dir=self.factory, reset_runtime=True), blank=True)
        self.assertEqual(result["status"], "written_and_verified")
        self.assertIn("write_flash", runner.commands)

    def test_existing_backup_directory_is_never_overwritten(self):
        directory = self.root / "backup"
        directory.mkdir()
        sentinel = directory / "sentinel.txt"
        sentinel.write_text("keep")
        with self.assertRaises(flash.FlashError):
            self.execute()
        self.assertEqual(sentinel.read_text(), "keep")

    def test_full_backup_before_write_and_readback_after(self):
        result, runner = self.execute()
        self.assertEqual(runner.commands, ["version", "read_mac", "flash_id", "get_security_info", "read_flash", "write_flash", "verify_flash"])
        directory = self.root / "backup"
        manifest = json.loads((directory / "backup-manifest.json").read_text())
        self.assertEqual(manifest["size"], flash.FLASH_SIZE)
        self.assertEqual(manifest["sha256"], flash.sha256((directory / "before-16mb.bin").read_bytes()))
        self.assertEqual(result["hardware_acceptance"], "not_tested")
        for index, args in enumerate(runner.calls[1:]):
            self.assertEqual(args[args.index("--after") + 1], "no_reset")
            self.assertEqual(args[args.index("--before") + 1], "default_reset" if index == 0 else "no_reset")
        verify = runner.calls[-1]
        self.assertIn(str(directory / "reset-runtime.bin"), verify)
        self.assertNotIn(str(directory / "preserved-runtime.bin"), verify)
        self.assertIn(str(directory / "preserved-factory.bin"), verify)

    def test_esptool_v5_uses_hyphenated_commands(self):
        _, runner = self.execute(version="5.0.2")
        self.assertIn("write-flash", runner.calls[-2])
        self.assertIn("verify-flash", runner.calls[-1])
        self.assertIn("--flash-mode", runner.calls[-2])
        self.assertIn("no-reset", runner.calls[-1])

    def test_write_failure_never_claims_success(self):
        runner = SimulatedEsptool(failure="write_flash")
        with self.assertRaises(flash.FlashError), contextlib.redirect_stdout(io.StringIO()):
            flash.execute_plan(self.plan(reset_runtime=True), python="python", port="COM6", expect_mac=MAC,
                               backup_dir=self.root / "backup", runner=runner)
        self.assertNotIn("verify_flash", runner.commands)
        self.assertFalse((self.root / "backup" / "result.json").exists())
        self.assertTrue((self.root / "backup" / "before-16mb.bin").is_file())

    def test_verify_failure_never_claims_success(self):
        runner = SimulatedEsptool(failure="verify_flash")
        with self.assertRaises(flash.FlashError), contextlib.redirect_stdout(io.StringIO()):
            flash.execute_plan(self.plan(reset_runtime=True), python="python", port="COM6", expect_mac=MAC,
                               backup_dir=self.root / "backup", runner=runner)
        self.assertFalse((self.root / "backup" / "result.json").exists())


if __name__ == "__main__":
    unittest.main()