pipeline {
  agent none
  parameters {
    booleanParam(name: 'fpgadataflow', defaultValue: false, description: 'Run fpgadataflow tests')
    booleanParam(name: 'sanity', defaultValue: true, description: 'Run sanity hardware and unit tests')
    booleanParam(name: 'end2end', defaultValue: false, description: 'Run end2end tests')
  }
  stages {
    stage('Run Tests') {
      parallel {
        stage('Sanity - Build Hardware') {
          when {
            expression { return params['sanity'] }
          }
          agent {
            label 'finn-build'
          }
          environment {
            TEST_NAME = "bnn_build_sanity"
            FINN_HOST_BUILD_DIR = "${env.FINN_HOST_BUILD_DIR}/${env.TEST_NAME}"
          }
          steps {
            catchError(stageResult: 'FAILURE') {
              script {
                // Creates dir in finn clone to store build files for stashing
                sh "mkdir -p ${env.TEST_NAME}"
                cleanPreviousBuildFiles(env.FINN_HOST_BUILD_DIR)

                // Pass in the marker to run with pytest and the XML test results filename
                runDockerPytestWithMarker("sanity_bnn", "${env.TEST_NAME}", '')

                // Find the board's build files (bitstreams/xclbins) and zip for use on the boards themselves
                findCopyZip("Pynq-Z1", env.FINN_HOST_BUILD_DIR, env.TEST_NAME)
                findCopyZip("ZCU104", env.FINN_HOST_BUILD_DIR, env.TEST_NAME)
                findCopyZip("KV260_SOM", env.FINN_HOST_BUILD_DIR, env.TEST_NAME)
                findCopyZip("U250", env.FINN_HOST_BUILD_DIR, env.TEST_NAME)

                // Stash the test results file(s)
                stash name: "${env.TEST_NAME}", includes: "${env.TEST_NAME}.xml,${env.TEST_NAME}.html"

                // Use an env variable to help collect test results later in pipeline
                env.BNN_BUILD_SANITY = "SUCCESS"
              }
            }
          }
        }
        stage('Sanity - Unit Tests') {
          when {
            expression { params['sanity'] }
          }
          agent {
            label 'finn-build'
          }
          environment {
            TEST_NAME = "sanity_ut"
            FINN_HOST_BUILD_DIR = "${env.FINN_HOST_BUILD_DIR}/${env.TEST_NAME}"
          }
          steps {
            catchError(stageResult: 'FAILURE') {
              script {
                cleanPreviousBuildFiles(env.FINN_HOST_BUILD_DIR)

                // Multiple markers with pytest needs its own script
                createMultiMarkerScript("util or brevitas_export or streamline or transform or notebooks", "${env.TEST_NAME}", "--cov --cov-report=html:coverage_sanity_ut")
                sh './run-docker.sh ./run-tests.sh'

                // Stash the test results file(s)
                stash name: env.TEST_NAME, includes: "${env.TEST_NAME}.xml,${env.TEST_NAME}.html"

                // Use an env variable to help collect test results later in pipeline
                env.SANITY_UT = "SUCCESS"

                // Archive coverage report if successful
                archiveSuccessfulStage(env.SANITY_UT, "coverage_sanity_ut")
              }
            }
          }
        }
        stage('fpgadataflow Tests') {
          when {
            expression { params['fpgadataflow'] }
          }
          agent {
            label 'finn-build'
          }
          environment {
            TEST_NAME = "fpgadataflow"
            FINN_HOST_BUILD_DIR = "${env.FINN_HOST_BUILD_DIR}/${env.TEST_NAME}"
          }
          steps {
            catchError(stageResult: 'FAILURE') {
              script {
                cleanPreviousBuildFiles(env.FINN_HOST_BUILD_DIR)

                // Pass in the marker to run with pytest and the XML test results filename
                runDockerPytestWithMarker("fpgadataflow", "${env.TEST_NAME}", "--cov --cov-report=html:coverage_fpgadataflow")

                // Stash the test results file(s)
                stash name: env.TEST_NAME, includes: "${env.TEST_NAME}.xml,${env.TEST_NAME}.html"

                // Use an env variable to help collect test results later in pipeline
                env.FPGADATAFLOW_RESULT = "SUCCESS"

                // Archive coverage report if successful
                archiveSuccessfulStage(env.FPGADATAFLOW_RESULT, "coverage_fpgadataflow")
              }
            }
          }
        }
        stage('End2end') {
          when {
            expression { params['end2end'] }
          }
          agent {
            label 'finn-build'
          }
          environment {
            TEST_NAME = "end2end"
            FINN_HOST_BUILD_DIR = "${env.FINN_HOST_BUILD_DIR}/${env.TEST_NAME}"
          }
          steps {
            catchError(stageResult: 'FAILURE') {
              script {
                // Delete any build files from a previous build
                cleanPreviousBuildFiles(env.FINN_HOST_BUILD_DIR)

                // Pass in the marker to run with pytest and the XML test results filename
                runDockerPytestWithMarker(env.TEST_NAME, "${env.TEST_NAME}", '')

                // Stash the test results file(s)
                stash name: env.TEST_NAME, includes: "${env.TEST_NAME}.xml,${env.TEST_NAME}.html"

                // Use an env variable to help collect test results later in pipeline
                env.END2END_RESULT = "SUCCESS"
              }
            }
          }
        }
        stage('BNN end2end - U250') {
          when {
            expression { return params['end2end'] }
          }
          agent {
            label 'finn-build'
          }
          environment {
            BOARD = "U250"
            TEST_NAME = "bnn_build_full"
            FINN_HOST_BUILD_DIR = "${env.FINN_HOST_BUILD_DIR}/${env.TEST_NAME}_${env.BOARD}"
          }
          steps {
            catchError(stageResult: 'FAILURE') {
              script {
                // Creates dir in finn clone to store build files for stashing
                sh "mkdir -p ${env.TEST_NAME}"
                cleanPreviousBuildFiles(env.FINN_HOST_BUILD_DIR)

                // Pass in the marker to run with pytest and the XML test results filename
                runDockerPytestWithMarker("bnn_u250", "${env.TEST_NAME}_${env.BOARD}", '')
                findCopyZip(env.BOARD, env.FINN_HOST_BUILD_DIR, env.TEST_NAME)

                // Stash the test results file(s)
                stash name: "${env.TEST_NAME}_${env.BOARD}", includes: "${env.TEST_NAME}_${env.BOARD}.xml,${env.TEST_NAME}_${env.BOARD}.html"

                // Use an env variable to help collect test results later in pipeline
                env.BNN_BUILD_U250 = "SUCCESS"
              }
            }
          }
        }
        stage('BNN end2end - Pynq-Z1') {
          when {
            expression { return params['end2end'] }
          }
          agent {
            label 'finn-build'
          }
          environment {
            BOARD = "Pynq-Z1"
            TEST_NAME = "bnn_build_full"
            FINN_HOST_BUILD_DIR = "${env.FINN_HOST_BUILD_DIR}/${env.TEST_NAME}_${env.BOARD}"
          }
          steps {
            catchError(stageResult: 'FAILURE') {
              script {
                // Creates dir in finn clone to store build files for stashing
                sh "mkdir -p ${env.TEST_NAME}"
                cleanPreviousBuildFiles(env.FINN_HOST_BUILD_DIR)

                // Pass in the marker to run with pytest and the XML test results filename
                runDockerPytestWithMarker("bnn_pynq", "${env.TEST_NAME}_${env.BOARD}", '')
                findCopyZip(env.BOARD, env.FINN_HOST_BUILD_DIR, env.TEST_NAME)

                // Stash the test results file(s)
                stash name: "${env.TEST_NAME}_PynqZ1", includes: "${env.TEST_NAME}_${env.BOARD}.xml,${env.TEST_NAME}_${env.BOARD}.html"

                // Use an env variable to help collect test results later in pipeline
                env.BNN_BUILD_PYNQZ1 = "SUCCESS"
              }
            }
          }
        }
        stage('BNN end2end - ZCU104') {
          when {
            expression { return params['end2end'] }
          }
          agent {
            label 'finn-build'
          }
          environment {
            BOARD = "ZCU104"
            TEST_NAME = "bnn_build_full"
            FINN_HOST_BUILD_DIR = "${env.FINN_HOST_BUILD_DIR}/${env.TEST_NAME}_${env.BOARD}"
          }
          steps {
            catchError(stageResult: 'FAILURE') {
              script {
                // Creates dir in finn clone to store build files for stashing
                sh "mkdir -p ${env.TEST_NAME}"
                cleanPreviousBuildFiles(env.FINN_HOST_BUILD_DIR)

                // Pass in the marker to run with pytest and the XML test results filename
                runDockerPytestWithMarker("bnn_zcu104", "${env.TEST_NAME}_${env.BOARD}", '')
                findCopyZip(env.BOARD, env.FINN_HOST_BUILD_DIR, env.TEST_NAME)

                // Stash the test results file(s)
                stash name: "${env.TEST_NAME}_${env.BOARD}", includes: "${env.TEST_NAME}_${env.BOARD}.xml,${env.TEST_NAME}_${env.BOARD}.html"

                // Use an env variable to help collect test results later in pipeline
                env.BNN_BUILD_ZCU104 = "SUCCESS"
              }
            }
          }
        }
        stage('BNN end2end - KV260_SOM') {
          when {
            expression { return params['end2end'] }
          }
          agent {
            label 'finn-build'
          }
          environment {
            BOARD = "KV260_SOM"
            TEST_NAME = "bnn_build_full"
            FINN_HOST_BUILD_DIR = "${env.FINN_HOST_BUILD_DIR}/${env.TEST_NAME}_${env.BOARD}"
          }
          steps {
            catchError(stageResult: 'FAILURE') {
              script {
                // Creates dir in finn clone to store build files for stashing
                sh "mkdir -p ${env.TEST_NAME}"
                cleanPreviousBuildFiles(env.FINN_HOST_BUILD_DIR)

                // Pass in the marker to run with pytest and the XML test results filename
                runDockerPytestWithMarker("bnn_kv260", "${env.TEST_NAME}_${env.BOARD}", '')
                findCopyZip(env.BOARD, env.FINN_HOST_BUILD_DIR, env.TEST_NAME)

                // Stash the test results file(s)
                stash name: "${env.TEST_NAME}_${env.BOARD}", includes: "${env.TEST_NAME}_${env.BOARD}.xml,${env.TEST_NAME}_${env.BOARD}.html"

                // Use an env variable to help collect test results later in pipeline
                env.BNN_BUILD_KV260_SOM = "SUCCESS"
              }
            }
          }
        }
      }
    }
    stage('Check Stage Results') {
      agent {
        label 'finn-build'
      }
      steps {
        script {
          sh 'mkdir -p reports'
          cleanPreviousBuildFiles('reports')
          dir('reports') {
            // Only unstash for stages that ran
            unstashSuccessfulStage(env.SANITY_UT, "sanity_ut")
            unstashSuccessfulStage(env.FPGADATAFLOW_RESULT, "fpgadataflow")
            unstashSuccessfulStage(env.BNN_BUILD_SANITY, "bnn_build_sanity")
            unstashSuccessfulStage(env.END2END_RESULT, "end2end")
            unstashSuccessfulStage(env.BNN_BUILD_U250, "bnn_build_full_U250")
            unstashSuccessfulStage(env.BNN_BUILD_PYNQZ1, "bnn_build_full_PynqZ1")
            unstashSuccessfulStage(env.BNN_BUILD_ZCU104, "bnn_build_full_ZCU104")
            unstashSuccessfulStage(env.BNN_BUILD_KV260_SOM, "bnn_build_full_KV260_SOM")
          }

          // Combine individual HTML files to one single report
          sh './run-docker.sh pytest_html_merger -i reports/ -o reports/test_report_final.html'

          // Archive the XML & HTML test results
          archiveArtifacts artifacts: "reports/*.xml"
          archiveArtifacts artifacts: "reports/*.html"

          // Plot what XML files were created during the test run
          junit 'reports/*.xml'
        }
      }
    }
  }
}

