Jenkinsfile.reflash

pipeline {
    agent any

    parameters {
        choice(
            name: 'BOARD_TYPE',
            choices: ["i.mx8mp-frdm", "i.mx95-frdm"],
            description: 'Choose board type to flash'
        )
        booleanParam(
            name: 'FLASH_ALL_BOARDS',
            defaultValue: false,
            description: 'Flash all boards sequentially (ignores BOARD_TYPE selection)'
        )
        string(
            name: 'BSP_VERSION',
            defaultValue: 'q1-26',
            description: 'BSP version in format qx-yy (e.g., q1-26)'
        )
        string(
            name: 'RT_SDK_ARA2_VERSION',
            defaultValue: '2.0.4',
            description: 'rt-sdk-ara2 version (e.g., 2.0.4)'
        )
    }

    options {
        timestamps()
        timeout(time: 2, unit: 'HOURS')
    }

    environment {
        TARGET_USER = "root"
        SSH_OPTS = "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o ConnectTimeout=10 -o ServerAliveInterval=60"
        BSP_IMAGES_DIR = "${HOME}/bsp_images/${params.BSP_VERSION}"
        SDK_DEB = "rt-sdk-ara2_${params.RT_SDK_ARA2_VERSION}.deb"
        SDK_PATH = "${HOME}/bsp_images/sdks/${SDK_DEB}"
        // Network discovery settings
        ARP_SCAN_TIMEOUT = "180"  // seconds to wait for board discovery
    }

    stages {
        stage('Validate Parameters') {
            steps {
                script {
                    // Validate BSP version format
                    if (!params.BSP_VERSION.matches(/q\d+-\d+/)) {
                        error("Invalid BSP_VERSION format. Expected format: qx-yy (e.g., q1-26)")
                    }

                    // Check if BSP images directory exists
                    sh """
                        if [ ! -d "${BSP_IMAGES_DIR}" ]; then
                            echo "ERROR: BSP images directory not found: ${BSP_IMAGES_DIR}"
                            exit 1
                        fi
                    """

                    // Check if SDK deb exists
                    sh """
                        if [ ! -f "${SDK_PATH}" ]; then
                            echo "ERROR: rt-sdk-ara2 deb not found: ${SDK_PATH}"
                            exit 1
                        fi
                    """

                    // Determine which boards to flash
                    if (params.FLASH_ALL_BOARDS) {
                        env.BOARDS_TO_FLASH = "i.mx8mp-frdm,i.mx95-frdm"
                    } else {
                        env.BOARDS_TO_FLASH = params.BOARD_TYPE
                    }

                    echo "Boards to flash: ${env.BOARDS_TO_FLASH}"
                    echo "BSP Version: ${params.BSP_VERSION}"
                    echo "rt-sdk-ara2 Version: ${params.RT_SDK_ARA2_VERSION}"
                }
            }
        }

        stage('Flash Boards') {
            steps {
                script {
                    def boards = env.BOARDS_TO_FLASH.split(',')

                    for (board in boards) {
                        echo "=========================================="
                        echo "Processing board: ${board}"
                        echo "=========================================="

                        // Lock resource by label (board type)
                        lock(variable: 'LOCKED_BOARD', extra: [[label: board, quantity: 1]], resource: null) {
                            def boardIP = env.LOCKED_BOARD0_IP
                            def boardName = env.LOCKED_BOARD0_NAME  // e.g., "DP107297"

                            if (!boardIP) {
                                error("Could not determine IP for board with label: ${board}. Check that the resource has an 'IP' property configured.")
                            }

                            echo "Board label: ${board}"
                            echo "Board name: ${boardName}"
                            echo "Board IP: ${boardIP}"

                            flashBoard(board, boardIP)
                        }
                    }
                }
            }
        }
    }

    post {
        always {
            echo "Reflash pipeline completed"
        }
        failure {
            echo "Reflash pipeline failed - manual intervention may be required"
        }
    }
}

