#!/usr/bin/env python3
# (c) 2024-26 Prof. Flavio ABREU ARAUJO. All rights reserved.
#
# OTA Update Client for NEnG CircuitPython Instruments
# Sends firmware updates to devices via SCPI-over-TCP

import base64
import contextlib
import hashlib
import os
import socket
import subprocess
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional


class OTAClient:
    """
    Client for sending OTA (Over-The-Air) updates to CircuitPython instruments.

    Supports:
    - Secure authenticated updates
    - Checksum verification (SHA256, CRC32)
    - Progress tracking
    - Automatic rollback on failure
    - TCP connection management

    Usage:
        client = OTAClient("192.168.1.100", password="mypassword")
        success = client.upload_file("code_pid_async.py", "/path/to/new/code.py")
        if success:
            client.reboot_device()
    """

    def __init__(
        self,
        host: str,
        port: int = 5025,
        password: Optional[str] = None,
        timeout: int = 30,
        chunk_size: int = 8192,  # Increased from 1024 for faster uploads
    ):
        """
        Initialize OTA client.

        Args:
            host: Device IP address or hostname
            port: SCPI-over-TCP port (default: 5025)
            password: OTA authentication password
            timeout: Socket timeout in seconds
            chunk_size: Data chunk size in bytes (default: 8192, optimized for ESP32)
        """
        self.host = host
        self.port = port
        self.password = password
        self.timeout = timeout
        self.chunk_size = chunk_size
        self.sock = None
        self._authenticated = False

    def connect(self):
        """Connect to device"""
        if self.sock:
            self.close()

        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.settimeout(self.timeout)
            self.sock.connect((self.host, self.port))
            print(f"Connected to {self.host}:{self.port}")
            return True
        except Exception as e:
            print(f"Connection failed: {e}")
            return False

    def close(self):
        """Close connection"""
        if self.sock:
            with contextlib.suppress(Exception):
                self.sock.close()
            self.sock = None
        self._authenticated = False

    def send_command(self, command: str, expect_response: bool = True) -> Optional[str]:
        """
        Send SCPI command and optionally read response.

        Args:
            command: SCPI command string
            expect_response: Whether to read response (queries need this)

        Returns:
            Response string if expect_response, None otherwise
        """
        if not self.sock:
            raise RuntimeError("Not connected - call connect() first")

        try:
            # Send command
            cmd_bytes = (command.strip() + "\n").encode("utf-8")
            self.sock.sendall(cmd_bytes)

            # Read response if expected
            if expect_response:
                response = b""
                while True:
                    chunk = self.sock.recv(4096)
                    if not chunk:
                        break
                    response += chunk
                    if b"\n" in chunk:
                        break

                return response.decode("utf-8").strip()

            return None

        except socket.timeout as e:
            raise RuntimeError("Command timeout - device may be unresponsive") from e
        except Exception as e:
            raise RuntimeError(f"Command failed: {e}") from e

    def authenticate(self) -> bool:
        """Authenticate for OTA updates"""
        if not self.password:
            print("Warning: No password set - authentication may fail")
            return True

        try:
            response = self.send_command(f":SYST:UPD:AUTH {self.password}")
            if response and "OK" in response:
                print("Authentication successful")
                self._authenticated = True
                return True
            else:
                print(f"Authentication failed: {response}")
                return False
        except Exception as e:
            print(f"Authentication error: {e}")
            return False

    def calculate_checksum(self, file_path: Path, algo: str = "sha256") -> str:
        """
        Calculate file checksum.

        Args:
            file_path: Path to file
            algo: Algorithm ('sha256' or 'crc32')

        Returns:
            Checksum as hex string
        """
        if algo == "sha256":
            h = hashlib.sha256()
            with open(file_path, "rb") as f:
                while chunk := f.read(8192):
                    h.update(chunk)
            return h.hexdigest()

        elif algo == "crc32":
            import binascii

            crc = 0
            with open(file_path, "rb") as f:
                while chunk := f.read(8192):
                    crc = binascii.crc32(chunk, crc)
            return f"{crc & 0xFFFFFFFF:08x}"

        else:
            raise ValueError(f"Unsupported algorithm: {algo}")

    def upload_file(
        self,
        target_filename: str,
        local_path: str,
        algo: str = "sha256",
        auto_reboot: bool = False,
        verify_first: bool = True,
        verbose: bool = True,
    ) -> bool:
        """
        Upload file to device via OTA.

        Args:
            target_filename: Target filename on device (e.g., 'code_pid_async.py')
            local_path: Local file path to upload
            algo: Checksum algorithm ('sha256' or 'crc32')
            auto_reboot: Automatically reboot device after successful update
            verify_first: Verify connection before starting update
            verbose: Show detailed progress (default: True)

        Returns:
            True if upload successful, False otherwise
        """
        local_file = Path(local_path)

        if not local_file.exists():
            print(f"Error: File not found: {local_path}")
            return False

        # Get file info
        file_size = local_file.stat().st_size
        if verbose:
            print(f"Preparing to upload {local_file.name} ({file_size} bytes) as {target_filename}")

        # Calculate checksum
        if verbose:
            print(f"Calculating {algo} checksum...")
        checksum = self.calculate_checksum(local_file, algo)
        if verbose:
            print(f"Checksum: {checksum}")

        # Connect if not already connected
        if not self.sock and not self.connect():
            return False

        # Verify connection
        if verify_first:
            try:
                idn = self.send_command("*IDN?")
                if verbose:
                    print(f"Device: {idn}")
            except Exception as e:
                print(f"Connection verification failed: {e}")
                return False

        # Authenticate
        if not self._authenticated and not self.authenticate():
            return False

        try:
            # Start update
            if verbose:
                print("Starting OTA update...")
            start_cmd = f":SYST:UPD:START {target_filename},{file_size},{checksum},{algo}"
            response = self.send_command(start_cmd)
            if verbose:
                print(f"Device response: {response}")

            if not response or "ERROR" in response:
                if not verbose:
                    print(f"  Failed to start: {response}")
                return False

            # Send file in chunks
            if verbose:
                print("Uploading file...")
            bytes_sent = 0
            start_time = time.time()
            last_progress_pct = 0

            with open(local_file, "rb") as f:
                while True:
                    chunk = f.read(self.chunk_size)
                    if not chunk:
                        break

                    # Encode as base64
                    encoded = base64.b64encode(chunk).decode("ascii")

                    # Send chunk
                    response = self.send_command(f":SYST:UPD:DATA {encoded}")

                    if not response or "ERROR" in response:
                        if verbose:
                            print(f"\nUpload failed: {response}")
                        self.send_command(":SYST:UPD:ABORT", expect_response=False)
                        return False

                    bytes_sent += len(chunk)

                    # Only show progress in verbose mode
                    if verbose:
                        progress = (bytes_sent / file_size) * 100
                        # Only print progress every 5% or on completion (reduces overhead)
                        if progress >= last_progress_pct + 5 or bytes_sent == file_size:
                            elapsed = time.time() - start_time
                            rate = bytes_sent / elapsed if elapsed > 0 else 0
                            print(
                                f"  Progress: {bytes_sent}/{file_size} bytes "
                                f"({progress:.1f}%) @ {rate / 1024:.1f} KB/s",
                                end="\r",
                            )
                            last_progress_pct = int(progress / 5) * 5  # Round down to nearest 5%

            if verbose:
                print()  # Newline after progress

            # Commit update
            if verbose:
                print("Committing update (verifying checksum)...")
            response = self.send_command(":SYST:UPD:COMMIT")
            if verbose:
                print(f"Device response: {response}")

            if not response or "ERROR" in response:
                if not verbose:
                    print(f"  Commit failed: {response}")
                return False

            if verbose:
                duration = time.time() - start_time
                print(f"✓ Upload successful in {duration:.1f}s")
                print("Update is staged - reboot device to apply changes")

            # Auto-reboot if requested
            if auto_reboot:
                print("Rebooting device...")
                self.reboot_device()

            return True

        except Exception as e:
            print(f"Upload error: {e}")
            with contextlib.suppress(Exception):
                self.send_command(":SYST:UPD:ABORT", expect_response=False)
            return False

    def query_update_status(self) -> dict:
        """
        Query current update status.

        Returns:
            Dictionary with status information
        """
        try:
            response = self.send_command(":SYST:UPD:STATUS?")
            if not response:
                return {}

            # Parse key:value pairs
            status = {}
            for pair in response.split(","):
                if ":" in pair:
                    key, value = pair.split(":", 1)
                    status[key] = value

            return status

        except Exception as e:
            print(f"Status query failed: {e}")
            return {}

    def rollback(self) -> bool:
        """
        Trigger rollback to previous version.

        Returns:
            True if rollback successful
        """
        if not self._authenticated and not self.authenticate():
            return False

        try:
            response = self.send_command(":SYST:UPD:ROLLBACK")
            print(f"Rollback response: {response}")
            return response and "OK" in response

        except Exception as e:
            print(f"Rollback error: {e}")
            return False

    def delete_file(self, filename: str, verbose: bool = True) -> bool:
        """
        Delete file from device.

        Args:
            filename: File path to delete (e.g., "lib/old_module.py")
            verbose: Show detailed output

        Returns:
            True if deletion successful, False otherwise
        """
        # Ensure filename starts with /
        if not filename.startswith("/"):
            filename = "/" + filename

        if not self.sock and not self.connect():
            return False

        # Authenticate if needed
        if not self._authenticated and not self.authenticate():
            return False

        try:
            if verbose:
                print(f"Deleting {filename}...")
            response = self.send_command(f":SYST:FILE:DELETE {filename}")

            if not response or "ERROR" in response:
                if verbose:
                    print(f"  Failed: {response}")
                return False

            if verbose:
                print(f"  {response}")
            return True

        except Exception as e:
            if verbose:
                print(f"  Error: {e}")
            return False

    def reboot_device(self):
        """Reboot device (apply updates)"""
        try:
            print("Sending reboot command...")
            self.send_command(":SYST:RES", expect_response=False)
            print("Device is rebooting...")
            time.sleep(1)
            self.close()
        except Exception as e:
            print(f"Reboot command sent (connection closed: {e})")
            self.close()

    def __enter__(self):
        """Context manager entry"""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit"""
        self.close()


class GitToDeviceMapper:
    """Maps git repository paths to CircuitPython device paths.

    Supports all NEnG firmware projects.  Pass *device* to select which
    CircuitPython project and shared libraries are included.

    Supported device values:
        ``TMP117``, ``3Dmag``, ``BTS7960``, ``RelayBank16``, ``RTDMax31865``
    """

    # Per-device configuration: CircuitPython base dir + shared libraries
    DEVICE_PROFILES: dict[str, dict] = {
        "TMP117": {
            "source_base": "TMP117/CircuitPython",
            "libraries": ["neng_scpi_base", "neng_tmp117", "neng_pid_controller", "neng_bts7960"],
        },
        "3Dmag": {
            "source_base": "3Dmag/CircuitPython",
            "libraries": ["neng_scpi_base", "neng_3dmag"],
        },
        "BTS7960": {
            "source_base": "BTS7960/CircuitPython",
            "libraries": ["neng_scpi_base", "neng_bts7960"],
        },
        "RelayBank16": {
            "source_base": "RelayBank16/CircuitPython",
            "libraries": ["neng_scpi_base", "neng_relaybank16"],
        },
        "RTDMax31865": {
            "source_base": "RTDMax31865/CircuitPython",
            "libraries": ["neng_scpi_base", "neng_rtdmax31865"],
        },
    }

    # Exclusion patterns (files that should never be deployed via OTA)
    EXCLUDE_PATTERNS = [
        ".md",
        "__pycache__",
        ".pyc",
        "boot_out.txt",
        "log.txt",
        ".DS_Store",
        "help/",
        "backup/",
        "update/",
        ".Trashes",
        ".fseventsd",
    ]

    def __init__(self, device: str = "TMP117") -> None:
        device = device.strip()
        if device not in self.DEVICE_PROFILES:
            raise ValueError(
                f"Unknown device '{device}'. Supported: {', '.join(sorted(self.DEVICE_PROFILES))}"
            )
        profile = self.DEVICE_PROFILES[device]
        self.source_base: str = profile["source_base"]
        self.library_sources: list[str] = profile["libraries"]

        # Symlink directories in lib/ that should be excluded
        # (use source files from neng_* root dirs instead)
        self.symlink_dirs: list[str] = [
            f"{self.source_base}/lib/{lib}" for lib in self.library_sources
        ]

    def git_to_device_path(self, git_path: str) -> Optional[str]:
        """
        Convert git file path to CircuitPython device path.

        Args:
            git_path: Path from git diff (e.g., 'neng_scpi_base/ota_scpi_server.py')

        Returns:
            Device path (e.g., '/lib/neng_scpi_base/ota_scpi_server.py') or None if excluded

        Examples (device="TMP117"):
            TMP117/CircuitPython/code.py → /code.py
            neng_scpi_base/ota_scpi_server.py → /lib/neng_scpi_base/ota_scpi_server.py
            TMP117/CircuitPython/README.md → None (excluded)
        """
        # Application files (<Device>/CircuitPython/*)
        if git_path.startswith(f"{self.source_base}/"):
            rel_path = git_path[len(self.source_base) + 1 :]  # Remove prefix

            # Check exclusion patterns
            if self._is_excluded(rel_path):
                return None

            return f"/{rel_path}"

        # Library files (neng_*/*)
        for lib in self.library_sources:
            if git_path.startswith(f"{lib}/"):
                # Check exclusion patterns
                if self._is_excluded(git_path):
                    return None

                # Map to /lib/<library>/<file>
                return f"/lib/{git_path}"

        # Not a deployable file
        return None

    def _is_excluded(self, path: str) -> bool:
        """Check if file should be excluded from OTA deployment."""
        return any(pattern in path for pattern in self.EXCLUDE_PATTERNS)

    def get_deployable_files(
        self, git_paths: list[str], for_deletion: bool = False, repo_root: Optional[Path] = None
    ) -> dict[str, str]:
        """
        Filter and map list of git paths to device paths.

        Args:
            git_paths: List of file paths from git diff (relative to repo root)
            for_deletion: If True, skip existence checks (for deleted files)
            repo_root: Absolute path to repo root; used to resolve git_paths to
                       real filesystem paths for existence/directory checks.

        Returns:
            Dict mapping device_path -> local_path (absolute) for files to upload/delete
        """
        result = {}
        for git_path in git_paths:
            # Skip symlink directories in lib/ (we use source files from neng_* instead)
            if git_path in self.symlink_dirs:
                continue

            # Resolve to an absolute path for filesystem checks
            local_path = str(Path(repo_root) / git_path) if repo_root else git_path

            # Skip files that don't exist locally (but not when processing deletions)
            if not for_deletion and not os.path.exists(local_path):
                continue

            # Skip directories (but not files in those directories)
            if not for_deletion and os.path.isdir(local_path):
                continue

            device_path = self.git_to_device_path(git_path)
            if device_path:
                result[device_path] = local_path

        return result


class GitDeployManager:
    """Manages git-aware OTA deployments."""

    def __init__(self, repo_path: str = ".", device: str = "TMP117"):
        # Auto-detect the actual git repo root so pathspecs work correctly
        # regardless of which directory the GUI / CLI was launched from.
        try:
            result = subprocess.run(
                ["git", "rev-parse", "--show-toplevel"],
                cwd=repo_path,
                capture_output=True,
                text=True,
                check=True,
            )
            self.repo_path = Path(result.stdout.strip())
        except subprocess.CalledProcessError:
            self.repo_path = Path(repo_path).resolve()
        self.device = device
        self.mapper = GitToDeviceMapper(device=device)

    def get_current_commit(self) -> str:
        """Get current HEAD commit hash."""
        result = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=self.repo_path,
            capture_output=True,
            text=True,
            check=True,
        )
        return result.stdout.strip()

    def _git_pathspecs(self) -> list[str]:
        """Return the git pathspec list for the current device profile."""
        specs = [f"{self.mapper.source_base}/"]
        specs.extend(f"{lib}/" for lib in self.mapper.library_sources)
        return specs

    def get_changed_files(
        self, from_commit: str, to_commit: str = "HEAD"
    ) -> tuple[list[str], list[str]]:
        """
        Get list of changed and deleted files between commits.

        Args:
            from_commit: Base commit hash
            to_commit: Target commit hash (default: HEAD)

        Returns:
            Tuple of (changed_files, deleted_files)
        """
        # Get all changes with status
        result = subprocess.run(
            [
                "git",
                "diff",
                "--name-status",
                from_commit,
                to_commit,
                "--",
                *self._git_pathspecs(),
            ],
            cwd=self.repo_path,
            capture_output=True,
            text=True,
            check=True,
        )

        changed_files = []
        deleted_files = []

        for line in result.stdout.strip().split("\n"):
            if not line:
                continue

            parts = line.split("\t", 1)
            if len(parts) != 2:
                continue

            status, filepath = parts
            if status.startswith("D"):
                # File was deleted
                deleted_files.append(filepath)
            else:
                # File was added or modified
                changed_files.append(filepath)

        return changed_files, deleted_files

    def check_working_tree_clean(self) -> tuple[bool, list[str]]:
        """
        Check if working tree has uncommitted changes.

        Returns:
            Tuple of (is_clean, list_of_modified_files)
        """
        result = subprocess.run(
            [
                "git",
                "status",
                "--porcelain",
                "--",
                *self._git_pathspecs(),
            ],
            cwd=self.repo_path,
            capture_output=True,
            text=True,
        )
        modified = [line for line in result.stdout.strip().split("\n") if line]
        return len(modified) == 0, modified

    def generate_version_file(self, commit_hash: str, dirty: bool = False) -> str:
        """
        Generate version.txt content.

        Args:
            commit_hash: Full git commit hash
            dirty: Whether working tree had uncommitted changes

        Returns:
            Content string for version.txt file
        """
        timestamp = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
        deployer = f"{os.getenv('USER', 'unknown')}@{socket.gethostname()}"

        return f"""commit={commit_hash}
