#!/usr/bin/env bash

# Required APT packages to be installed (supports regex matching)
REQUIRED_APT_PACKAGES=(
    "epc670-tofhat"
    "epc-drivers-dkms"
    "libepctofcam-dev"
)

# Required kernel modules to be loaded (supports regex matching)
REQUIRED_MODULES=(
    "cam_epc670"
    "clk_si5351"
    "rp1_cfe|rp1_cfe_downstream"
    "i2c_dev"
    "mtdblock"
)

# Kernel modules that must not produce dmesg errors
DMESG_MODULES_NOERR=(
    "epc670"
    "si5351"
    "rp1-cfe"
)

# Kernel modules that must not produce dmesg warnings
DMESG_MODULES_NOWARN=(
    "si5351"
    "rp1-cfe"
)

# Mountpoint for SPI flash
SPI_FLASH_MOUNT="/mnt/spi_flash"
SPI_CALIBRATION_DIR="$SPI_FLASH_MOUNT"/calibration

# Local directory for calibration data
USR_EPCTOFCAM_DIR="/usr/local/share/epctofcam"
USR_DIAGNOSTICS_DIR="$USR_EPCTOFCAM_DIR"/diagnostics
USR_CALIBRATION_DIR="$USR_EPCTOFCAM_DIR"/calibration
USR_CALIBRATION_FILE="$USR_CALIBRATION_DIR"/calibration_info.json

