#!/bin/bash
shopt -s extglob
set -eu

cleanup() {
rm -f "$teststdout" "$teststderr" "$camerastdout" \
      "$cameratmpimg" "$wifistdout" "$wifistderr" \
      "$power1stdout" "$power2stdout"
}

trap cleanup EXIT

analyze_image() {
v4l2-ctl --device /dev/video0 --stream-mmap --stream-count=1 --set-fmt-video=pixelformat=MJPG --stream-to="$cameratmpimg" 2> /dev/null &&
export cameratmpimg

python3 -c '
from os import environ
import matplotlib.image as mpimg
import numpy as np
img = mpimg.imread(environ["cameratmpimg"])

data = np.unique(img.reshape(-1, img.shape[-1]), axis=0, return_counts=True)[0]

print(len(data))' > "$camerastdout" 2>/dev/null &
}

camera_test_status() {
  if [ -z "$(cat "$camerastdout")" ]; then
    echo "Scanning...";
  else
    if [ "$(cat "$camerastdout")" -gt 20000 ]; then
      echo "OK ""$(cat "$camerastdout")";
    else
      echo "Bad camera ""$(cat "$camerastdout")";
    fi
  fi
}

test_camera() {
  env CACA_DRIVER=ncurses mpv --really-quiet -frames 200 --vo=caca av://v4l2:/dev/video0
}

product_name() {
 cat /sys/class/dmi/id/product_name
}

test_disks() {
  local all_output disk_devices error_output letter mask smartctl_code \
        smartctl_output

  disk_devices=$(facter --json disks \
                   | jq -r '.disks | keys | .[]' \
                   | grep -E '^(mmcblk|nvme|sd|vd|xvda)')

  all_output=''
  error_output=''

  for disk in $disk_devices; do
    smartctl_code=0
    smartctl_output=$(smartctl -a "/dev/${disk}" 2>/dev/null) || smartctl_code=$?

    all_output="${all_output}${smartctl_output}

"

    if [ "$smartctl_code" -eq 0 ]; then
      error_output="${error_output}No errors on /dev/${disk} ::

"
    else
      error_output="${error_output}Some errors on /dev/${disk} ::
"
      mask=1
      for i in 0 1 2 3 4 5 6 7; do
        case "$i" in

          0) letter='P'; errormsg='Command line did not parse.';;
          1) letter='O'; errormsg='Device open failed.';;
          2) letter='B'; errormsg='Some SMART or other ATA command to the disk failed.';;
          3) letter='F'; errormsg='SMART status check returned "DISK FAILING".';;
          4) letter='A'; errormsg='We found prefail Attributes <= threshold.';;
          5) letter='T'; errormsg='Some Attributes have been <= threshold in the past.';;
          6) letter='E'; errormsg='The device error log contains records of errors.';;
          7) letter='S'; errormsg='The device self-test log contains records of errors.';;
        esac
        if [ "$((($smartctl_code & $mask) && 1))" -eq 1 ]; then
          error_output="${error_output}   ${letter}  :: ${errormsg}
"
        else
          error_output="${error_output}  (${letter}) :: -
"
        fi
        mask=$(($mask << 1))
      done
      error_output="${error_output}
"
    fi
  done

  dialog --no-collapse --msgbox "${error_output}${all_output}" -1 -1
}

bt_test_status() {
  a="$(wc --lines "$teststderr" | awk '{print $1}')"
  b="$(wc --lines "$teststdout" | awk '{print $1}')"
  if [ "$a" -gt "0" ]; then
    echo "ERROR"
  elif [ "$b" -gt "1" ]; then
    echo "OK"
  else
    echo "Scanning..."
  fi
}

bt_test_results() {
  clear
  cat "$teststderr" "$teststdout"
  sleep 2
}

