#!/usr/bin/env python3
"""Read-only backup helper for the TimbreESP ESP32.

This script never writes to the ESP32 and never calls erase_flash.  It uses the
PlatformIO esptool installation to read raw flash regions into a timestamped,
private backup directory.

Default regions (data scope):
  bootloader  0x01000, 0x08000
  partitions  0x08000, 0x01000
  nvs         0x09000, 0x05000
  otadata     0x0E000, 0x02000
  littlefs    0x350000, 0xB0000

Use --scope full to read the complete 4 MiB flash in one pass and then extract
all named regions from that image.  The full image includes application slots
as well as the data regions.

Examples:
  python3 scripts/backup_device.py --dry-run
  python3 scripts/backup_device.py --scope full
  python3 scripts/backup_device.py --port /dev/cu.usbserial-11220 --scope data
"""

from __future__ import annotations

import argparse
import datetime as dt
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path
from typing import Iterable


REGIONS: tuple[tuple[str, int, int], ...] = (
    ("bootloader.bin", 0x01000, 0x08000),
    ("partitions.bin", 0x08000, 0x01000),
    ("nvs.bin", 0x09000, 0x05000),
    ("otadata.bin", 0x0E000, 0x02000),
    ("littlefs.bin", 0x350000, 0x0B0000),
    ("app0.bin", 0x10000, 0x1A0000),
    ("app1.bin", 0x1B0000, 0x1A0000),
)

FULL_FLASH_SIZE = 0x400000
EXPECTED_PARTITIONS = ("nvs", "otadata", "app0", "app1", "spiffs")


class BackupError(RuntimeError):
    """A recoverable backup failure."""


def first_existing(candidates: Iterable[Path]) -> Path | None:
    for candidate in candidates:
        if candidate.is_file():
            return candidate
    return None


def find_pio() -> Path:
    home = Path.home()
    candidates = [
        home / ".platformio/penv/bin/pio",
        home / ".platformio/penv/bin/platformio",
        home / ".platformio/penv/Scripts/platformio.exe",
    ]
    found = first_existing(candidates)
    if found is not None:
        return found
    executable = shutil.which("pio") or shutil.which("platformio")
    if executable:
        return Path(executable)
    raise BackupError(
        "No se encontro PlatformIO. Abre el proyecto desde VS Code o instala "
        "PlatformIO Core."
    )


def find_pio_python() -> Path:
    home = Path.home()
    candidates = [
        home / ".platformio/penv/bin/python",
        home / ".platformio/penv/bin/python3",
        home / ".platformio/penv/Scripts/python.exe",
    ]
    found = first_existing(candidates)
    if found is not None:
        return found
    executable = shutil.which("python3") or shutil.which("python")
    if executable:
        return Path(executable)
    raise BackupError("No se encontro un interprete de Python para esptool.")


def find_esptool() -> Path:
    home = Path.home()
    package_root = home / ".platformio/packages"
    candidates = sorted(package_root.glob("tool-esptoolpy*/esptool.py"))
    found = first_existing(candidates)
    if found is not None:
        return found
    executable = shutil.which("esptool.py")
    if executable:
        return Path(executable)
    raise BackupError("No se encontro tool-esptoolpy dentro de ~/.platformio.")


def find_partition_tool() -> Path | None:
    home = Path.home()
    candidates = home.glob(
        ".platformio/packages/framework-arduinoespressif32/tools/gen_esp32part.py"
    )
    return first_existing(candidates)


def run_command(
    command: list[str],
    *,
    label: str,
    log_dir: Path | None = None,
    capture: bool = False,
) -> str:
    print(f"[{label}] {' '.join(command)}", flush=True)
    started = time.monotonic()
    if capture:
        result = subprocess.run(
            command,
            check=False,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
        )
        output = result.stdout
        if result.returncode != 0:
            print(output, file=sys.stderr)
            raise BackupError(f"{label} fallo con codigo {result.returncode}")
        return output

    if log_dir is None:
        result = subprocess.run(command, check=False)
    else:
        log_dir.mkdir(parents=True, exist_ok=True)
        log_path = log_dir / f"{label}.log"
        with log_path.open("w", encoding="utf-8") as log:
            result = subprocess.run(
                command,
                check=False,
                stdout=log,
                stderr=subprocess.STDOUT,
            )
        if result.returncode != 0:
            print(f"Ultimas lineas de {log_path}:", file=sys.stderr)
            try:
                lines = log_path.read_text(encoding="utf-8", errors="replace").splitlines()
                print("\n".join(lines[-30:]), file=sys.stderr)
            except OSError:
                pass
            raise BackupError(f"{label} fallo con codigo {result.returncode}")

    print(f"[{label}] completado en {time.monotonic() - started:.1f}s", flush=True)
    return ""


