"""Plan or safely flash the EchoEar ESP32-S3 16 MiB A/B firmware.

No serial port is opened unless the command is ``flash --execute``. All writes
require a matching physical MAC, a fresh full-flash backup, and verified input
images. The device remains in its bootloader after completion: reset it manually.
Full USB firmware writes require --reset-runtime and subsequent Wi-Fi provisioning
to clear old OTA version/journal metadata. Identity-only writes may preserve NVS.
Requires Python 3.10+; hardware commands additionally need esptool 4.x or 5.x.
"""
from __future__ import annotations

import argparse
from dataclasses import dataclass
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path, PureWindowsPath
import re
import struct
import subprocess
import sys

from device_prepare import private_directory

PROJECT = Path(__file__).resolve().parents[1]
FLASH_SIZE = 0x1000000
FACTORY_OFFSET = 0xE4E000
FACTORY_SIZE = 0x6000
NVS_OFFSET = 0x9000
NVS_SIZE = 0x6000
TABLE_OFFSET = 0x8000
SECTOR = 0x1000
# name: type, subtype, offset, size, flags. Exact supported 16 MiB A/B layout.
LAYOUT = {
    "nvs": (1, 2, 0x9000, 0x6000, 0),
    "phy_init": (1, 1, 0xF000, 0x1000, 0),
    "ota_0": (0, 0x10, 0x10000, 0x400000, 0),
    "ota_1": (0, 0x11, 0x410000, 0x400000, 0),
    "otadata": (1, 0, 0x810000, 0x2000, 0),
    "model": (1, 0x82, 0x812000, 0x80000, 0),
    "coredump": (1, 3, 0x91E000, 0x10000, 0),
    "storage": (1, 0x83, 0x92E000, 0x200000, 0),
    "eaf": (1, 0x82, 0xB2E000, 0x320000, 0),
    "ble_factory": (1, 2, FACTORY_OFFSET, FACTORY_SIZE, 0),
}
REGIONS = {
    "bootloader": (0, 0x8000),
    "partition-table": (TABLE_OFFSET, SECTOR),
    "app": (0x10000, 0x400000),
    "otadata": (0x810000, 0x2000),
    "model": (0x812000, 0x80000),
    "eaf": (0xB2E000, 0x320000),
}


class FlashError(ValueError):
    """An input or hardware safety check failed."""


def require(condition, message):
    if not condition:
        raise FlashError(message)


def number(value):
    require(type(value) in (str, int), "Expected an integer or hexadecimal string")
    try:
        return int(value, 0) if isinstance(value, str) else value
    except ValueError as exc:
        raise FlashError("Invalid integer in manifest") from exc


def normalize_mac(value):
    require(isinstance(value, str), "MAC must be a string")
    compact = re.sub("[:-]", "", value or "")
    require(bool(re.fullmatch(r"[0-9a-fA-F]{12}", compact)), "MAC must contain 12 hex digits")
    require(compact.lower() not in ("000000000000", "ffffffffffff"), "Invalid MAC address")
    return ":".join(compact[i:i + 2].upper() for i in range(0, 12, 2))


def sha256(data):
    return hashlib.sha256(data).hexdigest()


def load_json(path):
    try:
        value = json.loads(path.read_text(encoding="utf-8-sig"))
    except (OSError, ValueError) as exc:
        raise FlashError(f"Cannot read JSON: {path}") from exc
    require(isinstance(value, dict), f"Expected a JSON object: {path}")
    return value


def object_field(parent, name):
    value = parent.get(name)
    require(isinstance(value, dict), f"Missing or invalid object: {name}")
    return value


def contained_file(directory, relative):
    require(isinstance(relative, str) and relative, "Missing image filename")
    win = PureWindowsPath(relative)
    require(not win.drive and not win.is_absolute() and not Path(relative).is_absolute(),
            "Absolute image paths are not allowed in manifests")
    require(".." not in win.parts and ".." not in Path(relative).parts,
            "Image path must not escape its build/factory directory")
    path = (directory / relative).resolve()
    require(path.is_relative_to(directory.resolve()), "Image symlink escapes its directory")
    require(path.is_file(), f"Image is missing: {path}")
    return path


