Zynq_Flasher/src/bitstream_programmer.py

"""Bitstream programming module for Zynq Flasher GUI.

Downloads a .bit bitstream to the Zynq PL fabric and runs a .elf
executable on the PS ARM core.

Uses xsct (Xilinx Software Command Line Tool) as the primary tool
for JTAG operations, with Vivado batch mode as a fallback.
Compatible with Xilinx 2018.3 — 2023.2.
"""

from __future__ import annotations

import os
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from config_manager import Config
from vitis_checker import get_xsct_path, get_vivado_path

# ── Types ───────────────────────────────────────────────────────────


@dataclass
class BitstreamResult:
    """Result of a bitstream programming or ELF execution operation."""

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


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

# ── Helpers ─────────────────────────────────────────────────────────


def _tcl_path(path: str | Path) -> str:
    """Convert a filesystem path to TCL-safe format.

    TCL interprets backslashes as escape sequences (\\f, \\t, \\n).
    Forward slashes work on all platforms including Windows.
    """
    return str(path).replace("\\", "/")


def _find_ps7_init_tcl(config: Config) -> str:
    """Find the ps7_init.tcl file for PS initialization.

    Checks:
    1. Same directory as bootloader BIT file (hw_platform)
    2. ps7_init_tcl_path in config (if added later)

    Args:
        config: Application configuration.

    Returns:
        Path to ps7_init.tcl, or empty string if not found.
    """
    # 1. Check config attribute (if added)
    if hasattr(config, 'ps7_init_tcl_path') and config.ps7_init_tcl_path:
        p = Path(config.ps7_init_tcl_path)
        if p.exists():
            return str(p)

    # 2. Same directory as BIT file (typical hw_platform layout)
    if config.bootloader_bit_path:
        bit_dir = Path(config.bootloader_bit_path).parent
        candidates = [
            bit_dir / "ps7_init.tcl",
        ]
        for c in candidates:
            if c.exists():
                return str(c)

    return ""


def _run_command(
    cmd: list[str],
    timeout: int = 120,
    callback: ProgressCallback = None,
) -> subprocess.CompletedProcess[str]:
    """Run a subprocess command with progress reporting.

    Args:
        cmd: Command and arguments.
        timeout: Maximum seconds.
        callback: Optional callback(status, message).

    Returns:
        CompletedProcess result.

    Raises:
        subprocess.TimeoutExpired: If command exceeds timeout.
    """
    from subprocess_utils import run_streaming

    if callback:
        callback("running", f"Executing: {' '.join(cmd)}")

    try:
        return run_streaming(cmd, timeout=timeout, callback=callback)
    except subprocess.TimeoutExpired:
        if callback:
            callback("error", f"{cmd[0]} timed out after {timeout}s")
        raise


# ── Bitstream Download ──────────────────────────────────────────────


def program_bitstream(
    config: Config,
    bit_path: str | Path,
    callback: ProgressCallback = None,
) -> BitstreamResult:
    """Download a .bit bitstream to the Zynq PL (FPGA fabric).

    Inits PS (ps7_init) before FPGA config to keep the JTAG DAP alive,
    allowing subsequent ELF download without reconnection issues.

    Args:
        config: Application configuration.
        bit_path: Path to the .bit file.
        callback: Optional progress callback.

    Returns:
        BitstreamResult with status.
    """
    bit_path = Path(bit_path)
    if not bit_path.exists():
        return BitstreamResult(
            step="bitstream",
            success=False,
            message=f"Bitstream file not found: {bit_path}",
        )

    if callback:
        callback("start", f"Loading bitstream: {bit_path.name}...")

    # Find ps7_init.tcl for PS initialization
    ps7_tcl = _find_ps7_init_tcl(config)

    # Try xsct first (preferred)
    xsct = get_xsct_path(
        xilinx_root=getattr(config, "xilinx_path", ""),
    )
    if xsct:
        timeout = config.step_timeouts.get("bitstream", 60)
        return _program_with_xsct(xsct, bit_path, callback, ps7_tcl, timeout)

    # Fallback to Vivado
    vivado = get_vivado_path(
        xilinx_root=getattr(config, "xilinx_path", ""),
    )
    if vivado:
        return _program_with_vivado(vivado, bit_path, callback)

    return BitstreamResult(
        step="bitstream",
        success=False,
        message="Neither xsct nor vivado found in PATH",
    )


