mp_sync.py

#!/usr/bin/env python
"""
This script is responsible for managing and syncing files with a MicroPython board.
It provides utility functions to build, compile, and upload files to the board while handling
dependencies and maintaining a manifest of synchronized files.

The script includes functionality for:
- Identifying eligible source files for inclusion.
- Building and preparing files for upload to the MicroPython board.
- Compiling Python files into .mpy files for optimized execution.
- Managing the file sync manifest for tracking changes.
- Computing file checksums for validation.

Dependencies:
- adafruit-ampy
- click
- mpy-cross

Constants:
- DEFAULT_SERIAL_PORT: Default serial port for connecting to the board.
- BUILD_DIR: Directory to store build artifacts.
- MANIFEST_PATH: Path to the JSON manifest file used for synchronization.
- FILE_EXTENSIONS: List of file extensions eligible for inclusion.
- SRC_EXCLUSIONS: List of files or directories to exclude during file scanning.
"""

import logging
import os
import shutil
import subprocess
import sys
import time
from contextlib import suppress
from pathlib import Path
from typing import Dict, List, Set, Union, Optional
import json
import hashlib

import click
from ampy.files import DirectoryExistsError, Files
from ampy.pyboard import Pyboard

# Configure logging.
logger = logging.getLogger(__name__)
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

# Centralized configuration and well-named constants
DEFAULT_SERIAL_PORT: str = "/dev/ttyACM0"
BUILD_DIR_NAME: str = "build"
BUILD_DIR: Path = Path(BUILD_DIR_NAME)
MANIFEST_PATH: Path = Path(".sync-manifest.json")
FILE_EXTENSIONS: List[str] = [".py", ".css", ".html", ".js", ".tpl", ".mpy"]
SPECIAL_ENTRY_FILES: Set[str] = {"boot.py", "main.py"}
MPY_COMPILER: str = "mpy-cross"
HASH_CHUNK_SIZE: int = 1024 * 256

# Exclude the build directory and any other directories/files.
SRC_EXCLUSIONS: List[Union[str, Path]] = [
    "venv",
    ".venv",
    "local",
    ".",
    Path(__file__).name,
    "index_tpl.py",
    "sftp-sync.py",
    ".sync-manifest.json",
    BUILD_DIR_NAME,
]


def soft_reset(board: Pyboard) -> None:
    """Perform a soft reset of the MCU board."""
    logger.info("Sending soft reset command.")
    try:
        board.serial.write(b"\x03\x04")
    except Exception as error:
        logger.error("Failed to send soft reset command: %s", error)


def wait_for_board(delay: float = 0.5) -> None:
    """Wait briefly for the MCU device to be ready for REPL commands."""
    time.sleep(delay)


def _run_mpremote_repl() -> None:
    """Run mpremote to switch the board to REPL."""
    logger.info("Running mpremote to set REPL...")
    try:
        # Launch mpremote with the repl command and set the realtime clock to local time.
        subprocess.run(["mpremote", "rtc", "--set", "repl"], check=True)
    except subprocess.CalledProcessError as e:
        logger.error("mpremote command failed: %s", e)
        raise
    except FileNotFoundError as e:
        logger.error("mpremote command not found: %s", e)
        raise


def _should_include_file(path: Path) -> bool:
    """Decide whether a file should be included for build/sync."""
    if not path.is_file():
        return False
    if path.suffix not in FILE_EXTENSIONS:
        return False
    # Skip files that are in any of the excluded directories or names.
    if any(exc in path.parts for exc in map(str, SRC_EXCLUSIONS)):
        return False
    return True


def _ensure_clean_build_dir() -> None:
    """Ensure the build directory exists and is empty."""
    if BUILD_DIR.exists():
        try:
            shutil.rmtree(BUILD_DIR)
        except Exception as error:
            logger.error("Error removing build directory %s: %s", BUILD_DIR, error)
            raise
    BUILD_DIR.mkdir(parents=True, exist_ok=True)