def partition_table(data):
    """Validate IDF's 32-byte partition records and the mandatory MD5 record."""
    require(len(data) in (0xC00, 0x1000), "Partition table must be 3072 or 4096 bytes")
    result = {}
    for pos in range(0, 0xC00, 32):
        record = data[pos:pos + 32]
        if record[:16] == b"\xeb\xeb" + b"\xff" * 14:
            require(record[16:] == hashlib.md5(data[:pos]).digest(), "Partition-table MD5 mismatch")
            require(data[pos + 32:] == b"\xff" * (len(data) - pos - 32),
                    "Unsupported signed or nonempty partition-table padding")
            require(result == LAYOUT, "Partition layout is not the supported 16 MiB / 4 MiB A/B layout")
            return result
        require(record[:2] == b"\xaa\x50", "Missing partition-table entry or MD5 record")
        _, kind, subtype, offset, size, label, flags = struct.unpack("<HBBII16sI", record)
        try:
            name = label.split(b"\0", 1)[0].decode("ascii")
        except UnicodeError as exc:
            raise FlashError("Non-ASCII partition name") from exc
        require(name and name not in result, "Duplicate or empty partition name")
        require(offset % SECTOR == 0 and size > 0 and size % SECTOR == 0,
                "Unaligned partition")
        require(offset >= 0x9000 and offset + size <= FLASH_SIZE, "Partition out of flash bounds")
        for old in result.values():
            require(offset + size <= old[2] or offset >= old[2] + old[3], "Overlapping partitions")
        result[name] = (kind, subtype, offset, size, flags)
    raise FlashError("Missing partition-table MD5")


