Zynq_Flasher/src/reboot_manager.py

"""Reboot manager for Zynq Flasher GUI.

Handles rebooting the Zynq device via JTAG, UDP, or serial.
Also provides UDP-based version query.
"""

from __future__ import annotations

import os
import socket
import subprocess
import tempfile
import time
from dataclasses import dataclass
from typing import Callable

from config_manager import Config


@dataclass
class RebootResult:
    """Result of a reboot operation."""

    success: bool
    message: str
    method: str = ""
    output: str = ""


@dataclass
class UdpVersionResult:
    """Result of a UDP version query."""

    success: bool
    message: str
    version: str = ""


ProgressCallback = Callable[[str, str], None] | None

# ── UDP helpers ──────────────────────────────────────────────────────


def _udp_port_from_ip(ip: str) -> int:
    """Derive UDP command port from the last octet of the IP address.

    Port = last_octet - 3 (e.g., 192.168.100.11 → 8).

    Args:
        ip: Zynq IP address string.

    Returns:
        UDP port number, or -1 on parse failure.
    """
    try:
        octets = ip.strip().split(".")
        last = int(octets[3])
        return last - 3
    except (ValueError, IndexError):
        return -1


def _send_udp_command(
    ip: str,
    command: bytes,
    timeout: float = 5.0,
    buffer_size: int = 4096,
) -> tuple[bool, str]:
    """Send a UDP command to the Zynq and receive the response.

    Args:
        ip: Zynq IP address.
        command: Command bytes to send (e.g., b'reboot()').
        timeout: Receive timeout in seconds.
        buffer_size: Max receive buffer size.

    Returns:
        (success, response_text).
    """
    port = _udp_port_from_ip(ip)
    if port <= 0:
        return False, f"Invalid IP for UDP port derivation: {ip}"

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(timeout)
    try:
        sock.sendto(command, (ip, port))
        data, addr = sock.recvfrom(buffer_size)
        return True, data.decode("utf-8", errors="replace").strip()
    except socket.timeout:
        return False, f"UDP timeout ({timeout}s)"
    except OSError as e:
        return False, f"UDP error: {e}"
    finally:
        sock.close()


# ── Reboot methods ───────────────────────────────────────────────────


def reboot_via_serial(
    config: Config,
    callback: ProgressCallback = None,
) -> RebootResult:
    """Reboot the Zynq device via serial port.

    Sends reboot() through the serial console.

    Args:
        config: Application configuration.
        callback: Optional progress callback.

    Returns:
        RebootResult with status.
    """
    if not config.serial_port:
        return RebootResult(
            success=False,
            message="No serial port configured",
            method="serial",
        )

    if callback:
        callback("start", f"Sending reboot via serial {config.serial_port}...")

    try:
        import serial
        from serial_monitor import normalize_port

        with serial.Serial(
            normalize_port(config.serial_port),
            config.serial_baudrate,
            timeout=5,
        ) as ser:
            ser.reset_input_buffer()
            # Send Ctrl+C to break out of any running process
            ser.write(b"\x03")
            time.sleep(0.5)
            # Send reboot command as function call
            ser.write(b"reboot()\r\n")
            ser.reset_input_buffer()

        if callback:
            callback("progress", "reboot() sent via serial")

        return RebootResult(
            success=True,
            message="reboot() sent via serial",
            method="serial",
        )
    except ImportError:
        return RebootResult(
            success=False,
            message="pyserial not installed",
            method="serial",
        )
    except serial.SerialException as e:
        return RebootResult(
            success=False,
            message=f"Serial error: {e}",
            method="serial",
        )
    except OSError as e:
        return RebootResult(
            success=False,
            message=f"OS error: {e}",
            method="serial",
        )