timestamp={timestamp}
deployer={deployer}
dirty={str(dirty).lower()}
"""


def deploy_from_git(
    client: OTAClient,
    from_commit: Optional[str] = None,
    force_full: bool = False,
    dry_run: bool = False,
    auto_reboot: bool = False,
    algo: str = "crc32",
    auto_confirm: bool = False,
    device: str = "TMP117",
) -> bool:
    """
    Deploy firmware using git-aware incremental update.

    Args:
        client: OTA client (already connected)
        from_commit: Source commit (None = query device)
        force_full: Force full deployment (ignore git diff)
        dry_run: Show what would be uploaded without uploading
        auto_confirm: Skip confirmation prompt (use with caution)
        auto_reboot: Reboot device after successful deployment
        algo: Checksum algorithm (default: crc32 for CircuitPython)
        device: Target device type (TMP117, 3Dmag, BTS7960, RelayBank16, RTDMax31865)

    Returns:
        True if deployment successful
    """
    git_mgr = GitDeployManager(device=device)
    print(f"Device profile: {device} ({git_mgr.mapper.source_base})")

    try:
        # 1. Get current HEAD
        local_commit = git_mgr.get_current_commit()
        print(f"Local HEAD: {local_commit[:7]}")
    except subprocess.CalledProcessError as e:
        print(f"Error: Failed to get git commit - {e}")
        print("Make sure you're in a git repository")
        return False

    # 2. Check working tree
    is_clean, modified = git_mgr.check_working_tree_clean()
    if not is_clean:
        print("\n⚠️  Warning: Working tree has uncommitted changes:")
        for mod in modified[:5]:  # Show first 5
            print(f"  {mod}")
        if len(modified) > 5:
            print(f"  ... and {len(modified) - 5} more")
        print("These changes will be marked as dirty=true in version.txt\n")

    # 3. Get device commit (if not forcing full)
    device_commit = None
    if not force_full and not from_commit:
        try:
            device_commit = client.send_command(":SYST:VERS:GIT?")
            if device_commit and device_commit != "unknown":
                print(f"Device version: {device_commit[:7]}")
            else:
                print("Device has no git version info (first deployment)")
                device_commit = None
        except Exception as e:
            print(f"Warning: Could not query device version: {e}")
            device_commit = None

    # 4. Determine what files to upload
    upload_map = {}
    delete_map = {}  # Files to delete (empty for full deployment)

    if force_full or (not from_commit and not device_commit):
        print("\n📦 Full deployment mode - uploading all tracked files")
        # Get all tracked files in deployment areas
        try:
            result = subprocess.run(
                [
                    "git",
                    "ls-files",
                    *git_mgr._git_pathspecs(),
                ],
                cwd=git_mgr.repo_path,
                capture_output=True,
                text=True,
                check=True,
            )
            all_files = [line for line in result.stdout.strip().split("\n") if line]
            upload_map = git_mgr.mapper.get_deployable_files(all_files, repo_root=git_mgr.repo_path)
        except subprocess.CalledProcessError as e:
            print(f"Error: git ls-files failed - {e}")
            return False

    else:
        # Incremental deployment
        base_commit = from_commit or device_commit
        print(f"\n🔍 Computing changes from {base_commit[:7]} to {local_commit[:7]}...")

        # Get changed and deleted files
        try:
            changed_files, deleted_files = git_mgr.get_changed_files(base_commit, local_commit)
        except subprocess.CalledProcessError as e:
            print(f"Error: git diff failed - {e}")
            print("This may happen if commits have no common ancestor (different branches/repos)")
            print("Try --force-full or --from-commit with a valid commit hash")
            return False

        if not changed_files and not deleted_files:
            print("✓ No changes detected - deployment not needed")
            print(f"Device is already at {local_commit[:7]}")
            return True

        print(
            f"Found {len(changed_files)} changed file(s)"
            + (f", {len(deleted_files)} deleted file(s)" if deleted_files else "")
        )
        upload_map = git_mgr.mapper.get_deployable_files(changed_files, repo_root=git_mgr.repo_path)
        delete_map = git_mgr.mapper.get_deployable_files(
            deleted_files, for_deletion=True, repo_root=git_mgr.repo_path
        )

    if not upload_map and not delete_map:
        print("✓ No deployable changes after filtering")
        return True

    # 5. Show upload plan
    if upload_map:
        print(f"\n📤 Files to upload ({len(upload_map)}):")
        total_size = 0
        for device_path, local_path in sorted(upload_map.items()):
            size = os.path.getsize(local_path)
            total_size += size
            size_kb = size / 1024
            print(f"  {device_path:50s} ({size_kb:6.1f} KB)")

        print(f"\nTotal upload size: {total_size / 1024:.1f} KB")
    else:
        total_size = 0

    if delete_map:
        print(f"\n🗑️  Files to delete ({len(delete_map)}):")
        for device_path in sorted(delete_map.keys()):
            print(f"  {device_path:50s}")

    if dry_run:
        print("\n🔍 [DRY RUN] Would upload these files (not actually uploading)")
        return True

    # 6. Confirm
    if not auto_confirm:
        try:
            response = input("\n🚀 Continue with deployment? [y/N] ")
            if not response.lower().startswith("y"):
                print("Deployment cancelled")
                return False
        except (EOFError, KeyboardInterrupt):
            print("\nDeployment cancelled")
            return False
    else:
        print("\n🚀 Auto-confirming deployment (--yes flag)\n")

    # 7. Upload files (batch mode with quiet output)
    print("\n" + "=" * 60)
    print("Starting OTA Upload (Batch Mode)")
    print("=" * 60 + "\n")

    success_count = 0
    failed_files = []
    start_time = time.time()

    for idx, (device_path, local_path) in enumerate(sorted(upload_map.items()), 1):
        # Extract filename from device path
        target_filename = device_path.lstrip("/")

        # Quiet mode: just show filename and status on one line
        print(f"[{idx}/{len(upload_map)}] {target_filename:50s} ... ", end="", flush=True)

        success = client.upload_file(
            target_filename=target_filename,
            local_path=local_path,
            algo=algo,
            auto_reboot=False,
            verify_first=(idx == 1),  # Only verify first file
            verbose=False,  # Quiet mode for batch upload
        )

        if success:
            success_count += 1
            print("✓ uploaded")
        else:
            failed_files.append(target_filename)
            print("✗ failed")

    # 8. Delete files (if any)
    deleted_count = 0
    if delete_map and success_count == len(upload_map):
        print("\n" + "=" * 60)
        print("Deleting Removed Files")
        print("=" * 60 + "\n")

        for idx, device_path in enumerate(sorted(delete_map.keys()), 1):
            target_filename = device_path.lstrip("/")
            print(
                f"[{idx}/{len(delete_map)}] {target_filename:50s} ... ",
                end="",
                flush=True,
            )

            success = client.delete_file(target_filename, verbose=False)

            if success:
                deleted_count += 1
                print("✓ deleted")
            else:
                failed_files.append(f"{target_filename} (delete)")
                print("✗ failed")

    # 9. Upload version.txt (only if all files succeeded)
    if success_count == len(upload_map) and (not delete_map or deleted_count == len(delete_map)):
        print(
            f"[{len(upload_map) + 1}/{len(upload_map) + 1}] version.txt{' ' * 40} ... ",
            end="",
            flush=True,
        )

        # Generate version file content
        version_content = git_mgr.generate_version_file(local_commit, dirty=not is_clean)

        # Write to temporary file
        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as f:
            f.write(version_content)
            temp_version_path = f.name

        try:
            success = client.upload_file(
                target_filename="version.txt",
                local_path=temp_version_path,
                algo=algo,
                auto_reboot=False,
                verify_first=False,
                verbose=False,
            )

            if success:
                print("✓ uploaded")
            else:
                print("✗ failed")
        finally:
            os.unlink(temp_version_path)
    else:
        print("\n⚠️  Skipping version.txt update (some files failed)\n")

    # 10. Summary
    total_time = time.time() - start_time
    print("\n" + "=" * 60)
    print("Deployment Summary:")
    print(f"  Uploaded: {success_count}/{len(upload_map)} files")
    if delete_map:
        print(f"  Deleted: {deleted_count}/{len(delete_map)} files")
    print(f"  Total time: {total_time:.1f}s")
    if total_size > 0:
        print(f"  Average: {total_size / 1024 / total_time:.1f} KB/s")
    if failed_files:
        print(f"  Failed: {', '.join(failed_files)}")
    print(f"  New version: {local_commit[:7]}")
    if not is_clean:
        print("  ⚠️  Deployed with uncommitted changes (dirty=true)")
    print("=" * 60 + "\n")

    # 11. Reboot if requested
    all_succeeded = success_count == len(upload_map) and (
        not delete_map or deleted_count == len(delete_map)
    )
    if auto_reboot and all_succeeded:
        print("🔄 Rebooting device to apply changes...")
        client.reboot_device()
        print("✓ Git-aware OTA deployment completed")
    elif success_count == len(upload_map):
        print("✓ Deployment successful - reboot device to apply changes")
    elif failed_files:
        print("✗ Deployment completed with errors")
    else:
        print("✓ Deployment successful - reboot device to apply changes")

    return all_succeeded


def main():
    """CLI entry point for OTA updates"""
    import argparse

    parser = argparse.ArgumentParser(
        description="OTA Update Client for NEnG Instruments",
        epilog="Examples:\n"
        "  Auto-detect:     update-firmware --git-deploy --password 'pass' --reboot\n"
        "  Single file:     update-firmware 192.168.1.100 code.py --password 'pass' --reboot\n"
        "  Batch update:    update-firmware 192.168.1.100 code.py boot.py --password 'pass' --reboot\n"
        "  Git (TMP117):    update-firmware 192.168.1.100 --git-deploy --password 'pass' --reboot\n"
        "  Git (3Dmag):     update-firmware 192.168.1.100 --git-deploy --device 3Dmag --reboot\n"
        "  Git (any):       update-firmware 192.168.1.100 --git-deploy --device RTDMax31865 --reboot\n"
        "  Git dry-run:     update-firmware 192.168.1.100 --git-deploy --dry-run",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "host",
        nargs="?",
        help="Device IP address or hostname (auto-detects if not specified)",
    )
    parser.add_argument(
        "files",
        nargs="*",
        help="Local file(s) to upload (not used with --git-deploy)",
    )
    parser.add_argument(
        "--target",
        help="Target filename on device (only for single file updates)",
    )
    parser.add_argument("--password", default=None, help="OTA authentication password")
    parser.add_argument("--port", type=int, default=5025, help="SCPI port (default: 5025)")
    parser.add_argument(
        "--algo",
        choices=["sha256", "crc32"],
        default="crc32",
        help="Checksum algorithm (crc32 is faster on CircuitPython)",
    )
    parser.add_argument("--reboot", action="store_true", help="Auto-reboot after all updates")
    parser.add_argument(
        "--chunk-size",
        type=int,
        default=8192,
        help="Chunk size in bytes (default: 8192, optimized for WiFi)",
    )

    # Git deployment arguments
    parser.add_argument(
        "--git-deploy",
        action="store_true",
        help="Git-aware deployment (upload only changed files)",
    )
    parser.add_argument(
        "--from-commit",
        help="Base commit for git diff (default: query device)",
    )
    parser.add_argument(
        "--force-full",
        action="store_true",
        help="Force full deployment (ignore git diff, upload all files)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be uploaded without uploading",
    )
    parser.add_argument(
        "--yes",
        action="store_true",
        help="Auto-confirm deployment without prompting (git-deploy only)",
    )
    parser.add_argument(
        "--scan-network",
        metavar="CIDR",
        help="Network to scan for devices (e.g., 10.104.32.0/24). Auto-detects if not specified.",
    )
    parser.add_argument(
        "--device",
        choices=sorted(GitToDeviceMapper.DEVICE_PROFILES),
        default="TMP117",
        help="Target device type for git deploy (default: TMP117). "
        "Supported: " + ", ".join(sorted(GitToDeviceMapper.DEVICE_PROFILES)),
    )

    args = parser.parse_args()

    # Auto-detect device if host not specified
    if not args.host:
        print("🔍 No host specified - scanning network for devices...")
        import concurrent.futures

        from .device_scanner import get_all_local_networks, scan_network

        # Use specified network or auto-detect all networks
        if args.scan_network:
            networks = [args.scan_network]
            print(f"   Scanning {args.scan_network}...")
        else:
            networks = get_all_local_networks()
            if not networks:
                print("✗ Could not auto-detect local networks")
                print("  Please specify device IP address or use --scan-network")
                return 1
            print(f"   Auto-detected {len(networks)} network(s):")
            for net in networks:
                print(f"     • {net}")
            print("   Scanning all networks...")

        # Scan all networks concurrently
        devices = []
        with concurrent.futures.ThreadPoolExecutor(max_workers=len(networks)) as executor:
            future_to_network = {
                executor.submit(scan_network, network=net, verbose=False): net for net in networks
            }

            for future in concurrent.futures.as_completed(future_to_network):
                try:
                    result = future.result()
                    devices.extend(result)
                except Exception:
                    # Silently ignore scanning errors for individual networks
                    pass

        if not devices:
            print("✗ No devices found on any network")
            print(f"  Scanned: {', '.join(networks)}")
            print("  Please check:")
            print("    • Device is powered on")
            print("    • Device is connected to WiFi")
            print("    • Device is on the same network as your computer")
            return 1

        if len(devices) == 1:
            args.host = devices[0][0]
            device_id = devices[0][1]
            print(f"✓ Found 1 device: {args.host}")
            print(f"  {device_id}")
        else:
            print(f"✓ Found {len(devices)} devices:")
            for idx, (ip, idn) in enumerate(devices, 1):
                print(f"  [{idx}] {ip:15s} - {idn}")

            # Prompt user to select
            while True:
                try:
                    choice = input(f"\nSelect device [1-{len(devices)}]: ").strip()
                    idx = int(choice)
                    if 1 <= idx <= len(devices):
                        args.host = devices[idx - 1][0]
                        print(f"✓ Selected: {args.host}")
                        break
                    else:
                        print(f"  Invalid choice. Enter 1-{len(devices)}")
                except (ValueError, KeyboardInterrupt, EOFError):
                    print("\n✗ Selection cancelled")
                    return 1

        print()

    # Validate arguments
    if args.git_deploy:
        # Git deployment mode - files not required
        if args.target:
            print("Error: --target cannot be used with --git-deploy")
            return 1
        if args.files:
            print("Error: File arguments not allowed with --git-deploy")
            return 1
        if args.from_commit and args.force_full:
            print("Error: --from-commit and --force-full are mutually exclusive")
            return 1
    else:
        # Manual file upload mode - files required
        if not args.files:
            print("Error: File arguments required (or use --git-deploy)")
            return 1
        if args.target and len(args.files) > 1:
            print("Error: --target can only be used with a single file")
            return 1
        if args.dry_run:
            print("Error: --dry-run only works with --git-deploy")
            return 1
        if args.from_commit:
            print("Error: --from-commit only works with --git-deploy")
            return 1
        if args.force_full:
            print("Error: --force-full only works with --git-deploy")
            return 1

    # Create client
    client = OTAClient(
        host=args.host,
        port=args.port,
        password=args.password,
        chunk_size=args.chunk_size,
    )

    try:
        with client:
            # Git deployment mode
            if args.git_deploy:
                return (
                    0
                    if deploy_from_git(
                        client=client,
                        from_commit=args.from_commit,
                        force_full=args.force_full,
                        dry_run=args.dry_run,
                        auto_reboot=args.reboot,
                        algo=args.algo,
                        auto_confirm=args.yes,
                        device=args.device,
                    )
                    else 1
                )

            # Batch upload mode
            if len(args.files) > 1:
                print(f"\n{'=' * 60}")
                print(f"Batch OTA Update: {len(args.files)} file(s)")
                print(f"{'=' * 60}\n")

                success_count = 0
                failed_files = []

                for idx, file_path in enumerate(args.files, 1):
                    file_name = Path(file_path).name
                    print(f"[{idx}/{len(args.files)}] Uploading {file_name}...")

                    success = client.upload_file(
                        target_filename=file_name,
                        local_path=file_path,
                        algo=args.algo,
                        auto_reboot=False,  # Don't reboot between files
                    )

                    if success:
                        success_count += 1
                        print(f"✓ {file_name} uploaded successfully")

                        # Query and display status after successful upload
                        status = client.query_update_status()
                        if status:
                            state = status.get("state", "unknown")
                            target = status.get("target", "none")
                            print(f"  Status: {state} (target: {target})")
                        print()
                    else:
                        failed_files.append(file_name)
                        print(f"✗ {file_name} failed\n")

                # Summary
                print(f"{'=' * 60}")
                print("Batch Upload Summary:")
                print(f"  Successful: {success_count}/{len(args.files)}")
                if failed_files:
                    print(f"  Failed: {', '.join(failed_files)}")
                print(f"{'=' * 60}\n")

                # Reboot if requested and at least one file succeeded
                if args.reboot and success_count > 0:
                    print("Rebooting device to apply changes...")
                    client.reboot_device()
                    print("\n✓ Batch OTA update completed")
                    return 0 if success_count == len(args.files) else 1
                elif success_count > 0:
                    print("Updates staged - reboot device to apply changes")
                    return 0 if success_count == len(args.files) else 1
                else:
                    print("\n✗ All updates failed")
                    return 1

            # Single file mode
            else:
                target = args.target if args.target else Path(args.files[0]).name

                success = client.upload_file(
                    target_filename=target,
                    local_path=args.files[0],
                    algo=args.algo,
                    auto_reboot=args.reboot,
                )

                if success:
                    print("\n✓ OTA update completed successfully")
                    return 0
                else:
                    print("\n✗ OTA update failed")
                    return 1

    except KeyboardInterrupt:
        print("\n\nAborted by user")
        return 130
    except Exception as e:
        print(f"\nFatal error: {e}")
        return 1


if __name__ == "__main__":
    import sys

    sys.exit(main())