def list_src_files() -> List[Path]:
    """
    List source files to be included in the build.

    Returns:
        A list of Path objects representing the source files.
    """
    src_files: List[Path] = []
    for p in Path().rglob("*"):
        if not _should_include_file(p):
            continue
        logger.debug("Found source file: %s", p)
        # Ensure the destination directory exists.
        build_path = BUILD_DIR / p
        build_path.parent.mkdir(parents=True, exist_ok=True)
        src_files.append(p)
    return src_files


def make_build() -> None:
    """Create the build directory with all eligible source files."""
    _ensure_clean_build_dir()

    # Walk through files that are not in excluded directories.
    for p in Path().rglob("*"):
        if not _should_include_file(p):
            continue
        dst_path = BUILD_DIR / p
        dst_path.parent.mkdir(parents=True, exist_ok=True)
        try:
            shutil.copy(p, dst_path)
        except Exception as error:
            logger.error("Error copying file %s to %s: %s", p, dst_path, error)
            raise


def compile_build() -> List[Path]:
    """
    Compile Python files in the build directory to .mpy files.

    Returns:
        A list of file paths (as Path objects) that are to be uploaded.
    """
    compiled_files: List[Path] = []
    for src_path in BUILD_DIR.rglob("*"):
        if not src_path.is_file():
            continue
        if src_path.suffix == ".py" and src_path.name not in SPECIAL_ENTRY_FILES:
            try:
                subprocess.check_output([MPY_COMPILER, str(src_path)], stderr=subprocess.PIPE)
                logger.debug("Compiled %s successfully.", src_path)
            except subprocess.CalledProcessError as e:
                logger.error(
                    "Compilation failed for %s. Exit code: %s. Stdout: %s. Stderr: %s",
                    src_path,
                    e.returncode,
                    e.output.decode(sys.getfilesystemencoding()),
                    e.stderr.decode(sys.getfilesystemencoding()),
                )
                # Skip this file on failure.
                continue
            except FileNotFoundError as e:
                logger.error("%s not found: %s", MPY_COMPILER, e)
                raise
            try:
                src_path.unlink()
            except Exception as error:
                logger.error("Failed to remove source file %s: %s", src_path, error)
                raise
            # Update the path to point to the compiled .mpy file.
            src_path = src_path.with_suffix(".mpy")
        compiled_files.append(src_path)
    return compiled_files


def load_manifest() -> Dict[str, Dict[str, Union[int, float, str]]]:
    """
    Load the local sync manifest mapping relative path -> {size, mtime, sha256}.
    """
    if not MANIFEST_PATH.exists():
        return {}
    try:
        with open(MANIFEST_PATH, "r", encoding="utf-8") as f:
            data = json.load(f)
            if not isinstance(data, dict):
                return {}
            return data
    except Exception as e:
        logger.warning("Failed to read manifest %s: %s. Starting fresh.", MANIFEST_PATH, e)
        return {}


def save_manifest(manifest: Dict[str, Dict[str, Union[int, float, str]]]) -> None:
    """
    Persist the manifest to disk atomically.
    """
    MANIFEST_PATH.parent.mkdir(parents=True, exist_ok=True)
    tmp = MANIFEST_PATH.with_suffix(".json.tmp")
    try:
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(manifest, f, indent=2, sort_keys=True)
        os.replace(tmp, MANIFEST_PATH)
    finally:
        with suppress(FileNotFoundError):
            tmp.unlink()


def file_sha256(path: Path, chunk_size: int = HASH_CHUNK_SIZE) -> str:
    """
    Compute SHA-256 of file content.
    """
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(chunk_size), b""):
            h.update(chunk)
    return h.hexdigest()


def make_dirs(files: Files, path: Union[str, Path], path_cache: Optional[Set[str]] = None) -> None:
    """
    Recursively create directories on the board corresponding to the given path.

    Args:
        files: The Files object used to interact with the board's filesystem.
        path: The relative path (str or Path) for which to ensure directory existence.
        path_cache: A set of already created directory paths to avoid duplicate mkdir calls.
    """
    cache: Set[str] = path_cache or set()

    relative_path = Path(str(path).replace(f"{BUILD_DIR_NAME}{os.sep}", ""))
    if relative_path.parent.as_posix() == ".":
        return

    parent_dir = relative_path.parent.as_posix()
    if parent_dir not in cache:
        logger.info("Creating directory on board: %s", parent_dir)
        make_dirs(files, parent_dir, cache)

    with suppress(DirectoryExistsError):
        posix_path = relative_path.parent.as_posix()
        if posix_path not in cache and posix_path != ".":
            try:
                files.mkdir(posix_path)
                logger.info("Created directory: %s", posix_path)
                cache.add(posix_path)
            except Exception as error:
                # Directory may already exist; avoid raising to keep sync flowing
                logger.warning("Failed to create directory %s: %s", posix_path, error)