test_keyboard() {
  local ruby_test_keyboard_script

ruby_test_keyboard_script=$(cat <<'EOF'
  def show_state
    system('clear')
    puts "Press keys until all key labels disappear."
    puts "You need to use the real keyboard, this does not work through ssh."
    puts "This will quit in case no keyboard input happens in 10 seconds."
    puts ""
    puts "Remember to check the physical layout of the keyboard! (QWERTY/QWERTZ/AZERTY etc.)"
    puts ""

    $lines.each { |line| puts "    #{ line }" }
  end

  BEGIN {
    all_keys = {
       1  => [ 0, 0,  3 ],    # ESC
       59 => [ 0, 5,  2 ],    # F1
       60 => [ 0, 8,  2 ],    # F2
       61 => [ 0, 11, 2 ],    # F3
       62 => [ 0, 14, 2 ],    # F4
       63 => [ 0, 17, 2 ],    # F5
       64 => [ 0, 20, 2 ],    # F6
       65 => [ 0, 23, 2 ],    # F7
       66 => [ 0, 26, 2 ],    # F8
       67 => [ 0, 29, 2 ],    # F9
       68 => [ 0, 32, 3 ],    # F10
       87 => [ 0, 36, 3 ],    # F11
       88 => [ 0, 40, 3 ],    # F12
      110 => [ 0, 44, 6 ],    # INSERT
       99 => [ 0, 51, 4 ],    # PSCR
      111 => [ 0, 56, 6 ],    # DELETE

       41 => [ 1,  0, 1 ],    # §
        2 => [ 1,  5, 1 ],    # 1
        3 => [ 1,  8, 1 ],    # 2
        4 => [ 1, 11, 1 ],    # 3
        5 => [ 1, 14, 1 ],    # 4
        6 => [ 1, 17, 1 ],    # 5
        7 => [ 1, 20, 1 ],    # 6
        8 => [ 1, 23, 1 ],    # 7
        9 => [ 1, 26, 1 ],    # 8
       10 => [ 1, 29, 1 ],    # 9
       11 => [ 1, 33, 1 ],    # 0
       12 => [ 1, 37, 1 ],    # +
       13 => [ 1, 41, 1 ],    # ´
       14 => [ 1, 45, 9 ],    # BACKSPACE
      102 => [ 1, 57, 4 ],    # HOME
       69 => [ 1, 67, 3 ],    # NUMLOCK
       98 => [ 1, 71, 1 ],    # Numpad /
       55 => [ 1, 73, 1 ],    # Numpad *
       74 => [ 1, 75, 1 ],    # Numpad -

       15 => [ 2, 0,  3 ],    # TAB
       16 => [ 2, 6,  1 ],    # Q
       17 => [ 2, 9,  1 ],    # W
       18 => [ 2, 12, 1 ],    # E
       19 => [ 2, 15, 1 ],    # R
       20 => [ 2, 18, 1 ],    # T
       21 => [ 2, 21, 1 ],    # Y
       22 => [ 2, 24, 1 ],    # U
       23 => [ 2, 27, 1 ],    # I
       24 => [ 2, 30, 1 ],    # O
       25 => [ 2, 34, 1 ],    # P
       26 => [ 2, 38, 1 ],    # Å
       27 => [ 2, 42, 1 ],    # ¨
      104 => [ 2, 57, 4 ],    # PgUp
       71 => [ 2, 67, 1 ],    # Numpad 7
       72 => [ 2, 69, 1 ],    # Numpad 8
       73 => [ 2, 71, 1 ],    # Numpad 9
       78 => [ 2, 75, 1 ],    # Numpad +

       58 => [ 3,  0, 4 ],    # CAPS
       30 => [ 3,  6, 1 ],    # A
       31 => [ 3,  9, 1 ],    # S
       32 => [ 3, 12, 1 ],    # D
       33 => [ 3, 15, 1 ],    # F
       34 => [ 3, 18, 1 ],    # G
       35 => [ 3, 21, 1 ],    # H
       36 => [ 3, 24, 1 ],    # J
       37 => [ 3, 27, 1 ],    # K
       38 => [ 3, 30, 1 ],    # L
       39 => [ 3, 34, 1 ],    # Ö
       40 => [ 3, 38, 1 ],    # Ä
       43 => [ 3, 42, 1 ],    # '
       28 => [ 3, 49, 5 ],    # ENTER
      109 => [ 3, 57, 4 ],    # PgDn
       75 => [ 3, 67, 1 ],    # Numpad 4
       76 => [ 3, 69, 1 ],    # Numpad 5
       77 => [ 3, 71, 1 ],    # Numpad 6

       42 => [ 4,  0, 4 ],    # LSFT
       86 => [ 4,  5, 1 ],    # <
       44 => [ 4,  7, 1 ],    # Z
       45 => [ 4, 10, 1 ],    # X
       46 => [ 4, 13, 1 ],    # C
       47 => [ 4, 16, 1 ],    # V
       48 => [ 4, 19, 1 ],    # B
       49 => [ 4, 22, 1 ],    # N
       50 => [ 4, 25, 1 ],    # M
       51 => [ 4, 28, 1 ],    # ,
       52 => [ 4, 31, 1 ],    # .
       53 => [ 4, 35, 1 ],    # -
       54 => [ 4, 48, 6 ],    # RSHIFT
      107 => [ 4, 57, 3 ],    # END
       79 => [ 4, 67, 1 ],    # Numpad 1
       80 => [ 4, 69, 1 ],    # Numpad 2
       81 => [ 4, 71, 1 ],    # Numpad 3

       29 => [ 5, 0,  4 ],    # CTRL
      125 => [ 5, 8,  3 ],    # WIN
       56 => [ 5, 12, 3 ],    # Alt
       57 => [ 5, 18, 5 ],    # SPACE
      100 => [ 5, 25, 5 ],    # AltGr
      127 => [ 5, 31, 3 ],    # HAM
       97 => [ 5, 35, 4 ],    # CTRL
      105 => [ 5, 44, 4 ],    # Left
      103 => [ 5, 49, 2 ],    # Up
      108 => [ 5, 52, 4 ],    # Down
      106 => [ 5, 57, 5 ],    # Right
       82 => [ 5, 67, 1 ],    # Numpad 0
       83 => [ 5, 69, 1 ],    # Numpad ,
       96 => [ 5, 71, 5 ],    # Numpad ENTER
    }

    $lines = [
      "ESC  F1 F2 F3 F4 F5 F6 F7 F8 F9 F10 F11 F12 INSERT PSCR DELETE",
      "§    1  2  3  4  5  6  7  8  9   0   +   ´   BACKSPACE   HOME      NUM / * -",
      "TAB   Q  W  E  R  T  Y  U  I  O   P   Å   ¨              PgUp      7 8 9   +",
      "CAPS  A  S  D  F  G  H  J  K  L   Ö   Ä   '      ENTER   PgDn      4 5 6",
      "LSFT < Z  X  C  V  B  N  M  ,  .   -            RSHIFT   END       1 2 3",
      "CTRL    WIN Alt   SPACE  AltGr HAM CTRL     Left Up Down Right     0 , ENTER" ]
  }

  if $_.match(/^keycode *(\d+) release/) then
    if all_keys.has_key?($1.to_i) then
      y, x, length = * all_keys[$1.to_i]
      (0..(length-1)).each do |i|
        $lines[y][x + i] = " "
      end
    end

    all_keys.delete($1.to_i)
  end

  show_state()

  exit(0) if all_keys.empty?

  END {
    if all_keys.empty? then
      puts ""
      puts "All keys were pressed, great!"
      exit(0)
    else
      warn ""
      warn "WARNING: some keys were not pressed!"
      exit(1)
    end
  }
EOF
)

  unbuffer showkey | ruby -ln -e "$ruby_test_keyboard_script"
}