def reboot_via_jtag(
    config: Config,
    callback: ProgressCallback = None,
) -> RebootResult:
    """Reboot the Zynq device via JTAG system reset.

    Connects to the Zynq over JTAG through hw_server and issues
    a full system reset (rst -system), then resumes ARM execution.

    Args:
        config: Application configuration.
        callback: Optional progress callback.

    Returns:
        RebootResult with status.
    """
    from vitis_checker import get_xsct_path

    xsct = get_xsct_path(
        xilinx_root=getattr(config, "xilinx_path", ""),
    )
    if not xsct:
        return RebootResult(
            success=False,
            message="xsct not found — Xilinx tools not available",
            method="jtag",
        )

    if callback:
        callback("start", f"JTAG reset via xsct ({xsct})...")

    # Build TCL: system reset then resume ARM execution (no stop)
    tcl_script = """\
puts "INFO: Connecting to hw_server..."
connect
after 500

puts "INFO: Selecting ARM target..."
targets -set -nocase -filter {name =~ "arm*#0"}
after 200

puts "INFO: Issuing system reset..."
rst -system
after 1500

puts "INFO: Resuming ARM execution..."
catch { con }

puts "INFO: System reset complete — disconnecting..."
disconnect
exit
"""

    try:
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".tcl", delete=False,
        ) as f:
            f.write(tcl_script)
            tcl_path = f.name

        if callback:
            callback("progress", "Running JTAG system reset...")

        from vitis_checker import build_tool_command
        result = subprocess.run(
            build_tool_command(xsct, tcl_path),
            capture_output=True,
            text=True,
            timeout=30,
        )
        os.unlink(tcl_path)

        success = result.returncode == 0
        output = (result.stdout + result.stderr).strip()

        if callback:
            callback(
                "complete" if success else "error",
                "JTAG reset " + ("succeeded" if success else "failed"),
            )

        return RebootResult(
            success=success,
            message=f"JTAG system reset {'OK' if success else 'failed'}",
            method="jtag",
            output=output[:2000],
        )
    except subprocess.TimeoutExpired:
        try:
            os.unlink(tcl_path)
        except OSError:
            pass
        return RebootResult(
            success=False,
            message="JTAG reset timed out (30s)",
            method="jtag",
        )
    except OSError as e:
        try:
            os.unlink(tcl_path)
        except OSError:
            pass
        return RebootResult(
            success=False,
            message=f"JTAG reset error: {e}",
            method="jtag",
        )


def reboot_via_udp(
    config: Config,
    callback: ProgressCallback = None,
) -> RebootResult:
    """Reboot the Zynq device via UDP command.

    Sends reboot() to the Zynq's derived UDP port (last-octet - 3).

    Args:
        config: Application configuration.
        callback: Optional progress callback.

    Returns:
        RebootResult with status.
    """
    ip = config.zynq_ip
    port = _udp_port_from_ip(ip)

    if callback:
        callback("start", f"UDP reboot → {ip}:{port}")

    ok, resp = _send_udp_command(ip, b"reboot()")

    if callback:
        if ok:
            callback("complete", f"UDP reboot OK: {resp}")
        else:
            callback("error", f"UDP reboot failed: {resp}")

    return RebootResult(
        success=ok,
        message=f"reboot() via UDP: {resp}" if ok else resp,
        method="udp",
        output=resp,
    )


# ── Version query ────────────────────────────────────────────────────


def get_version_via_udp(
    config: Config,
    callback: ProgressCallback = None,
) -> UdpVersionResult:
    """Query Zynq firmware version via UDP.

    Sends ver() to the Zynq's derived UDP port (last-octet - 3).

    Args:
        config: Application configuration.
        callback: Optional progress callback.

    Returns:
        UdpVersionResult with version string.
    """
    ip = config.zynq_ip
    port = _udp_port_from_ip(ip)

    if callback:
        callback("start", f"UDP version query → {ip}:{port}")

    ok, resp = _send_udp_command(ip, b"ver()")

    if callback:
        if ok:
            callback("complete", f"Version: {resp}")
        else:
            callback("error", f"Version query failed: {resp}")

    return UdpVersionResult(
        success=ok,
        message=f"Version: {resp}" if ok else resp,
        version=resp if ok else "",
    )


# ── Top-level dispatcher ─────────────────────────────────────────────


def reboot_zynq(
    config: Config,
    method: str = "auto",
    timeout: int | None = None,
    callback: ProgressCallback = None,
) -> RebootResult:
    """Reboot the Zynq device using the best available method.

    Tries JTAG first, falls back to UDP then serial.

    Args:
        config: Application configuration.
        method: Reboot method: 'auto', 'jtag', 'udp', or 'serial'.
        timeout: Seconds to wait for reboot (ssh-only, kept for compat).
        callback: Optional progress callback.

    Returns:
        RebootResult with status.
    """
    timeout = timeout or config.reboot_timeout

    if method == "jtag":
        return reboot_via_jtag(config, callback)
    elif method == "udp":
        return reboot_via_udp(config, callback)
    elif method == "serial":
        return reboot_via_serial(config, callback)

    # Auto: try JTAG first, then UDP, then serial
    result = reboot_via_jtag(config, callback)
    if result.success:
        return result

    if callback:
        callback("progress", "JTAG failed, trying UDP...")
    result = reboot_via_udp(config, callback)
    if result.success:
        return result

    if callback:
        callback("progress", "UDP failed, trying serial...")

    return reboot_via_serial(config, callback)