def sync(force: bool = False, port: str = DEFAULT_SERIAL_PORT) -> None:
    """
    Main function to synchronize the local build with the MicroPython board.

    If force is True, upload all built files regardless of previous hashes.
    If force is False, only upload files whose content hash changed since last sync.
    """
    logger.info("Connecting to board on port: %s", port)
    try:
        board = Pyboard(port)
    except Exception as error:
        logger.error("Failed to connect to board on port %s: %s", port, error)
        sys.exit(1)

    files = Files(board)
    wait_for_board()

    # We still gather board files for cleanup only; not for change detection
    board_files: Dict[str, int] = {}
    orphan_dirs: List[str] = []
    try:
        for entry in files.ls(long_format=True, recursive=True):
            parts = [x.strip().replace(" bytes", "") for x in entry.split("-")]
            if len(parts) < 2:
                continue
            file_name = parts[0]
            try:
                size = int(parts[1])
            except ValueError:
                size = 0
            if file_name.endswith(tuple(FILE_EXTENSIONS)):
                board_files[file_name.lstrip("/")] = size
            elif size == 0:
                orphan_dirs.append(file_name)
    except Exception as error:
        logger.error("Error processing board file list: %s", error)
        sys.exit(1)

    try:
        make_build()
    except Exception as error:
        logger.error("Error building local files: %s", error)
        sys.exit(1)

    to_upload: List[Path] = compile_build()

    # Load previous manifest at once
    previous_manifest = load_manifest()
    updated_manifest = dict(previous_manifest)

    # Remove orphan directories on the board.
    for o in orphan_dirs:
        if o != "/":
            logger.info("Deleting orphan directory: %s", o)
            with suppress(Exception):
                files.rmdir(o)

    # Delete files on the board that are not in the local build.
    for bf in list(board_files.keys()):
        if (BUILD_DIR / bf) not in to_upload:
            logger.info("Deleting file not present locally: %s", bf)
            with suppress(Exception):
                files.rm(f"/{bf}")
            # Also drop from manifest
            updated_manifest.pop(bf, None)

    # Upload files (hash based on actual upload bytes)
    dir_cache: Set[str] = set()

    for local_file in to_upload:
        rel = local_file.as_posix().replace(f"{BUILD_DIR_NAME}{os.sep}", "")
        if rel == ".":
            continue

        try:
            # Read the exact content we will upload and hash that
            with open(local_file, "rb") as fd:
                content = fd.read()
        except Exception as e:
            logger.error("Failed to read %s: %s", local_file, e)
            continue

        content_hash = hashlib.sha256(content).hexdigest()
        stat = local_file.stat()
        current_meta = {"size": stat.st_size, "mtime": stat.st_mtime, "sha256": content_hash}
        prev_meta = previous_manifest.get(rel)

        if not force and isinstance(prev_meta, dict) and prev_meta.get("sha256") == content_hash:
            # Unchanged payload; skip upload
            logger.debug("Skipping unchanged file: %s", rel)
            updated_manifest[rel] = current_meta  # refresh size/mtime if they differ locally
            continue

        # Ensure destination directories exist
        try:
            make_dirs(files, local_file, dir_cache)
        except Exception as error:
            logger.error("Error creating directories for %s: %s", local_file, error)
            continue

        logger.info("Uploading file: %s%s", rel, " (force)" if force else "")
        try:
            files.put(rel, content)
            updated_manifest[rel] = current_meta
        except Exception as error:
            logger.error("Failed to upload file %s: %s", rel, error)
            # Do not update manifest on failure
            continue

    # Prune entries that no longer exist in build/
    for path_in_manifest in list(updated_manifest.keys()):
        if not (BUILD_DIR / path_in_manifest).exists():
            updated_manifest.pop(path_in_manifest, None)

    save_manifest(updated_manifest)
    soft_reset(board)