/**
 * Discover a board on the network by its MAC address using arp-scan
 *
 * @param mac The MAC address to search for
 * @param subnet The subnet to scan (e.g., "192.168.1.0/24")
 * @param timeoutSeconds Maximum time to wait for discovery
 * @return The discovered IP address, or null if not found
 */
def discoverBoardByMAC(String mac, String subnet, int timeoutSeconds = 180) {
    def scanInterval = 10  // seconds between scans
    def maxAttempts = (timeoutSeconds / scanInterval).toInteger()
    
    echo "Scanning ${subnet} for MAC ${mac} (timeout: ${timeoutSeconds}s)..."
    
    for (int i = 0; i < maxAttempts; i++) {
        def ip = sh(
            script: """
                sudo arp-scan --ignoredups --quiet ${subnet} 2>/dev/null | \
                    grep -i '${mac}' | \
                    awk '{print \$1}' | \
                    head -1
            """,
            returnStdout: true
        ).trim()
        
        if (ip) {
            echo "Found MAC ${mac} at IP ${ip}"
            return ip
        }
        
        echo "Discovery attempt ${i + 1}/${maxAttempts} - board not found yet..."
        sleep(time: scanInterval, unit: 'SECONDS')
    }
    
    return null
}

/**
 * Calculate subnet from an IP address (assumes /24)
 */
def getSubnet(String ip) {
    def octets = ip.tokenize('.')
    return "${octets[0]}.${octets[1]}.${octets[2]}.0/24"
}