def _program_with_xsct(
    xsct_path: str,
    bit_path: Path,
    callback: ProgressCallback = None,
    ps7_init_tcl: str = "",
    timeout: int = 60,
) -> BitstreamResult:
    """Program bitstream using xsct's fpga command.

    Inits the PS (clocks, DDR, MIO) before FPGA configuration to keep
    the JTAG DAP alive. Without PS init, the DAP breaks after fpga,
    preventing subsequent ELF download.

    Args:
        xsct_path: Path to xsct executable.
        bit_path: Path to .bit file.
        callback: Optional progress callback.
        ps7_init_tcl: Path to ps7_init.tcl for PS initialization.
        timeout: Subprocess timeout in seconds.

    Returns:
        BitstreamResult with status.
    """
    if callback:
        callback("progress", "Initializing PS then programming FPGA...")

    # Build TCL: PS init → FPGA config → post-config → halt CPU → disable Flash boot
    lines = ["connect"]
    if ps7_init_tcl and Path(ps7_init_tcl).exists():
        lines.append("targets 1")
        lines.append(f'source "{_tcl_path(ps7_init_tcl)}"')
        lines.append("ps7_init")
        if callback:
            callback("progress", "PS initialized (clocks, DDR, MIO)")
    lines.append(f'fpga -file {{{_tcl_path(bit_path)}}}')
    if ps7_init_tcl and Path(ps7_init_tcl).exists():
        lines.append("ps7_post_config")
        if callback:
            callback("progress", "PS post-config (PL-PS interfaces)")
    # After FPGA config, CPU may be reset — halt it to prevent Flash code
    # from enabling MMU before the ELF step can download FSBL
    lines.append("targets 2")
    lines.append("after 100")
    lines.append("stop")
    # Prevent boot from Flash by disabling PCAP reconfiguration
    lines.append("mwr 0xF800025C 0x00000001")
    lines.append("exit")
    tcl_script = "\n".join(lines)

    import tempfile, os
    with tempfile.NamedTemporaryFile(mode='w', suffix='.tcl', delete=False) as f:
        f.write(tcl_script)
        tcl_file = f.name

    try:
        result = subprocess.run(
            [xsct_path, tcl_file],
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        os.unlink(tcl_file)
        success = result.returncode == 0
        output = (result.stdout + result.stderr).strip()

        if callback:
            callback(
                "complete" if success else "error",
                "Bitstream " + ("loaded" if success else "failed"),
            )

        return BitstreamResult(
            step="bitstream",
            success=success,
            message="Bitstream loaded successfully" if success else "Bitstream load failed",
            output=output[:4000],
        )
    except subprocess.TimeoutExpired:
        os.unlink(tcl_file)
        return BitstreamResult(
            step="bitstream",
            success=False,
            message="xsct bitstream download timed out",
        )


def _program_with_vivado(
    vivado_path: str,
    bit_path: Path,
    callback: ProgressCallback = None,
) -> BitstreamResult:
    """Program bitstream using Vivado Hardware Manager.

    Args:
        vivado_path: Path to vivado executable.
        bit_path: Path to .bit file.
        callback: Optional progress callback.

    Returns:
        BitstreamResult with status.
    """
    if callback:
        callback("progress", "Using Vivado Hardware Manager...")

    tcl_content = f"""
open_hw_manager
connect_hw_server
open_hw_target
current_hw_device [lindex [get_hw_devices] 0]
program_hw_device -file "{bit_path}"
refresh_hw_device
quit
"""
    tcl_path = bit_path.parent / "temp_program.tcl"
    try:
        with open(tcl_path, "w") as f:
            f.write(tcl_content)

        result = _run_command(
            [vivado_path, "-mode", "batch", "-source", str(tcl_path)],
            timeout=120,
            callback=callback,
        )
        success = result.returncode == 0
        output = (result.stdout + result.stderr).strip()

        return BitstreamResult(
            step="bitstream",
            success=success,
            message="Bitstream loaded successfully" if success else "Bitstream load failed",
            output=output[:4000],
        )
    finally:
        tcl_path.unlink(missing_ok=True)


# ── ELF Execution ───────────────────────────────────────────────────


def run_elf(
    config: Config,
    elf_path: str | Path,
    callback: ProgressCallback = None,
) -> BitstreamResult:
    """Download and run a .elf executable on the Zynq PS ARM core.

    First downloads FSBL (First Stage Bootloader) to OCM (0x0) to
    initialize clocks, DDR, MIO and force JTAG boot mode.  Then
    downloads the target ELF to DDR.

    Args:
        config: Application configuration.
        elf_path: Path to the .elf file.
        callback: Optional progress callback.

    Returns:
        BitstreamResult with status.
    """
    elf_path = Path(elf_path)
    if not elf_path.exists():
        return BitstreamResult(
            step="elf",
            success=False,
            message=f"ELF file not found: {elf_path}",
        )

    fsbl_path = config.fsbl_elf_path if hasattr(config, 'fsbl_elf_path') else ""
    if fsbl_path and not Path(fsbl_path).exists():
        if callback:
            callback("warning", f"FSBL not found: {fsbl_path}, will try direct dow")
        fsbl_path = ""

    if callback:
        callback("start", f"Running ELF: {elf_path.name}...")

    xsct = get_xsct_path(
        xilinx_root=getattr(config, "xilinx_path", ""),
    )
    if not xsct:
        return BitstreamResult(
            step="elf",
            success=False,
            message="xsct not found in PATH (required for CPU operations)",
        )

    timeout = config.step_timeouts.get("elf_download", 60)

    if fsbl_path:
        return _run_elf_via_fsbl(xsct, elf_path, fsbl_path, callback, timeout)
    else:
        return _run_elf_direct(xsct, elf_path, callback, timeout)


def _run_elf_via_fsbl(
    xsct_path: str,
    elf_path: Path,
    fsbl_path: str,
    callback: ProgressCallback = None,
    timeout: int = 60,
) -> BitstreamResult:
    """Download and run ELF via FSBL bootloader.

    FSBL loads to OCM (0x0) which is always accessible even when
    MMU is enabled.  FSBL initializes clocks/DDR/MIO and forces
    JTAG boot mode, making DDR accessible for the app ELF.

    Args:
        xsct_path: Path to xsct executable.
        elf_path: Path to target .elf file.
        fsbl_path: Path to FSBL .elf file.
        callback: Optional progress callback.
        timeout: Subprocess timeout in seconds.

    Returns:
        BitstreamResult with status.
    """
    import tempfile, os

    if callback:
        callback("progress", "Downloading FSBL to OCM...")

    tcl = "\n".join([
        "connect",
        "targets 2",
        "dow {%s}" % _tcl_path(fsbl_path),
        "puts \"FSBL downloaded\"",
        "con",
        "after 3000",
        "stop",
        "puts \"FSBL done, CPU halted\"",
        "dow {%s}" % _tcl_path(elf_path),
        "mwr 0xF800025C 0x00000001",
        "con",
        "puts \"App running\"",
        "exit",
    ])

    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".tcl", delete=False
    ) as tf:
        tf.write(tcl)
        tcl_file = tf.name

    try:
        result = subprocess.run(
            [xsct_path, tcl_file],
            capture_output=True, text=True, timeout=timeout,
        )
        os.unlink(tcl_file)
        output = (result.stdout + result.stderr).strip()

        success = result.returncode == 0 or any(
            p in output for p in (
                "App running", "Download successful",
                "Successfully downloaded",
            )
        )

        if callback:
            callback(
                "complete" if success else "error",
                "ELF " + ("executed via FSBL" if success else "failed"),
            )

        return BitstreamResult(
            step="elf",
            success=success,
            message="ELF executed via FSBL" if success else "ELF execution failed",
            output=output[:4000],
        )
    except subprocess.TimeoutExpired:
        os.unlink(tcl_file)
        return BitstreamResult(step="elf", success=False,
                               message="ELF download timed out")