# I2C addresses
I2C_EPC670=/sys/bus/i2c/devices/*-003c
I2C_SI5351=/sys/bus/i2c/devices/*-0060
I2C_TMP112_LED1=/sys/bus/i2c/devices/*-0048
I2C_TMP112_LED2=/sys/bus/i2c/devices/*-0049

# Clock generator debugfs directories
CLKGEN_MAINCLK_DEBUGFS_DIR="/sys/kernel/debug/clk/clk2/"
CLKGEN_DATACLK_DEBUGFS_DIR="/sys/kernel/debug/clk/clk1/"
CLKGEN_MODCLK_DEBUGFS_DIR="/sys/kernel/debug/clk/clk0/"

# Hardware monitoring API directories 
HWMON_EPC670=$I2C_EPC670/hwmon/hwmon*
HWMON_LED1=$I2C_TMP112_LED1/hwmon/hwmon*
HWMON_LED2=$I2C_TMP112_LED2/hwmon/hwmon*

# V4L2 API devices
VIDEO_DEV="/run/epc670/video"
VIDEO_SUBDEV="/run/epc670/subdev"

# Python script for testing the camera sensor
PYTHON_SCRIPT=/usr/lib/epc670-tofhat-utils/epc670-test-imager.py

# Counters for summary
PASS=0
FAIL=0
SKIP=0
FAILED_NAMES=()
FAILED_MSGS=()
SKIPPED_NAMES=()
SKIPPED_MSGS=()


report_pass() {
    local name="$1" msg="$2"
    echo -e "\033[0;32m[PASS]\033[0m $name${msg:+ - $msg}"
    (( PASS++ ))
}

report_fail() {
    local name="$1" msg="$2"
    echo -e "\033[0;31m[FAIL]\033[0m $name${msg:+ - $msg}"
    (( FAIL++ ))
    FAILED_NAMES+=("$name")
    FAILED_MSGS+=("$msg")
}

report_skip() {
    local name="$1" msg="$2"
    echo -e "\033[0;33m[SKIP]\033[0m $name${msg:+ - $msg}"
    (( SKIP++ ))
    SKIPPED_NAMES+=("$name")
    SKIPPED_MSGS+=("$msg")
}

report() {
    local name="$1" status="$2" msg="$3"
    if [[ "$status" == "0" ]]; then
        report_pass "$name" "$msg"
    else
        report_fail "$name" "$msg"
    fi
}


section() { echo; echo "--- $1 ---"; }


# generate "full-scale" test pattern (vertical stripes)
gen_test_pattern() {
    for i in {1..38400}; do
        echo -n -e \\x00\\x00\\xf0\\xff
    done
}


check_packages_installed() {
    section "Installed Packages"

    for pattern in "${REQUIRED_APT_PACKAGES[@]}"; do
        local match=$(dpkg --list | grep -E "^ii" \
            | grep -Eo "$pattern" | head -n 1)
        if [[ -n "$match" ]]; then
            report_pass "Package is installed" "$match"
        else
            report_fail "Package is installed" "$pattern"
        fi
    done
}

check_modules_loaded() {
    section "Kernel Modules"

    for pattern in "${REQUIRED_MODULES[@]}"; do
        local match=$(find "/sys/module" -maxdepth 1 -type d \
            | grep -Eo "$pattern" | head -n 1)
        if [[ -n "$match" ]]; then
            report_pass "Kernel module is loaded" "$match"
        else
            report_fail "Kernel module is loaded" "$pattern"
        fi
    done
}

check_kernel_ring_buffer() {
    section "Kernel Ring Buffer"

    for module in "${DMESG_MODULES_NOERR[@]}"; do
        ! dmesg --level=err | grep "$module"
        report "DMESG check for errors" $? $module
    done

    for module in "${DMESG_MODULES_NOWARN[@]}"; do
        ! dmesg --level=warn | grep "$module"
        report "DMESG check for warnings" $? $module
    done
}

check_spi_flash() {
    section "SPI Flash"

    findmnt --mountpoint "$SPI_FLASH_MOUNT" --types "jffs2" > /dev/null
    local mount_ok=$?
    report "SPI flash mounted" $mount_ok "$SPI_FLASH_MOUNT"

    if [[ $mount_ok -ne 0 ]]; then
        report_skip "SPI flash contains calibration dir" "SPI flash not mounted"
        report_skip "SPI flash rsynced" "SPI flash not mounted"
        return
    fi

    [[ -d "$SPI_CALIBRATION_DIR" ]]
    report "SPI flash contains calibration dir" $? "$SPI_CALIBRATION_DIR"

    diff -qr "$SPI_FLASH_MOUNT" "$USR_EPCTOFCAM_DIR"
    report "SPI flash rsynced" $? "$USR_EPCTOFCAM_DIR"
}

check_calibration_data() {
    section "Calibration Data"

    [[ -d "$USR_EPCTOFCAM_DIR" ]]
    report "EPCTofcam directory exists" $? "$USR_EPCTOFCAM_DIR"

    [[ -d "$USR_CALIBRATION_DIR" ]]
    report "Calibration directory exists" $? "$USR_CALIBRATION_DIR"

    [[ -f "$USR_CALIBRATION_FILE" ]]
    report "Calibration info file exists" $? "$USR_CALIBRATION_FILE"

    [[ -f "$USR_CALIBRATION_DIR"/grayscale_drnu.bin ]]
    report "Calibration DSNU image exists" $? "$USR_CALIBRATION_DIR"/grayscale_drnu.bin

    [[ -f "$USR_CALIBRATION_DIR"/grayscale_prnu.bin ]]
    report "Calibration PRNU image exists" $? "$USR_CALIBRATION_DIR"/grayscale_prnu.bin
    
    find "$USR_CALIBRATION_DIR"/offset_*.bin > /dev/null
    report "Calibration offset image exists" $? "$USR_CALIBRATION_DIR"/offset_*.bin

    find "$USR_CALIBRATION_DIR"/phase_error_LUT_*.bin > /dev/null
    report "Calibration DRNU LUT exists" $? "$USR_CALIBRATION_DIR"/phase_error_LUT_*.bin
}

check_i2c_devices() {
    section "I2C devices"

    grep -q "epc670" $I2C_EPC670/name
    report "I2C check address for epc670" $? $I2C_EPC670

    grep -q "si5351" $I2C_SI5351/name
    report "I2C check address for clock generator" $? $I2C_SI5351

    grep -q "tmp112" $I2C_TMP112_LED1/name
    report "I2C check address for led1 temp sensor" $? $I2C_TMP112_LED1

    grep -q "tmp112" $I2C_TMP112_LED2/name
    report "I2C check address for led2 temp sensor" $? $I2C_TMP112_LED1
}

check_clock_generator() {
    section "Clock Generator (si5351a)"

    local rate=$(cat "$CLKGEN_MAINCLK_DEBUGFS_DIR"/clk_rate)
    report "CLKGEN read main clock rate (debugfs)" $? "$rate Hz"

    local rate=$(cat "$CLKGEN_DATACLK_DEBUGFS_DIR"/clk_rate)
    report "CLKGEN read data clock rate (debugfs)" $? "$rate Hz"
    
    local rate=$(cat "$CLKGEN_MODCLK_DEBUGFS_DIR"/clk_rate)
    report "CLKGEN read modulation clock rate (debugfs)" $? "$rate Hz"

}

check_hwmon_temperatures() {
    section "HWMon (Temperature Sensors)"

    local temp_min=10000 # 10°C
    local temp_max=99000 # 99°C

    local temp=$(cat $HWMON_EPC670/temp1_input)
    (( $temp > temp_min && $temp < temp_max ))
    report "HWMON read epc670 temperature 1" $? "$temp mC°"

    local temp=$(cat $HWMON_EPC670/temp2_input)
    (( $temp > temp_min && $temp < temp_max ))
    report "HWMON read epc670 temperature 2" $? "$temp mC°"

    local temp=$(cat $HWMON_EPC670/temp3_input)
    (( $temp > temp_min && $temp < temp_max ))
    report "HWMON read epc670 temperature 3" $? "$temp mC°"

    local temp=$(cat $HWMON_EPC670/temp4_input)
    (( $temp > temp_min && $temp < temp_max ))
    report "HWMON read epc670 temperature 4" $? "$temp mC°"

    local temp=$(cat $HWMON_LED1/temp1_input)
    (( $temp > temp_min && $temp < temp_max ))
    report "HWMON read led1 temperature" $? "$temp mC°"

    local temp=$(cat $HWMON_LED2/temp1_input)
    (( $temp > temp_min && $temp < temp_max ))
    report "HWMON read led2 temperature" $? "$temp mC°"
}

check_v4l2_api() {
    section "V4L2 API"

    v4l2-ctl -d "$VIDEO_DEV"
    local dev_ok=$?
    report "V4L2 open device" $dev_ok

    v4l2-ctl -d "$VIDEO_SUBDEV"
    local subdev_ok=$?
    report "V4L2 open subdevice" $subdev_ok

    if [[ $dev_ok -ne 0 ]] || [[ $subdev_ok -ne 0 ]]; then
        report_skip "V4L2 get video format"
        report_skip "V4L2 set acquisiton mode"
        report_skip "V4L2 get acquisiton mode"
        report_skip "V4L2 set modulation frequency"
        report_skip "V4L2 get modulation frequency"
        report_skip "V4L2 set exposure"
        report_skip "V4L2 get exposure"
        report_skip "V4L2 capture test pattern"
        report_skip "V4L2 compare test pattern"
        return
    fi

    local dims=$(v4l2-ctl -d "$VIDEO_DEV" --get-fmt-video | \
     awk '/Width\/Height/ { split($NF, a, "/"); printf "%sx%s", a[1], a[2] }')
    [[ -n "$dims" ]]
    report "V4L2 get video format" $? "$dims"

    v4l2-ctl -d "$VIDEO_SUBDEV" --set-ctrl="acquisition_type=1" # grayscale
    report "V4L2 set acquisiton mode" $?

    v4l2-ctl -d "$VIDEO_SUBDEV" --get-ctrl="acquisition_type" | grep -q "grayscale"
    report "V4L2 get acquisiton mode" $?

    v4l2-ctl -d "$VIDEO_SUBDEV" --set-ctrl="mod_freq_hz=10000000" # 10 MHz
    report "V4L2 set modulation frequency" $?

    v4l2-ctl -d "$VIDEO_SUBDEV" --get-ctrl="mod_freq_hz" | grep -q "10000000"
    report "V4L2 get modulation frequency" $?

    v4l2-ctl -d "$VIDEO_SUBDEV" --set-ctrl="exposure=400" # 400 us
    report "V4L2 set exposure" $?

    v4l2-ctl -d "$VIDEO_SUBDEV" --get-ctrl="exposure" | grep -q "400"
    report "V4L2 get exposure" $?


    # Capture test pattern:

    # create temporary output file
    local img_test_pattern=$(mktemp)

    # launch stream in background
    v4l2-ctl -d "$VIDEO_DEV" \
        --stream-mmap=1 \
        --stream-count=1 \
        --stream-to="$img_test_pattern" \
        2>/dev/null &
    local stream_pid=$!

    # give camera some time to start the stream
    sleep 0.1

    # trigger test pattern "0" (full-scale vertical stripes)
    v4l2-ctl -d "$VIDEO_SUBDEV" --set-ctrl="test_pattern=0"
    report "V4L2 capture test pattern" $?
    
    # wait for background stream process to finish
    wait $stream_pid

    # compare to expected test pattern
    gen_test_pattern | cmp -s "$img_test_pattern"
    report "V4L2 compare test pattern" $?
    rm -f $img_test_pattern
}

check_sensor_response() {
    section "Sensor Response"

    python -c "import epc_tofcam_native"
    local deps_ok=$?
    report "SENSOR check python dependencies" $deps_ok "epc_tofcam_native"

    if [[ $deps_ok -ne 0 ]]; then
        report_skip "SENSOR check DSNU image"
        report_skip "SENSOR check amplitude exposure scaling"
        report_skip "SENSOR check amplitude exposure scaling"
        report_skip "SENSOR check amplitude-modfreq scaling"
        report_skip "SENSOR check grayscale-exposure scaling"
        report_skip "SENSOR check LED delay scaling"
        report_skip "SENSOR check grayscale-modfreq scaling"
        report_skip "SENSOR check reference images"
        return
    fi

    # for testing, switch to LED inside lens holder
    v4l2-ctl -d "$VIDEO_SUBDEV" --set-ctrl="led_select=0" &> /dev/null

    [[ -f "$USR_CALIBRATION_DIR"/grayscale_drnu.bin ]]
    local calib_ok=$?

    if [[ $calib_ok -ne 0 ]]; then
        report_skip "SENSOR check DSNU image" "missing calibration data"
        report_skip "SENSOR check grayscale exposure scaling" "missing calibration data"
        report_skip "SENSOR check grayscale modfreq scaling" "missing calibration data"
    else
        python "$PYTHON_SCRIPT" check_dsnu_image
        report "SENSOR check DSNU image" $?
            
        python "$PYTHON_SCRIPT" check_grayscale_exposure_scaling
        report "SENSOR check grayscale exposure scaling" $?

        python "$PYTHON_SCRIPT" check_grayscale_modfreq_scaling
        report "SENSOR check grayscale modfreq scaling" $?
    fi

    python "$PYTHON_SCRIPT" check_amplitude_exposure_scaling
    report "SENSOR check amplitude exposure scaling" $?

    python "$PYTHON_SCRIPT" check_amplitude_modfreq_scaling
    report "SENSOR check amplitude modfreq scaling" $?

    python "$PYTHON_SCRIPT" check_led_delay_scaling
    report "SENSOR check LED delay scaling" $?

    python "$PYTHON_SCRIPT" check_reference_images
    report "SENSOR check reference images" $?

    # reset uncommonly used controls to default for user convenience
    v4l2-ctl -d "$VIDEO_SUBDEV" --set-ctrl="led_select=1" &> /dev/null
    v4l2-ctl -d "$VIDEO_SUBDEV" --set-ctrl="led_delay=0" &> /dev/null
}

echo "epc670 Diagnostics"
echo "=================="

check_packages_installed
check_modules_loaded
check_kernel_ring_buffer
check_spi_flash
check_calibration_data
check_i2c_devices
check_clock_generator
check_hwmon_temperatures
check_v4l2_api
check_sensor_response

echo
echo "=================="
echo "SUMMARY"
echo "=================="
echo "  PASS: $PASS"
echo "  FAIL: $FAIL"
echo "  SKIP: $SKIP"

if (( FAIL > 0 )); then
    echo
    echo "Failed checks:"
    for i in "${!FAILED_NAMES[@]}"; do
        msg="${FAILED_MSGS[$i]}"
        echo -e "  \033[0;31m[FAIL]\033[0m ${FAILED_NAMES[$i]}${msg:+ - $msg}"
    done
fi
if (( SKIP > 0 )); then
    echo "Skipped checks:"
    for i in "${!SKIPPED_NAMES[@]}"; do
        msg="${SKIPPED_MSGS[$i]}"
        echo -e "  \033[0;33m[SKIP]\033[0m ${SKIPPED_NAMES[$i]}${msg:+ - $msg}"
    done
fi

(( FAIL > 0 )) && exit 1 || exit 0