def purge_board(port: str = DEFAULT_SERIAL_PORT) -> None:
    """
    Purges all files and directories from the MicroPython board connected on the given port.
    """
    logger.info("Purging all files from board on port: %s", port)
    try:
        board = Pyboard(port)
    except Exception as error:
        logger.error("Failed to connect to board on port %s: %s", port, error)
        return

    files = Files(board)
    wait_for_board()

    # Gather all entries
    entries = files.ls(long_format=True, recursive=True)

    file_paths: List[str] = []
    dir_paths: List[str] = []

    for entry in entries:
        parts = [x.strip().replace(" bytes", "") for x in entry.split("-")]
        if len(parts) < 2:
            continue
        path, size_str = parts[0], parts[1]
        try:
            size_int = int(size_str)
        except ValueError:
            continue

        # Anything with a non-zero size (or matching known extensions) is a file
        if size_int > 0 or path.endswith(tuple(FILE_EXTENSIONS)):
            file_paths.append(path)
        else:
            dir_paths.append(path)

    # First, remove all files
    for p in file_paths:
        if p in ("/", ""):
            continue
        logger.info("Deleting file: %s", p)
        with suppress(Exception):
            files.rm(p)

    # Then remove directories, deepest first
    for d in sorted(dir_paths, key=lambda p: p.count("/"), reverse=True):
        if d in ("/", "", "."):
            continue
        logger.info("Removing directory: %s", d)
        with suppress(Exception):
            files.rmdir(d)

    # Soft-reset to apply changes
    soft_reset(board)


# ---------------------------
# CLI using Click
# ---------------------------


def _set_verbosity(verbose: int) -> None:
    # 0 -> INFO (default), 1 -> DEBUG
    if verbose > 0:
        logging.getLogger().setLevel(logging.DEBUG)
        for h in logging.getLogger().handlers:
            h.setLevel(logging.DEBUG)


@click.group(context_settings=dict(help_option_names=["-h", "--help"]))
@click.option("--port", "-p", default=DEFAULT_SERIAL_PORT, show_default=True, help="Serial port of the board.")
@click.option("--verbose", "-v", count=True, help="Increase verbosity (-v for DEBUG).")
@click.pass_context
def cli(ctx: click.Context, port: str, verbose: int) -> None:
    """
    MicroPython MCU sync utility.
    """
    _set_verbosity(verbose)
    ctx.obj = {"port": port}


@cli.command("list-src")
def cli_list_src() -> None:
    """List local source files considered for the build."""
    files = list_src_files()
    for f in files:
        click.echo(f)


@cli.command("build")
def cli_build() -> None:
    """Build a local artifact directory (copy and compile)."""
    make_build()
    compiled = compile_build()
    click.echo(f"Prepared {len(compiled)} files under {BUILD_DIR}")


@cli.command("sync")
@click.option("--force", is_flag=True, help="Upload all files regardless of manifest.")
@click.option(
    "--auto-build/--no-auto-build", default=True, show_default=True, help="Automatically build before syncing."
)
@click.option(
    "--mpremote/--no-mpremote", default=False, show_default=True, help="Run mpremote to set REPL after successful sync."
)
@click.pass_context
def cli_sync(ctx: click.Context, force: bool, auto_build: bool, mpremote: bool) -> None:
    """Synchronize local build to the board."""
    port: str = ctx.obj["port"]

    # Optional automatic build
    if auto_build:
        make_build()
        compile_build()

    # Perform sync
    sync(force=force, port=port)
    click.echo("Sync complete.")

    # Optional mpremote step after successful sync
    if mpremote:
        try:
            time.sleep(1)
            _run_mpremote_repl()
        except Exception:
            sys.exit(1)


@cli.command("purge")
@click.pass_context
def cli_purge(ctx: click.Context) -> None:
    """Purge all files on the board."""
    port: str = ctx.obj["port"]
    purge_board(port=port)
    click.echo("Purge complete.")


if __name__ == "__main__":
    # Default behavior when run directly without subcommand: run CLI group
    cli()