def _run_elf_direct(
    xsct_path: str,
    elf_path: Path,
    callback: ProgressCallback = None,
    timeout: int = 60,
) -> BitstreamResult:
    """Download and run ELF directly (no FSBL, assumes DAP alive).

    Used when no FSBL ELF is configured.  Attempts direct dow
    which only works if the CPU is in a clean state with MMU off.

    Args:
        xsct_path: Path to xsct executable.
        elf_path: Path to .elf file.
        callback: Optional progress callback.
        timeout: Subprocess timeout in seconds.

    Returns:
        BitstreamResult with status.
    """
    import tempfile, os

    if callback:
        callback("progress", "Downloading ELF directly...")

    tcl = "\n".join([
        "connect",
        "targets 2",
        "dow {%s}" % _tcl_path(elf_path),
        "mwr 0xF800025C 0x00000001",
        "con",
        "exit",
    ])

    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".tcl", delete=False
    ) as tf:
        tf.write(tcl)
        tcl_file = tf.name

    try:
        result = subprocess.run(
            [xsct_path, tcl_file],
            capture_output=True, text=True, timeout=timeout,
        )
        os.unlink(tcl_file)
        output = (result.stdout + result.stderr).strip()

        success = result.returncode == 0 or any(
            p in output for p in (
                "Download successful", "Successfully downloaded",
            )
        )

        if callback:
            callback("complete" if success else "error",
                     "ELF " + ("executed" if success else "failed"))

        return BitstreamResult(
            step="elf",
            success=success,
            message="ELF executed" if success else "ELF execution failed",
            output=output[:4000],
        )
    except subprocess.TimeoutExpired:
        os.unlink(tcl_file)
        return BitstreamResult(step="elf", success=False,
                               message="ELF download timed out")