def esp_image(data, application=False):
    """Validate ESP image segments, XOR checksum, appended SHA256 and chip ID."""
    require(len(data) >= 64 and data[0] == 0xE9, "Not an ESP executable image")
    require(1 <= data[1] <= 16, "Invalid ESP image segment count")
    require(struct.unpack_from("<H", data, 12)[0] == 9, "Image is not for ESP32-S3")
    require(data[2] == 3 and data[3] == 0x4F, "Image must target DOUT / 16MB / 80MHz")
    require(data[23] == 1, "ESP image must include its validation SHA256")
    position, checksum = 24, 0xEF
    for _ in range(data[1]):
        require(position + 8 <= len(data), "Truncated ESP segment header")
        size = struct.unpack_from("<I", data, position + 4)[0]
        position += 8
        require(size > 0 and position + size <= len(data), "Truncated or empty ESP segment")
        for byte in data[position:position + size]:
            checksum ^= byte
        position += size
    checksum_pos = (position // 16 + 1) * 16 - 1
    require(checksum_pos + 33 == len(data), "Unexpected ESP image tail or signature")
    require(data[checksum_pos] == checksum, "ESP image checksum mismatch")
    require(data[-32:] == hashlib.sha256(data[:-32]).digest(), "ESP image SHA256 mismatch")
    if not application:
        return {}
    require(struct.unpack_from("<I", data, 32)[0] == 0xABCD5432, "App descriptor is missing")
    def cstring(offset, length):
        return data[offset:offset + length].split(b"\0", 1)[0].decode("utf-8", errors="replace")
    project = cstring(80, 32)
    require(project == "VolcRTCDemo", "Image is not the EchoEar VolcRTCDemo application")
    versions = set(re.findall(rb"\b\d+\.\d+\.\d+-controls\.\d+\b", data))
    require(len(versions) == 1, "Application must contain one identifiable controls firmware version")
    return {"project": project, "firmware_version": next(iter(versions)).decode("ascii"),
            "idf_app_version": cstring(48, 32), "idf_version": cstring(144, 32),
            "hardware_acceptance": "not_verified_by_this_tool"}


@dataclass
class Image:
    name: str
    offset: int
    data: bytes
    source: str

    def summary(self):
        return {"name": self.name, "offset": hex(self.offset), "size": len(self.data),
                "sha256": sha256(self.data), "source": self.source}


@dataclass
class Plan:
    images: list[Image]
    application: dict
    factory: dict | None
    identity_only: bool
    reset_runtime: bool

    def summary(self):
        return {"mode": "identity-only" if self.identity_only else "full-firmware",
                "chip": "esp32s3", "flash_size": FLASH_SIZE,
                "application": {"action": "retained_on_device", "firmware_version": None}
                if self.identity_only else self.application,
                "reference_build": self.application if self.identity_only else None,
                "factory": self.factory,
                "reset_runtime": self.reset_runtime,
                "execute_requires_reset_runtime": not self.identity_only,
                "runtime_effect": "Erase Wi-Fi, binding, session and OTA state; BLE provisioning required"
                if self.reset_runtime else "Preserve Wi-Fi, binding and session state",
                "images": [image.summary() for image in self.images],
                "completion": "Verify written bytes, stay in bootloader; manual reset required"}


def build_plan(build_dir, factory_dir=None, identity_only=False, reset_runtime=False, expect_mac=None):
    build_dir = Path(build_dir).resolve()
    manifest = load_json(build_dir / "flasher_args.json")
    require(object_field(manifest, "extra_esptool_args").get("chip") == "esp32s3", "Build target must be esp32s3")
    settings = object_field(manifest, "flash_settings")
    require(settings.get("flash_mode") == "dout" and settings.get("flash_freq") == "80m"
            and settings.get("flash_size") in ("detect", "16MB"), "Unsupported build flash settings")
    flash_files = manifest.get("flash_files")
    require(isinstance(flash_files, dict) and len(flash_files) == len(REGIONS),
            "Build must contain exactly bootloader, app, partition-table, otadata, model and eaf")
    mapped = {number(key): value for key, value in flash_files.items()}
    require(len(mapped) == len(REGIONS), "Duplicate flash offsets")
    images, app_info = [], {}
    for name, (offset, maximum) in REGIONS.items():
        entry = object_field(manifest, name)
        require(number(entry.get("offset")) == offset and mapped.get(offset) == entry.get("file"),
                f"Build metadata for {name} disagrees with supported address")
        require(entry.get("encrypted") in (False, "false"), "Encrypted build images are unsupported")
        path = contained_file(build_dir, entry.get("file"))
        data = path.read_bytes()
        require(0 < len(data) <= maximum, f"{name} image does not fit its allocated region")
        if name == "partition-table":
            partition_table(data)
        elif name == "bootloader":
            esp_image(data)
        elif name == "app":
            app_info = esp_image(data, application=True)
        elif name == "otadata":
            require(len(data) == maximum and data == b"\xff" * maximum,
                    "Expected the IDF initial empty OTA selection image")
        if not identity_only:
            images.append(Image(name, offset, data, str(path)))
    factory_summary = None
    if factory_dir:
        factory_dir = Path(factory_dir).resolve()
        factory = load_json(factory_dir / "manifest.json")
        require(factory.get("schema_version") == 1 and factory.get("chip") == "esp32s3"
                and number(factory.get("flash_size")) == FLASH_SIZE, "Unsupported factory manifest")
        require(isinstance(factory.get("device_id"), str)
                and bool(re.fullmatch(r"KP01-[0-9A-F]{12}", factory["device_id"])),
                "Invalid factory device ID")
        mac = normalize_mac(factory.get("expected_mac"))
        if expect_mac:
            require(mac == normalize_mac(expect_mac), "Factory manifest MAC does not match --expect-mac")
        entry = object_field(factory, "image")
        require(entry.get("file") == "ble_factory.bin" and number(entry.get("offset")) == FACTORY_OFFSET
                and number(entry.get("size")) == FACTORY_SIZE, "Incorrect factory image location/size")
        path = contained_file(factory_dir, entry.get("file"))
        data = path.read_bytes()
        require(len(data) == FACTORY_SIZE and sha256(data) == entry.get("sha256"),
                "Factory image size or SHA256 mismatch")
        require(data != b"\xff" * FACTORY_SIZE and data != b"\0" * FACTORY_SIZE,
                "Factory image is empty")
        images.append(Image("ble_factory", FACTORY_OFFSET, data, str(path)))
        factory_summary = {"device_id": factory["device_id"], "expected_mac": mac}
    require(not identity_only or factory_summary is not None, "--identity-only requires --factory-dir")
    if reset_runtime:
        images.append(Image("reset-runtime", NVS_OFFSET, b"\xff" * NVS_SIZE, "explicit --reset-runtime"))
    images.sort(key=lambda image: image.offset)
    previous_end = 0
    for image in images:
        end = (image.offset + len(image.data) + SECTOR - 1) // SECTOR * SECTOR
        require(image.offset % SECTOR == 0 and image.offset >= previous_end and end <= FLASH_SIZE,
                "Flash regions or their erase sectors overlap")
        previous_end = end
    return Plan(images, app_info, factory_summary, identity_only, reset_runtime)


class Esptool:
    def __init__(self, python, port, baud, runner=subprocess.run):
        self.python, self.port, self.baud, self.runner = python, port, baud, runner
        self.major = 4
        self.first = True

    def invoke(self, args):
        result = self.runner([self.python, "-X", "utf8", "-m", "esptool"] + args,
                             stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                             text=True, encoding="utf-8", errors="replace", check=False)
        if result.returncode:
            raise FlashError(f"esptool failed ({result.returncode}); no later commands were run:\n{result.stdout}")
        return result.stdout

    def check_version(self):
        output = self.invoke(["version"])
        match = re.search(r"(?<!\d)([45])\.\d+(?:\.\d+)?", output)
        require(bool(match), "Install esptool 4.x or 5.x in the selected Python environment")
        self.major = int(match.group(1))

    def command(self, command, *args):
        before = "default_reset" if self.first else "no_reset"
        self.first = False
        prefix = ["--chip", "esp32s3", "--port", self.port, "--baud", str(self.baud),
                  "--before", before, "--after", "no_reset"]
        if self.major == 5:
            prefix = [x.replace("_", "-") if x in ("default_reset", "no_reset") else x for x in prefix]
            command = command.replace("_", "-")
            args = tuple(x.replace("_", "-") if x.startswith("--flash_") else x for x in args)
        return self.invoke(prefix + [command] + list(args))


def save_json(path, value):
    with path.open("x", encoding="utf-8") as output:
        json.dump(value, output, ensure_ascii=True, indent=2)
        output.write("\n")
    os.chmod(path, 0o600)


def execute_plan(plan, *, python, port, expect_mac, backup_dir, baud=460800, runner=subprocess.run):
    require(plan.identity_only or plan.reset_runtime,
            "Full USB firmware flashing requires --reset-runtime to remove stale OTA version/journal metadata; "
            "Wi-Fi must be provisioned again. For an ordinary update preserving settings, use cloud OTA.")
    require(bool(port) and not port.startswith("-") and "\0" not in port, "An explicit serial port is required")
    require(baud in (115200, 230400, 460800, 921600), "Unsupported baud rate")
    expected = normalize_mac(expect_mac)
    if plan.factory:
        require(plan.factory["expected_mac"] == expected, "Factory manifest MAC mismatch")
    backup_dir = Path(backup_dir).resolve()
    require(not backup_dir.exists(), "Backup directory already exists; select a new private directory")
    tool = Esptool(python, port, baud, runner)
    tool.check_version()
    private_directory(backup_dir)
    save_json(backup_dir / "plan.json", plan.summary())
    # Capture exact validated bytes, avoiding input mutation between planning and writing.
    staged = []
    for image in plan.images:
        target = backup_dir / (image.name + ".bin")
        with target.open("xb") as output:
            output.write(image.data)
        os.chmod(target, 0o600)
        staged.extend([hex(image.offset), str(target)])
    log = []
    def command(name, *args):
        try:
            output = tool.command(name, *args)
        except FlashError as exc:
            log.append({"command": name, "error": str(exc)})
            raise
        log.append({"command": name, "output": output})
        print(f"Completed {name}", flush=True)
        return output
    try:
        probe = command("read_mac")
        macs = re.findall(r"(?im)^\s*MAC:\s*([0-9a-f]{2}(?::[0-9a-f]{2}){5})\b", probe)
        require({normalize_mac(mac) for mac in macs} == {expected},
                "Connected device MAC does not match --expect-mac")
        require(bool(re.search(r"ESP32[- ]?S3", probe, re.I)), "Connected device is not ESP32-S3")
        flash = command("flash_id")
        require(bool(re.search(r"Detected flash size:\s*16\s*MB\b", flash, re.I)), "Connected device flash is not 16MB")
        security = command("get_security_info")
        require(bool(re.search(r"Secure Boot:\s*Disabled", security, re.I))
                and bool(re.search(r"Flash Encryption:\s*Disabled", security, re.I)),
                "Secure Boot / flash encryption is enabled or could not be confirmed disabled")
        backup = backup_dir / "before-16mb.bin"
        command("read_flash", "0x0", hex(FLASH_SIZE), str(backup))
        require(backup.is_file() and backup.stat().st_size == FLASH_SIZE, "Incomplete full-flash backup; refusing to write")
        original = backup.read_bytes()
        os.chmod(backup, 0o600)
        save_json(backup_dir / "backup-manifest.json", {
            "schema_version": 1, "created_at": datetime.now(timezone.utc).isoformat(),
            "expected_mac": expected, "chip": "esp32s3", "flash_size": FLASH_SIZE,
            "file": backup.name, "size": len(original), "sha256": sha256(original),
        })
        table = original[TABLE_OFFSET:TABLE_OFFSET + 0xC00]
        if table == b"\xff" * 0xC00:
            require(not plan.identity_only and plan.factory is not None,
                    "Blank partition table needs full firmware and --factory-dir")
        else:
            partition_table(table)
        preserved = []
        for name, offset, size, enabled in (
                ("preserved-runtime", NVS_OFFSET, NVS_SIZE, not plan.reset_runtime),
                ("preserved-factory", FACTORY_OFFSET, FACTORY_SIZE, plan.factory is None)):
            if enabled:
                target = backup_dir / (name + ".bin")
                with target.open("xb") as output:
                    output.write(original[offset:offset + size])
                os.chmod(target, 0o600)
                preserved.extend([hex(offset), str(target)])
        # Recheck staged bytes immediately before the first destructive command.
        for index, image in enumerate(plan.images):
            require(Path(staged[2 * index + 1]).read_bytes() == image.data, "Staged image changed; refusing to write")
        command("write_flash", "--flash_mode", "keep", "--flash_size", "keep", "--flash_freq", "keep", *staged)
        command("verify_flash", *(staged + preserved))
        result = {"status": "written_and_verified", "expected_mac": expected,
                  "backup_dir": str(backup_dir),
                  "firmware_version": None if plan.identity_only else plan.application["firmware_version"],
                  "firmware_action": "preserved" if plan.identity_only else "written",
                  "hardware_acceptance": "not_tested", "next_step": "Manually reset device to boot firmware"}
        save_json(backup_dir / "result.json", result)
        return result
    finally:
        save_json(backup_dir / "esptool-log.json", log)


def parser():
    result = argparse.ArgumentParser(description=__doc__)
    result.add_argument("command", nargs="?", default="plan", choices=("plan", "ports", "flash"))
    result.add_argument("--build-dir", type=Path, default=PROJECT / "build")
    result.add_argument("--factory-dir", type=Path)
    result.add_argument("--identity-only", action="store_true")
    result.add_argument("--reset-runtime", action="store_true", help="Required for full USB firmware writes; erase ordinary NVS including Wi-Fi, binding, session and OTA metadata")
    result.add_argument("--execute", action="store_true")
    result.add_argument("--port")
    result.add_argument("--expect-mac")
    result.add_argument("--backup-dir", type=Path)
    result.add_argument("--python", default=sys.executable, help="Python with esptool 4.x or 5.x and pyserial installed")
    result.add_argument("--baud", type=int, choices=(115200, 230400, 460800, 921600), default=460800)
    return result


def main(argv=None):
    cli = parser()
    args = cli.parse_args(argv)
    try:
        require(not args.execute or args.command == "flash", "--execute is only valid with the flash command")
        if args.command == "ports":
            code = ("import json; from serial.tools import list_ports; "
                    "print(json.dumps([{'port':p.device,'description':p.description,'hwid':p.hwid} "
                    "for p in list_ports.comports()],ensure_ascii=True,indent=2))")
            subprocess.run([args.python, "-X", "utf8", "-c", code], check=True)
            return 0
        plan = build_plan(args.build_dir, args.factory_dir, args.identity_only, args.reset_runtime, args.expect_mac)
        if args.command == "flash":
            require(args.execute and args.port and args.expect_mac and args.backup_dir,
                    "flash requires --execute --port --expect-mac --backup-dir; use plan for an offline preview")
            result = execute_plan(plan, python=args.python, port=args.port, expect_mac=args.expect_mac,
                                  backup_dir=args.backup_dir, baud=args.baud)
        else:
            result = {"status": "plan_only_no_hardware_access", **plan.summary()}
        print(json.dumps(result, ensure_ascii=True, indent=2))
        return 0
    except (ValueError, OSError, subprocess.SubprocessError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())