---
# This will pass through the Coral PCIe Accelerator to the specific Proxmox virtual machine
- name: Set Coral device facts
  set_fact:
    device_id: "1ac1:089a"                    # Coral PCIe device ID
    proxmox_virtual_machine_passthrough: 101  # The ID of your Virtual Machine in Proxmox

- name: Check the Coral PCI device is found on the system
  become: true
  shell:
    cmd: lspci -nnk | grep {{ device_id}}
  register: coral_found
  ignore_errors: true

# If device ID not found
- when: device_id not in coral_found.stdout
  block:
    - name: Fail, cannot continue
      ansible.builtin.fail:
        msg: "Coral Device ID `{{ device_id }}` not found in `lspci -nnk`. Now failing."

# If device ID has been found
- name: Coral Device ID `{{ device_id }}` has been found. Continuing to pass the device through to container {{ proxmox_virtual_machine_passthrough }}
  when: device_id in coral_found.stdout
  block:
    - name: Result from the `lspci -nnk`
      debug: msg="{{ coral_found.stdout }}"

    - name: Updating Coral facts
      set_fact:
        # Split by ` System peripheral` and store the first value which is the PCI device bus number
        device_bus_number: "{{ coral_found.stdout.split(' System peripheral')[0] }}"

    # Device bus iommu group check
    - name: Device bus iommu group check
      when: device_bus_number is defined
      block:
        - name: Make the device bus is in an iommu group
          become: true
          shell:
            cmd: find /sys/kernel/iommu_groups/ -type l | grep {{ device_bus_number }}
          register: bus_iommu_groups_find

        - name: Result from `find /sys/kernel/iommu_groups/ -type l`
          debug: msg="{{ bus_iommu_groups_find.stdout }}"

        - name: Updating Coral facts
          set_fact:
            # Remove the `/sys/kernel/iommu_groups/` text and split by `/` and store the first value which is the iommu group ID number
            bus_iommu_group_id: "{{ (bus_iommu_groups_find.stdout|replace('/sys/kernel/iommu_groups/','')).split('/')[0] }}"

        - name: Storing group ID number to facts
          debug: msg="{{ bus_iommu_group_id }}"

        - name: Now check to make sure no other devices are using iommu group ID `{{ bus_iommu_group_id }}`. Passthrough won't work if more than one device in the same group.
          become: true
          shell:
            cmd: find /sys/kernel/iommu_groups/ -type l | grep /{{ bus_iommu_group_id }}/
          register: bus_iommu_groups_folder_find

        - name: Result from `find /sys/kernel/iommu_groups/ -type l | grep /{{ bus_iommu_group_id }}/`
          debug: msg="{{ bus_iommu_groups_folder_find.stdout }}"


        # Update Grub to force break-up of iommu devices if the device is sharing a group with other devices
        - name: If more than one device in the iommu group ID {{ bus_iommu_group_id }}
          when: bus_iommu_groups_folder_find.stdout_lines | length > 1
          block:
            - name: Note
              debug: msg="It appears more than one device in the iommu group ID {{ bus_iommu_group_id }}. This will require a GRUB amendment..."

            - name: Question
              ansible.builtin.pause:
                prompt: "Do you wish to update GRUB? (yes/no)"
              register: update_grub

            - name: Updating GRUB
              when: update_grub.user_input | lower == 'yes'
              block:
                - name: Set GRUB options for Intel systems to force break-up of iommu devices
                  become: true
                  ansible.builtin.lineinfile:
                    path: /etc/default/grub
                    regexp: '^GRUB_CMDLINE_LINUX_DEFAULT='
                    line: 'GRUB_CMDLINE_LINUX_DEFAULT="quiet intel_iommu=on pcie_acs_override=downstream,multifunction"'
                    backrefs: yes
                  when: ansible_processor[1] == 'GenuineIntel'

                - name: Set GRUB options for AMD systems to force break-up of iommu devices
                  become: true
                  ansible.builtin.lineinfile:
                    path: /etc/default/grub
                    regexp: '^GRUB_CMDLINE_LINUX_DEFAULT='
                    line: 'GRUB_CMDLINE_LINUX_DEFAULT="quiet amd_iommu=on pcie_acs_override=downstream,multifunction"'
                    backrefs: yes
                  when: ansible_processor[1] == 'AuthenticAMD'

                - name: Run `update-grub`
                  become: true
                  ansible.builtin.command:
                    cmd: update-grub

                - name: Rebooting machine
                  become: true
                  ansible.builtin.reboot:
                    reboot_timeout: 300
                    msg: "Rebooting machine in 5 seconds"
                    # reboot_timeout: 3600

                - name: Re-check to make the device bus is in an iommu group
                  become: true
                  shell:
                    cmd: find /sys/kernel/iommu_groups/ -type l | grep {{ device_bus_number }}
                  register: bus_iommu_groups_find

                - name: Result from `find /sys/kernel/iommu_groups/ -type l`
                  debug: msg="{{ bus_iommu_groups_find.stdout }}"

                - name: Updating Coral facts
                  set_fact:
                    # Remove the `/sys/kernel/iommu_groups/` text and split by `/` and store the first value which is the iommu group ID number
                    bus_iommu_group_id: "{{ (bus_iommu_groups_find.stdout|replace('/sys/kernel/iommu_groups/','')).split('/')[0] }}"

                - name: Storing group ID number to facts
                  debug: msg="{{ bus_iommu_group_id }}"

                - name: Now re-checking again to make sure no other devices are using iommu group ID `{{ bus_iommu_group_id }}`.
                  become: true
                  shell:
                    cmd: find /sys/kernel/iommu_groups/ -type l | grep /{{ bus_iommu_group_id }}/
                  register: recheck_bus_iommu_groups_folder_find

                - name: Result from `find /sys/kernel/iommu_groups/ -type l | grep /{{ bus_iommu_group_id }}/`
                  debug: msg="{{ recheck_bus_iommu_groups_folder_find.stdout }}"

                # If still more than one line
                - when: recheck_bus_iommu_groups_folder_find.stdout_lines | length > 1
                  block:
                    - name: Fail, cannot continue
                      ansible.builtin.fail:
                        msg: "Updating group to force break-up of iommu devices has not been successful. Devices are still grouped."


    # Blacklist the device on the Promox host only if the Device ID is found, and the Device Bus Number and iommu Group ID is defined from the previous block
    - name: Check if the blacklist file exists
      become: true
      stat:
        path: /etc/modprobe.d/blacklist-apex.conf
      register: blacklist_apex_file
    
    - name: Blacklist the device on the Promox host
      when: device_id and device_bus_number and bus_iommu_group_id is defined and blacklist_apex_file.stat.exists == False
      block:
        - name: Create a new blacklist file
          become: true
          ansible.builtin.copy:
            content: |
                    blacklist gasket
                    blacklist apex
                    options vfio-pci ids={{ device_id }}
            dest: /etc/modprobe.d/blacklist-apex.conf

        - name: Update initramfs
          become: true
          ansible.builtin.command:
            cmd: >
              update-initramfs -u -k all

        - name: Rebooting machine
          become: true
          ansible.builtin.reboot:
            reboot_timeout: 300
            msg: "Rebooting machine in 5 seconds"
            # reboot_timeout: 3600

    # Verify that the apex and gasket modules did not load (they should be blacklisted, therefore should not appear)
    - name: Verify that the apex and gasket modules did not load
      become: true
      shell:
        cmd: lsmod | grep apex
      register: apex_found
      ignore_errors: true

    # If this does appear, it has not been blacklisted correctly
    - when: apex_found.stdout | length > 0
      block:
        - name: Fail, cannot continue
          ansible.builtin.fail:
            msg: "Apex and gasket modules has been found when running `lsmod | grep apex`. The device may not have been blacklisted correctly. Now failing."
        

    # Add the device to the VM if the device is blacklists, and the VM ID and Device Bus Number is defined
    - name: Add the device to the VM  ID `{{ proxmox_virtual_machine_passthrough }}`
      when: apex_found.stdout | length == 0 and proxmox_virtual_machine_passthrough is defined and device_bus_number is defined
      become: true
      shell:
        cmd: qm set {{ proxmox_virtual_machine_passthrough }} -hostpci0 {{ device_bus_number }}