run_mouse_test() {
  clear
  echo '>>>'
  echo ">>> Now running evtest ${1} for 4 seconds"
  echo '>>>'
  sleep 1
  timeout -k 1 4 evtest "$1" || true
  sleep 1
}

test_mouse() {
  local dialog_args i mouse_choice mouse_device_count mouse_event_devices \
        mouse_devpath

  mouse_event_devices=$(
    sed 's/=/ /' /proc/bus/input/devices \
      | awk '$1 == "H:" && $2 == "Handlers" && /mouse/ {
               for (i = 3; i <= NF; i++) { if ($i ~ /^event/) { print $i } } }
             ')

  if [ -z "$mouse_event_devices" ]; then
    dialog --msgbox 'No mouse devices could be found.' 5 36
    return 1
  fi

  mouse_device_count="$(printf "%s\n" "$mouse_event_devices" | wc -l)"
  if [ "$mouse_device_count" -eq 1 ]; then
    run_mouse_test "/dev/input/${mouse_event_devices}" || true
  else
    dialog_args=''
    i=1
    for mouse_event_dev in $mouse_event_devices; do
      dialog_args="${dialog_args} $i Mouse_${i}"
      i=$(($i + 1))
    done
    dialog_args="${dialog_args} exit Exit"

    default_mouse=1
    while true; do
      mouse_choice=$(dialog --stdout \
                            --default-item "$default_mouse" \
                            --menu 'Choose a mouse device' \
                      $(($mouse_device_count + 8)) 36 $(($mouse_device_count + 1)) \
                      $dialog_args) || true
      if [ -z "$mouse_choice" -o "$mouse_choice" = 'exit' ]; then
        break
      fi

      mouse_devpath=''
      i=1
      for mouse_event_dev in $mouse_event_devices; do
        if [ "$i" = "$mouse_choice" ]; then
          mouse_devpath="/dev/input/${mouse_event_dev}"
          default_mouse="$(($i + 1))"
          if [ "$default_mouse" -gt "$mouse_device_count" ]; then
            default_mouse='exit'
          fi
        fi
        i=$(($i + 1))
      done

      if [ -z "$mouse_devpath" ]; then
        echo 'Odd error, could not find mouse event device path' >&2
        return 1
      fi

      run_mouse_test "$mouse_devpath" || true
    done
  fi

  return 0
}