void cleanPreviousBuildFiles(String buildDir) {
  // Delete any build files from a previous build
  // Previous build folders affect findCopyZip() and can cause the stage to fail
  if (!buildDir.empty) {
      sh "rm -rf ${buildDir}"
  }
}

void createMultiMarkerScript(String markers, String testResultsFilename, String additionalOptions) {
  // Passing multiple markers when running ./run-docker.sh does not work with bash.
  // Therefore, create a script to maintain the single quotes that surround the markers
  sh """echo "#!/bin/bash
python -m pytest -m \'${markers}\' --junitxml=${testResultsFilename}.xml --html=${testResultsFilename}.html --self-contained-html ${additionalOptions}" >> run-tests.sh
    """

  // Give permissions to script
  sh 'chmod 777 run-tests.sh'
}

void runDockerPytestWithMarker(String marker, String testResultsFilename, String additionalOptions) {
  sh """./run-docker.sh python -m pytest -m ${marker} --junitxml=${testResultsFilename}.xml --html=${testResultsFilename}.html --self-contained-html ${additionalOptions}"""
}

def findBoardBuildFiles(String searchDir, String dirToFind) {
  def result = sh(script: "find $searchDir -type d -name \"$dirToFind*\"", returnStdout: true).trim()
  if (result.empty) {
      error "Directory containing '$dirToFind' not found."
  }
	return result
}

void findCopyZip(String board, String findDir, String copyDir) {
  def buildDir = findBoardBuildFiles(findDir, "hw_deployment_${board}")
  sh "cp -r ${buildDir}/${board} ${copyDir}/"
  dir(copyDir) {
    sh "zip -r ${board}.zip ${board}/"
    sh "mkdir -p ${env.ARTIFACT_DIR}/${copyDir}/"
    sh "cp ${board}.zip ${env.ARTIFACT_DIR}/${copyDir}/"
  }
}

void unstashSuccessfulStage(String stageEnvVariableSet, String stashName) {
  if (stageEnvVariableSet) {
    unstash stashName
  }
}

void archiveSuccessfulStage(String stageEnvVariableSet, String folder) {
  if (stageEnvVariableSet) {
    archiveArtifacts artifacts: "${folder}/**/*"
  }
}
