---
# This role is for developer machines. The role for vehicle images is imaging/driving_os/roles/nvidia/.

- name: Ensure service facts have been gathered
  service_facts:

- name: Fail on horizon view VDI virtual machines
  fail:
    msg: "The --nvidia flag is not suported on Virtual Desktops. If a different driver version is required please open an ISD ticket by sending an email to ithelp@zoox.com."
  when: "'viewagent' in services"

- block:

  # Following deps are installed when you run apt-get install nvidia-375 libcuda-375
  - name: Install Nvidia driver dependencies
    apt:
      name:
        - "dkms"
        - "lib32gcc1"
        - "libc6-i386"
        - "libvdpau1"
        - "bbswitch-dkms"
        - "screen-resolution-extra"
      force: yes
      install_recommends: no
      cache_valid_time: "{{ apt_cache_valid_time }}"

  - name: Remove old Nvidia packages
    apt:
      name:
        - "libcuda1-???"
        - "nvidia-???"
        - "nvidia-settings"
      state: absent
      purge: yes

  - name: Check driver version
    command: "nvidia-smi --query-gpu=driver_version --format=csv,noheader -i 0"
    register: driver_version
    ignore_errors: yes

  - name: Prepare installer block
    block:
      - name: Download nvidia driver
        get_url:
          url: "{{ nvidia_download_url }}"
          dest: /opt
          mode: 755
          checksum: "sha256:{{ nvidia_SHA256 }}"

      - name: Extract nvidia driver
        command: "./{{ nvidia_package_version }}.run --extract-only --silent"
        args:
          chdir: "/opt"

      - name: Uninstall previous driver
        command: "/opt/{{ nvidia_package_version }}/nvidia-installer --silent --uninstall"

      - name: Patch dkms conf
        patch:
          src: 001-ignore-preempt-rt.patch
          dest: "/opt/{{ nvidia_package_version }}/kernel/dkms.conf"

      - name: Overwrite nvidia installer
        get_url:
          url: "{{ nvidia_installer_download_url }}"
          dest: /opt/{{ nvidia_package_version }}/nvidia-installer
          mode: 755
          checksum: "sha256:{{ nvidia_installer_SHA256 }}"
          force: yes

      - name: Install nvidia driver
        command: "/opt/{{ nvidia_package_version }}/nvidia-installer --silent --no-x-check --dkms"

    always:
      - name: Clean nvidia install files
        file:
          path: "{{ item }}"
          state: absent
        with_items:
          - /opt/{{ nvidia_package_version }}.run
          - /opt/{{ nvidia_package_version }}

    when: driver_version is failed or driver_version.stdout.find(nvidia_version) == -1

  - name: Install Nvidia related utilities
    apt:
      name:
        - "nvidia-prime"
      install_recommends: no
      cache_valid_time: "{{ apt_cache_valid_time }}"

  become: yes