def flashBoard(String boardType, String boardIP) {
    def bootImage = "boot-${boardType}"
    def rootImage = "root-${boardType}"
    def bootDevice = getBootDevice(boardType)
    def bootPartition = getBootPartition(boardType)
    
    // Variables to store board identity
    def boardMAC = ""
    def originalIP = boardIP
    def subnet = getSubnet(boardIP)
    def currentIP = boardIP  // This will be updated after discovery

    echo "Boot image: ${bootImage}"
    echo "Root image: ${rootImage}"
    echo "Boot device: ${bootDevice}"
    echo "Boot partition: ${bootPartition}"

    // Verify images exist
    sh """
        if [ ! -f "${BSP_IMAGES_DIR}/${bootImage}" ]; then
            echo "ERROR: Boot image not found: ${BSP_IMAGES_DIR}/${bootImage}"
            exit 1
        fi
        if [ ! -f "${BSP_IMAGES_DIR}/${rootImage}" ]; then
            echo "ERROR: Root image not found: ${BSP_IMAGES_DIR}/${rootImage}"
            exit 1
        fi
    """

    stage("Capture Board Identity - ${boardType}") {
        echo "Capturing board MAC address before flashing..."
        
        boardMAC = sh(
            script: "ssh ${SSH_OPTS} ${TARGET_USER}@${boardIP} 'cat /sys/class/net/eth0/address'",
            returnStdout: true
        ).trim()
        
        if (!boardMAC || !boardMAC.matches(/([0-9a-fA-F]{2}:){5}[0-9a-fA-F]{2}/)) {
            error("Failed to capture valid MAC address from board. Got: '${boardMAC}'")
        }
        
        echo "Board Identity Captured:"
    }

    stage("Erase Bootloader - ${boardType}") {
        echo "Connecting to board and erasing bootloader..."

        // Enable boot partition write access and erase bootloader
        sh """
            ssh ${SSH_OPTS} ${TARGET_USER}@${boardIP} "
                echo 'Enabling boot partition write access...'
                echo 0 > /sys/block/${bootPartition}/force_ro || true
                echo 0 > /sys/block/${bootDevice}/${bootPartition}/force_ro || true

                echo 'Erasing bootloader...'
                dd if=/dev/zero of=/dev/${bootPartition} bs=1M count=4 conv=fsync
                
                echo 'Bootloader erased successfully'
                sync
            "
        """
    }

    stage("Reboot to Download Mode - ${boardType}") {
        echo "Rebooting board to enter download mode..."

        // Reboot and don't wait for response (board will lose connectivity)
        sh """
            ssh ${SSH_OPTS} ${TARGET_USER}@${boardIP} "reboot" || true
        """

        // Wait for board to enter download mode
        echo "Waiting for board to enter download mode..."
        sleep(time: 30, unit: 'SECONDS')
    }

    stage("Flash with UUU - ${boardType}") {
        echo "Flashing board using UUU..."

        // Wait for UUU to detect the board
        def uuuDetected = false
        def maxAttempts = 30

        for (int i = 0; i < maxAttempts && !uuuDetected; i++) {
            def result = sh(
                script: "uuu -lsusb 2>&1 | grep -i 'mx8\\|mx95\\|SDP'",
                returnStatus: true
            )
            if (result == 0) {
                uuuDetected = true
                echo "Board detected in download mode"
            } else {
                echo "Waiting for board to appear in download mode... (${i + 1}/${maxAttempts})"
                sleep(time: 10, unit: 'SECONDS')
            }
        }

        if (!uuuDetected) {
            error("Board did not appear in download mode within timeout")
        }

        // Flash the board using UUU
        sh """
            cd ${BSP_IMAGES_DIR}

            echo "Flashing bootloader and rootfs..."
            sudo uuu -b emmc_all ${bootImage} ${rootImage}

            echo "Flash completed successfully"
            
            echo "Resetting board..."
            sudo uuu FB: acmd reset || sudo uuu SDP: reset || true
        """
    }

    stage("Discover Board - ${boardType}") {
        echo "Waiting for board to boot and acquire network..."
        
        // Initial wait for board to boot
        sleep(time: 45, unit: 'SECONDS')
        
        // Discover board by MAC address
        currentIP = discoverBoardByMAC(boardMAC, subnet, ARP_SCAN_TIMEOUT.toInteger())
        
        if (!currentIP) {
            error("Failed to discover board with MAC ${boardMAC} on subnet ${subnet}")
        }
        
        echo "Board Found after Reboot"
    }

    stage("Restore Static IP - ${boardType}") {
        if (currentIP != originalIP) {
            echo "Restoring original static IP: ${originalIP}"
            
            sh """
                ssh ${SSH_OPTS} ${TARGET_USER}@${currentIP} '
                    cat > /etc/systemd/network/10-static.network << EOF
[Match]
Name=eth0

[Network]
Address=${originalIP}/24
EOF
                    systemctl restart systemd-networkd
                '
            """
            
            // Wait for network to reconfigure
            sleep(time: 10, unit: 'SECONDS')
            
            // Verify new IP is accessible
            def verifyResult = sh(
                script: "ssh ${SSH_OPTS} ${TARGET_USER}@${originalIP} 'echo IP restored successfully'",
                returnStatus: true
            )
            
            if (verifyResult != 0) {
                echo "WARNING: Could not verify restored IP ${originalIP}, continuing with ${currentIP}"
            } else {
                currentIP = originalIP
                echo "Static IP restored successfully: ${currentIP}"
            }
        } else if (currentIP == originalIP) {
            echo "Board retained original IP ${originalIP}, no restoration needed"
        } else {
            echo "Static IP restoration disabled, using DHCP-assigned IP: ${currentIP}"
        }
    }

    stage("Install uv and rt-sdk-ara2 - ${boardType}") {
        echo "Installing uv"
        sh """
            ssh ${SSH_OPTS} ${TARGET_USER}@${currentIP} "curl -LsSf https://astral.sh/uv/install.sh | sh"
        """
        echo "Installing rt-sdk-ara2..."

        // Copy SDK deb to board
        sh """
            scp ${SSH_OPTS} ${SDK_PATH} ${TARGET_USER}@${currentIP}:/tmp/${SDK_DEB}
        """

        // Set correct time and install SDK
        sh """
            ssh ${SSH_OPTS} ${TARGET_USER}@${currentIP} "
                echo 'Setting system time...'
                date -s '@\$(date +%s)'

                echo 'Installing rt-sdk-ara2...'
                dpkg -i /tmp/${SDK_DEB}

                echo 'Cleaning up...'
                rm -f /tmp/${SDK_DEB}

                echo 'Installation completed successfully'
            "
        """
    }

    stage("Mount NFS Models - ${boardType}") {
        echo "Mounting LLM models from NFS server..."

        sh """
            ssh ${SSH_OPTS} ${TARGET_USER}@${currentIP} '
                SERVER="10.171.82.129"
                REMOTE="/srv/nfs/shared"
                TARGET="/usr/share/llm"

                echo "Creating temporary mount point..."
                mkdir -p /mnt/tmp

                echo "Mounting NFS share..."
                mountpoint -q /mnt/tmp || mount "\$SERVER:\$REMOTE" /mnt/tmp

                echo "Mounting model directories and adding to fstab..."
                for dir in /mnt/tmp/*; do
                    if [ -d "\$dir" ]; then
                        name=\$(basename "\$dir")
                        echo "  Setting up \$name..."
                        mkdir -p "\$TARGET/\$name"
                        
                        # Add to fstab if not already present
                        FSTAB_ENTRY="\$SERVER:\$REMOTE/\$name \$TARGET/\$name nfs defaults,_netdev 0 0"
                        grep -qF "\$TARGET/\$name" /etc/fstab || echo "\$FSTAB_ENTRY" >> /etc/fstab
                        
                        # Mount now
                        mountpoint -q "\$TARGET/\$name" || mount "\$TARGET/\$name"
                    fi
                done

                echo "Cleaning up temporary mount..."
                umount /mnt/tmp

                echo "NFS models mounted successfully"
                echo "Mounted directories:"
                ls -la \$TARGET/
            '
        """
    }

    stage("Final Reboot - ${boardType}") {
        echo "Performing final reboot..."

        sh """
            ssh ${SSH_OPTS} ${TARGET_USER}@${currentIP} "reboot" || true
        """

        // Wait for board to come back up
        sleep(time: 30, unit: 'SECONDS')

        // Re-discover board after reboot (IP might change if static IP wasn't persisted)
        def finalIP = discoverBoardByMAC(boardMAC, subnet, 120)
        
        if (!finalIP) {
            error("Board did not come back up after final reboot")
        }
        
        currentIP = finalIP
        echo "Board ${boardType} successfully reflashed and ready at ${currentIP}"

        // Verify installation
        sh """
            ssh ${SSH_OPTS} ${TARGET_USER}@${currentIP} "
                echo '=== Board Info ==='
                uname -a
                echo ''
                echo '=== Network Config ==='
                ip addr show eth0
                echo ''
                echo '=== rt-sdk-ara2 Status ==='
                dpkg -l | grep rt-sdk-ara2 || echo 'Package not found in dpkg list'
            "
        """
        
        // Final IP consistency check
        if (currentIP != originalIP) {
            echo "WARNING: Final IP ${currentIP} differs from original ${originalIP}"
            echo "Static IP configuration may not have persisted across reboot"
        }
    }

    stage("Update ARA firmware ${boardType}") {
        echo "Updating ARA firmware..."

        sh """
            ssh ${SSH_OPTS} ${TARGET_USER}@${currentIP} "program_flash.sh" || true
        """
    }

    echo "============================================"
    echo "Board ${boardType} successfully reflashed!"
    echo "Final IP: ${currentIP}"
    echo "============================================"
}

def getBootDevice(String boardType) {
    // Return the eMMC device name based on board type
    switch (boardType) {
        case "i.mx8mp-frdm":
            return "mmcblk2"
        case "i.mx95-frdm":
            return "mmcblk0"
        default:
            return "mmcblk2"
    }
}

def getBootPartition(String boardType) {
    // Return the boot partition name based on board type
    switch (boardType) {
        case "i.mx8mp-frdm":
            return "mmcblk2boot0"
        case "i.mx95-frdm":
            return "mmcblk0boot0"
        default:
            return "mmcblk2boot0"
    }
}