test_audio() {
  local alsacard alsamicdev alsaspkctl alsaspkctl2 alsaspkdev alsavolume
  alsacard='PCH'
  alsamicdev='default'
  alsaspkctl='Master'
  alsaspkctl2=''
  alsaspkdev='default'
  alsavolume=''

  ### Some device model/gen specific configurations
  case "$(product_name)" in
    "HP EliteBook 8"[345]"0 G6" | "HP EliteBook 8"[345]"0 G"[78]" Notebook PC")
      ### With these, speaker test works on G6-G8 devices. Mic test doesn't work just yet, but maybe soon.
      alsacard='0'
      alsaspkctl='Master'
      alsaspkctl2='Speaker'
      alsavolume="85%"
      ;;
    "HP EliteBook 820 G1"|"HP EliteBook 820 G2"|"HP EliteBook 840 G1"|"HP EliteBook 840 G2" )
      alsaspkdev='hw:1,0'
      alsavolume="85%"
      ;;
    *)
      ### Use the defaults specified above for all other models + settings that don't need to be modified
      ;;
  esac

  clear

  ### Use the default volume the system has (that usually works fine),
  ### unless 'alsavolume' is set (when a device is known to have very low volume by default).
  if [ -n "${alsavolume}" ]; then
    amixer --card "${alsacard}" set ${alsaspkctl} ${alsavolume}
  fi

  if [ -n "${alsaspkctl2}" ]; then
    amixer --card "${alsacard}" set ${alsaspkctl2} ${alsavolume}
  fi

  ### On some devices, VUMeter does not register changes in audio levels if we play the test with normal sensitivity. Probably something to do with the microphones location.
  ### We do not want to break our ears by boosting the audio itself, so lets just boost microphone input by a couple levels (a few dozen dB).
  amixer --card "${alsacard}" set 'Internal Mic Boost' 2
  aplay --device="${alsaspkdev}" \
        /usr/share/sounds/puavo/speakertest.wav & \
  arecord --vumeter=stereo -f cd -d 4 --device="${alsamicdev}" /dev/null && \
  echo 'Recording finished.'

  sleep 2
}

test_wifi() {
  ifconfig wlan0 up || true

  iwlist wlan0 scan 2>"$wifistderr" \
    | ruby -lne '
        BEGIN { networks = {} }

        if $_.match(/Signal level=(.*?) dBm/) then
          signal_level = $1
        end

        if $_.match(/ESSID:"(.*?)"/) then
          essid = $1
          if signal_level then
            networks[essid] = signal_level
          end
        end

        END {
          networks_list = networks.keys.sort_by do |i,j|
                            networks[i].to_i <= networks[j].to_i
                          end
          networks_list.each do |essid|
            printf("%-30s %s dBm\n", essid, networks[essid])
          end
        }
      '
}