# ── Full Workflow (combined BIT+ELF in one xsct session) ──────────


def full_bitstream_program(
    config: Config,
    bit_path: str | Path,
    elf_path: str | Path,
    callback: ProgressCallback = None,
) -> list[BitstreamResult]:
    """Execute full bootloader workflow in one xsct session.

    Flow:
      1. rst -processor (bypass BootROM, block Flash boot)
      2. dow fsbl → con → stop (FSBL initializes clocks/DDR/MIO)
      3. fpga -f bit (configure PL)
      4. dow app → con (run bootloader)

    Args:
        config: Application configuration.
        bit_path: Path to the .bit file.
        elf_path: Path to the .elf file.
        callback: Optional progress callback.

    Returns:
        List of BitstreamResult for each step.
    """
    import tempfile, os

    bit_path = Path(bit_path)
    elf_path = Path(elf_path)
    fsbl_path = config.fsbl_elf_path if hasattr(config, 'fsbl_elf_path') else ""

    # Validate files
    if not bit_path.exists():
        return [BitstreamResult(step="bitstream", success=False,
                message=f"BIT not found: {bit_path}")]
    if not elf_path.exists():
        return [BitstreamResult(step="elf", success=False,
                message=f"ELF not found: {elf_path}")]
    if not fsbl_path or not Path(fsbl_path).exists():
        return [BitstreamResult(step="elf", success=False,
                message=f"FSBL not found: {fsbl_path or '(not configured)'}")]

    xsct = get_xsct_path(
        xilinx_root=getattr(config, "xilinx_path", ""),
    )
    if not xsct:
        return [BitstreamResult(step="bitstream", success=False,
                message="xsct not found")]

    if callback:
        callback("start", "Initializing Zynq via JTAG (FSBL override)...")

    # Build TCL with preemptive system reset + DAP error recovery
    tcl_lines = [
        'puts "INFO: Starting Zynq initialization via JTAG override..."',
        "connect",
        "after 500",
    ]

    # Preemptive system reset (safe on clean board; may help after flash)
    tcl_lines.extend([
        '# Preemptive system reset to clean up any leftover state',
        'if {[catch {',
        '  targets -set -nocase -filter {name =~ "arm*#0"}',
        '  rst -system',
        '  after 1000',
        '  stop',
        '  puts "INFO: Preemptive system reset OK"',
        '}]} { puts "INFO: Preemptive system reset skipped (no ARM target)" }',
        '',
    ])

    # DAP error check — if ARM cores not visible, try recovery via DAP
    tcl_lines.extend([
        '# Check for DAP errors and attempt recovery',
        'set target_list [targets]',
        'set has_arm 0',
        'foreach t $target_list { if {[regexp -nocase {arm|cortex} $t]} { set has_arm 1 } }',
        'if {!$has_arm} {',
        '  puts "WARNING: ARM cores not found — DAP may be in error state"',
        '  # Show DAP status for logging',
        '  foreach t $target_list { puts "  TARGET: $t" }',
        '  # Attempt recovery: select DAP (target 1) and issue system reset',
        '  puts "INFO: Attempting DAP recovery via system reset..."',
        '  catch {',
        '    targets 1',
        '    rst -system',
        '    after 1000',
        '    stop',
        '  }',
        '  # Disconnect and reconnect to refresh JTAG chain',
        '  catch { disconnect; after 1000; connect; after 500 }',
        '  # Check if ARM cores reappeared',
        '  set target_list2 [targets]',
        '  set has_arm2 0',
        '  foreach t $target_list2 {',
        '    puts "  AFTER RECOVERY: $t"',
        '    if {[regexp -nocase {arm|cortex} $t]} { set has_arm2 1 }',
        '  }',
        '  if {!$has_arm2} {',
        '    puts "FATAL: DAP recovery failed — board requires power cycle"',
        '    puts "FATAL: Unplug/replug power, then retry Step 3"',
        '    exit 1',
        '  }',
        '  puts "INFO: ARM cores recovered after DAP reset"',
        '}',
        '',
    ])

    # Core flow
    tcl_lines.extend([
        'targets -set -nocase -filter {name =~ "arm*#0"}',
        'puts "INFO: Resetting processor (Bypassing BootROM Flash boot)..."',
        "rst -processor",
        "after 500",
        'puts "INFO: Downloading FSBL..."',
        'dow {%s}' % _tcl_path(fsbl_path),
        'puts "INFO: Running FSBL for hardware initialization..."',
        "con",
        "after 2000",
        "stop",
        'puts "INFO: Hardware initialization completed."',
        'puts "INFO: Downloading Bitstream..."',
        'fpga -f {%s}' % _tcl_path(bit_path),
        "after 1000",
        'targets -set -nocase -filter {name =~ "arm*#0"}',
        'puts "INFO: Downloading Application ELF..."',
        'dow {%s}' % _tcl_path(elf_path),
        'puts "INFO: Executing Application..."',
        "con",
        'puts "SUCCESS: Zynq target is running the specified application."',
        "exit",
    ])

    tcl = "\n".join(tcl_lines)

    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".tcl", delete=False
    ) as tf:
        tf.write(tcl)
        tcl_file = tf.name

    timeout = config.step_timeouts.get("bitstream", 60) + \
              config.step_timeouts.get("elf_download", 60)

    try:
        if callback:
            callback("progress", "Executing combined JTAG flow...")

        result = subprocess.run(
            [xsct, tcl_file],
            capture_output=True, text=True, timeout=timeout,
        )
        os.unlink(tcl_file)
        output = (result.stdout + result.stderr).strip()

        success = result.returncode == 0 and "SUCCESS" in output

        # Detect DAP errors for logging
        dap_errors = [l for l in output.split("\n") if "DAP" in l and "error" in l.lower()]
        if dap_errors:
            if callback:
                for e in dap_errors:
                    callback("error", f"DAP: {e.strip()[:120]}")
            success = False

        if callback:
            callback(
                "complete" if success else "error",
                "Bootloader " + ("running" if success else "failed"),
            )

        return [
            BitstreamResult(
                step="bitstream",
                success=success,
                message=(
                    "DAP error — board needs power cycle"
                    if dap_errors else
                    "Bitstream loaded" if success else "Combined flow failed"
                ),
                output=output[:4000],
            ),
            BitstreamResult(
                step="elf",
                success=success,
                message=(
                    "DAP error — board needs power cycle"
                    if dap_errors else
                    "ELF executed" if success else "Combined flow failed"
                ),
                output=output[:4000],
            ),
        ]
    except subprocess.TimeoutExpired:
        os.unlink(tcl_file)
        return [
            BitstreamResult(step="bitstream", success=False,
                            message="Combined JTAG flow timed out"),
            BitstreamResult(step="elf", success=False,
                            message="Combined JTAG flow timed out"),
        ]