---
# Coral PCIe Passthrough - This will configure the Coral PCIe device in Frigate
#   
# After running this play, setup the device in Frigate's config: https://docs.frigate.video/configuration/object_detectors/
# detectors:
#   coral:
#     type: edgetpu
#     device: pci
#   

- name: Coral PCIe passthrough Question
  ansible.builtin.pause:
    prompt: "Do you wish to setup Coral Passthrough from the host? (yes/no)"
  register: setup_coral_passthrough

- name: Setting up Coral PCIe passthrough (Yes)
  when: setup_coral_passthrough.user_input | lower == 'yes'
  block:
    - name: Set Coral device facts
      set_fact:
        device_id: "1ac1:089a"                    # Coral PCIe device ID

    - name: Check the Coral PCI device is found on the system
      become: true
      shell:
        cmd: lspci -nnk | grep {{ device_id}}
      register: coral_found
      ignore_errors: true

    # If device ID not found
    - when: device_id not in coral_found.stdout
      block:
        - name: "Note"
          debug: msg="Coral Device ID `{{ device_id }}` not found in `lspci -nnk`."

    # If device ID is found then continue
    - when: device_id in coral_found.stdout
      block:
        - name: "Note"
          debug: msg="Coral Device ID `{{ device_id }}` found in `lspci -nnk`."

        - name: "Result from lspci -nnk | grep {{ device_id}}"
          debug: msg="{{ coral_found.stdout }}"

        - name: Check for Apex drivers installed on the system
          become: true
          shell:
            cmd: lsmod | grep apex
          register: apex_drivers
          ignore_errors: true

        - name: "Result from lsmod | grep apex"
          debug: msg="{{ apex_drivers.stdout }}"

        # If not drivers found, continue to install PCIe drivers
        - when: apex_drivers.stdout | length == 0
          block:
            - name: Add repo
              become: true
              shell: echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list
            
            - stat: path=/usr/share/keyrings/cloud.google.com.apt-key.gpg
              register: repo_key

            - name: Install repo key
              when: not repo_key.stat.exists
              become: true
              shell: curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg -o /usr/share/keyrings/cloud.google.com.apt-key.gpg --dearmor
              
            - name: Update the cache. Only if the last one is more than 3600 seconds ago (apt-get update)
              ansible.builtin.apt:
                update_cache: yes
                cache_valid_time: 3600

            - name: Install required packages for Build, PCIe driver and Edge TPU runtime
              become: true
              apt:
                name:
                  - devscripts
                  - debhelper
                  - dh-dkms
                  - git
                state: present

            - name: Clone Git repo for Gasket driver
              become: true
              ansible.builtin.git:
                repo: "https://github.com/google/gasket-driver.git"
                dest: "gasket-driver"
                force: yes

            - name: "Build and install Gasket driver"
              become: true
              ansible.builtin.shell: |
                cd gasket-driver/ \
                  && debuild -us -uc -tc -b \
                  && cd ../ \
                  && dpkg -i gasket-dkms_1.0-18_all.deb

            - name: Install additional required packages
              become: true
              apt:
                name:
                  - libedgetpu1-std
                state: present

            - name: "Add the `apex` group"
              become: true
              ansible.builtin.shell: |
                sh -c "echo 'SUBSYSTEM==\"apex\", MODE=\"0660\", GROUP=\"apex\"' >> /etc/udev/rules.d/65-apex.rules" \
                  && groupadd apex \
                  && adduser $USER apex
              ignore_errors: true

            - debug:
                msg: "NOTE: Now rebooting the machine as PCIe drivers have been installed"

            - name: Rebooting machine
              become: true
              changed_when: true
              ansible.builtin.reboot:
                reboot_timeout: 300
                msg: "Rebooting machine in 5 seconds"
                # reboot_timeout: 3600

            - name: Re-check for Apex drivers installed on the system
              become: true
              shell:
                cmd: lsmod | grep apex
              register: recheck_apex_drivers
              ignore_errors: true

            - name: "Result from lsmod | grep apex"
              debug: msg="{{ recheck_apex_drivers.stdout }}"


        # If drivers are found from the original check, or the re-check
        - when: apex_drivers.stdout | length > 0 or recheck_apex_drivers.stdout | length > 0
          block:
            - name: Check for Apex devices
              become: true
              shell:
                cmd: ls -al /dev/apex_0
              register: apex_devices

            - name: "Result from ls -al /dev/apex_0"
              debug: msg="{{ apex_devices.stdout }}"

detectors:
  coral:
    type: edgetpu
    device: pci

devices:
    - /dev/apex_0:/dev/apex_0 # Passes a PCIe Coral, follow driver instructions here https://coral.ai/docs/m2/get-started/#2a-on-linux