wifi_status() {
  wifi="$(wc --lines "$wifistdout" | awk '{print $1}')"
  err="$(wc --lines "$wifistderr" | awk '{print $1}')"
  if [ "$wifi" -gt "0" ]; then
    echo "OK"
  elif [ "$err" -gt "0" ]; then
    echo "Bad wifi"
  elif [ "$wifi" -eq "0" ] && [ "$err" -eq "0" ]; then
    echo "Scanning..."
  fi
}

show_wifi() {
  a="$(wc --lines "$wifistdout" | awk '{print $1}')"
  if [ "$a" -eq "0" ]; then
    results="$(cat "$wifistderr")

No networks found.
Maybe a problem with wireless interface?"
    dialog --no-collapse --msgbox "$results" 9 48
    return 1
  fi
  results=$(cat "$wifistdout")
  dialog --no-collapse --msgbox "$results" 20 48
}

is_ac_connected() {
  [ "$(cat /sys/class/power_supply/AC/online 2>/dev/null)" -eq 1 ]
}

is_hub_connected() {
  [ -n "$(lsusb -t 2>/dev/null | grep 'hub' | grep 'Port 4')" ]
}

test_power() {
  while sleep 0.5
  do
    if is_ac_connected && is_hub_connected; then
      echo "Dock OK " > "$power1stdout"
      SECONDS=0
    fi
    if [ "$SECONDS" -gt "2" ]; then
      if is_ac_connected && ! is_hub_connected; then
        echo "AC OK" > "$power2stdout"
      fi
    fi
  done
}

power_status() {
  if [ -z "$(cat "$power1stdout")$(cat "$power2stdout")" ]; then
    echo "Scanning..."
  else
    echo "$(cat "$power1stdout")$(cat "$power2stdout")"
  fi
}

check_if_root() {
if [ "$(id -u)" -ne 0 ]; then
  echo 'Not run as root, some tests might not work!
Choose a test'
else
  echo 'Choose a test'
fi
}

teststdout=$(mktemp /tmp/puavo-test-hardware-bluetooth-stdout.XXXXXX)
teststderr=$(mktemp /tmp/puavo-test-hardware-bluetooth-stderr.XXXXXX)
camerastdout=$(mktemp /tmp/puavo-test-hardware-camera-stdout.XXXXXX)
cameratmpimg=$(mktemp /tmp/puavo-test-hardware-camera-tmpimg.XXXXXX)
wifistdout=$(mktemp /tmp/puavo-test-hardware-wifi-stdout.XXXXXX)
wifistderr=$(mktemp /tmp/puavo-test-hardware-wifi-stderr.XXXXXX)
power1stdout=$(mktemp /tmp/puavo-test-hardware-power1-stdout.XXXXXX)
power2stdout=$(mktemp /tmp/puavo-test-hardware-power2-stdout.XXXXXX)

analyze_image &

hcitool scan > "$teststdout" 2> "$teststderr" &
test_wifi > "$wifistdout" &
test_power &

mainmenu_default_item='audio'
while true; do
  response=$(dialog --stdout --default-item "$mainmenu_default_item" \
                    --menu "$(check_if_root)" 18 50 10 \
               audio      'Test audio devices' \
               disks      'Test disks' \
               keyboard   'Test keyboard' \
               mouse      'Test mouse/mice' \
               wifi       "$(wifi_status)" \
               bluetooth  "$(bt_test_status)" \
               camera     "$(camera_test_status)" \
               power      "$(power_status)" \
               quit       'Exit tests') || true

  case "$response" in
    audio)      mainmenu_default_item=keyboard;   test_audio      || true ;;
    disks)      mainmenu_default_item=keyboard;   test_disks      || true ;;
    keyboard)   mainmenu_default_item=mouse;      test_keyboard   || true ;;
    mouse)      mainmenu_default_item=quit;       test_mouse      || true ;;
    wifi)       mainmenu_default_item=quit;       show_wifi       || true ;;
    bluetooth)  mainmenu_default_item=quit;       bt_test_results || true ;;
    camera)     mainmenu_default_item=quit;       test_camera     || true ;;
    power)      mainmenu_default_item=quit;                          true ;;
    quit|'')    break ;;
    *)          echo 'Unknown option' >&2 ;;
  esac
done

exit 0