def detect_port(pio: Path) -> str:
    result = subprocess.run(
        [str(pio), "device", "list"],
        check=False,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    if result.returncode != 0:
        raise BackupError(f"No se pudo ejecutar 'pio device list':\n{result.stdout}")
    ports = re.findall(r"^(/dev/cu\S+)", result.stdout, flags=re.MULTILINE)
    serial_ports = [
        port
        for port in ports
        if "usbserial" in port.lower() or "wchusbserial" in port.lower()
    ]
    unique = sorted(set(serial_ports))
    if not unique:
        raise BackupError(
            "No se detecto ningun puerto USB serial. Desconecta/reconecta el "
            "ESP32 y vuelve a ejecutar el script."
        )
    if len(unique) > 1:
        raise BackupError(
            "Hay varios puertos USB serial conectados: "
            + ", ".join(unique)
            + ". Usa --port para indicar el correcto."
        )
    return unique[0]


def ensure_port_free(port: str) -> None:
    port_path = Path(port)
    if not port_path.exists():
        raise BackupError(f"El puerto {port} ya no existe. Vuelve a detectar el ESP32.")
    lsof = shutil.which("lsof")
    if lsof is None:
        print("[AVISO] lsof no esta disponible; no se pudo comprobar el bloqueo del puerto.")
        return
    result = subprocess.run(
        [lsof, port],
        check=False,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    owners = [line for line in result.stdout.splitlines() if line.strip()]
    if owners:
        print(port, file=sys.stderr)
        print("\n".join(owners), file=sys.stderr)
        raise BackupError(
            f"El puerto {port} esta ocupado. Cierra el Serial Monitor antes "
            "de ejecutar la copia de seguridad."
        )


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_private(path: Path, data: bytes) -> None:
    path.write_bytes(data)
    try:
        path.chmod(0o600)
    except OSError:
        pass


def extract_regions(full_image: Path, output_dir: Path) -> list[dict[str, object]]:
    entries: list[dict[str, object]] = []
    with full_image.open("rb") as source:
        for name, offset, size in REGIONS:
            source.seek(offset)
            data = source.read(size)
            if len(data) != size:
                raise BackupError(
                    f"La imagen completa no contiene {size} bytes para {name} "
                    f"en offset 0x{offset:X}."
                )
            target = output_dir / name
            write_private(target, data)
            entries.append(
                {
                    "file": name,
                    "offset": f"0x{offset:06X}",
                    "size": size,
                    "sha256": sha256_file(target),
                }
            )
            print(f"[extract] {name}: {size} bytes", flush=True)
    return entries


def decode_partitions(python: Path, partition_tool: Path, binary: Path, output: Path) -> str | None:
    if partition_tool is None:
        return "gen_esp32part.py no encontrado"
    command = [str(python), str(partition_tool), "--quiet", str(binary), str(output)]
    result = subprocess.run(
        command,
        check=False,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    if result.returncode != 0:
        return result.stdout.strip() or "No se pudo decodificar la tabla de particiones."
    return None


def make_readme(scope: str, port: str, baud: int, regions: list[dict[str, object]]) -> str:
    lines = [
        "TimbreESP backup (read-only)",
        "===========================",
        "",
        f"Scope: {scope}",
        f"Port: {port}",
        f"Baud: {baud}",
        "",
        "IMPORTANTES:",
        "- Este directorio contiene datos privado del dispositivo.",
        "- NVS puede contener credenciales WiFi y datos de autenticacion.",
        "- No ejecutar erase_flash sobre el ESP32 usando este directorio.",
        "- Guardar tambien una copia cifrada fuera de la Mac cuando sea posible.",
        "",
        "Regiones:",
    ]
    for region in regions:
        lines.append(
            f"- {region['file']}: offset {region['offset']}, {region['size']} bytes"
        )
    lines.extend(
        [
            "",
            "La tabla de particiones decodificada esta en partitions.csv.",
            "Ver manifest.json para hashes SHA-256.",
            "",
            "Restore: no se automatiza. Revisar offsets, tabla de particiones y",
            "tamano de flash antes de escribir cualquier imagen.",
        ]
    )
    return "\n".join(lines) + "\n"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--port", help="Puerto serie; si se omite se autodetecta")
    parser.add_argument(
        "--out",
        type=Path,
        help="Directorio de salida; por defecto Documents/TimbreESPv2-backups/...",
    )
    parser.add_argument(
        "--baud",
        type=int,
        default=115200,
        help="Baud de lectura (por defecto 115200, mas seguro para CH340/hubs)",
    )
    parser.add_argument(
        "--scope",
        choices=("data", "full"),
        default="data",
        help="data: regiones criticas; full: imagen completa de 4 MiB",
    )
    parser.add_argument(
        "--include-apps",
        action="store_true",
        help="En scope=data, leer tambien app0.bin y app1.bin",
    )
    parser.add_argument(
        "--no-stub",
        action="store_true",
        help="Usar el lector ROM sin stub; mas lento pero mas robusto con cables CH340",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Mostrar puerto y regiones sin acceder al ESP32",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    os.umask(0o077)
    output_dir: Path | None = None

    try:
        pio = find_pio()
        port = args.port or detect_port(pio)
        ensure_port_free(port)
        python = find_pio_python()
        esptool = find_esptool()
        partition_tool = find_partition_tool()

        regions = list(REGIONS[:5])
        if args.include_apps:
            regions.extend(REGIONS[5:])
        if args.scope == "full":
            regions = [("flash-full.bin", 0, FULL_FLASH_SIZE)]

        print(f"Puerto: {port}")
        print(f"Scope: {args.scope}")
        print("Regiones:")
        for name, offset, size in regions:
            print(f"  {name}: 0x{offset:06X}, {size} bytes")
        if args.dry_run:
            return 0

        if args.out is None:
            stamp = dt.datetime.now().strftime("%Y%m%d-%H%M%S")
            output_dir = Path.home() / "Documents" / "TimbreESPv2-backups" / f"backup-{stamp}"
        else:
            output_dir = args.out.expanduser().resolve()
        if output_dir.exists() and any(output_dir.iterdir()):
            raise BackupError(f"El directorio de salida ya existe y no esta vacio: {output_dir}")
        output_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
        log_dir = output_dir / "logs"
        log_dir.mkdir(parents=True, exist_ok=True, mode=0o700)

        base = [
            str(python),
            str(esptool),
            "--chip",
            "esp32",
            "--port",
            port,
            "--baud",
            str(args.baud),
            "--before",
            "default_reset",
            "--after",
            "hard_reset",
        ]
        if args.no_stub:
            base.append("--no-stub")

        chip_output = run_command(base + ["chip_id"], label="chip_id", capture=True)
        flash_output = run_command(base + ["flash_id"], label="flash_id", capture=True)
        (output_dir / "device-info.txt").write_text(
            chip_output + "\n" + flash_output, encoding="utf-8"
        )
        os.chmod(output_dir / "device-info.txt", 0o600)

        region_entries: list[dict[str, object]] = []
        if args.scope == "full":
            full_path = output_dir / "flash-full.bin"
            run_command(
                base
                + [
                    "read_flash",
                    f"0x{regions[0][1]:06X}",
                    f"0x{regions[0][2]:06X}",
                    str(full_path),
                ],
                label="read-flash-full",
                log_dir=log_dir,
            )
            if full_path.stat().st_size != FULL_FLASH_SIZE:
                raise BackupError(
                    f"La imagen completa mide {full_path.stat().st_size} bytes, "
                    f"se esperaban {FULL_FLASH_SIZE}."
                )
            os.chmod(full_path, 0o600)
            region_entries = extract_regions(full_path, output_dir)
            region_entries.append(
                {
                    "file": "flash-full.bin",
                    "offset": "0x000000",
                    "size": FULL_FLASH_SIZE,
                    "sha256": sha256_file(full_path),
                }
            )
        else:
            for index, (name, offset, size) in enumerate(regions, start=1):
                target = output_dir / name
                run_command(
                    base
                    + [
                        "read_flash",
                        f"0x{offset:06X}",
                        f"0x{size:06X}",
                        str(target),
                    ],
                    label=f"read-{index:02d}-{name.removesuffix('.bin')}",
                    log_dir=log_dir,
                )
                if target.stat().st_size != size:
                    raise BackupError(
                        f"{name} mide {target.stat().st_size} bytes, se esperaban {size}."
                    )
                os.chmod(target, 0o600)
                region_entries.append(
                    {
                        "file": name,
                        "offset": f"0x{offset:06X}",
                        "size": size,
                        "sha256": sha256_file(target),
                    }
                )
                print(f"[verify] {name}: {size} bytes", flush=True)

        partition_binary = output_dir / "partitions.bin"
        partition_csv = output_dir / "partitions.csv"
        decode_error = decode_partitions(python, partition_tool, partition_binary, partition_csv)
        if decode_error is None:
            os.chmod(partition_csv, 0o600)

        manifest = {
            "schema": 1,
            "created_local": dt.datetime.now().astimezone().isoformat(),
            "created_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
            "scope": args.scope,
            "port": port,
            "baud": args.baud,
            "chip": "esp32",
            "read_only": True,
            "regions": region_entries,
            "partition_decode_error": decode_error,
        }
        manifest_path = output_dir / "manifest.json"
        manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
        os.chmod(manifest_path, 0o600)
        readme = output_dir / "README.txt"
        readme.write_text(make_readme(args.scope, port, args.baud, region_entries), encoding="utf-8")
        os.chmod(readme, 0o600)

        print(f"Backup completada: {output_dir}")
        print("Manifest: manifest.json")
        print("No se realizo ninguna escritura en el ESP32.")
        return 0
    except (BackupError, OSError, subprocess.SubprocessError) as error:
        if output_dir is not None and output_dir.exists():
            try:
                marker = output_dir / "INCOMPLETE.txt"
                marker.write_text(
                    "Backup incompleto. No usar este directorio para restaurar.\n"
                    f"Error: {error}\n",
                    encoding="utf-8",
                )
                marker.chmod(0o600)
            except OSError:
                pass
        print(f"ERROR: {error}", file=sys.stderr)
        return 1


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