---
# GPU information gathering tasks

- name: Detect GPU devices using lspci
  shell: lspci | grep -iE "VGA|3D|Display" || echo "No GPU detected"
  register: system_info_gpu_lspci_raw
  changed_when: false
  tags: [gather, gpu]

- name: Gather detailed GPU information
  shell: lspci -v -s $(lspci | grep -iE "VGA|3D" | cut -d' ' -f1) 2>/dev/null || echo "No detailed GPU info available"
  register: system_info_gpu_detailed_raw
  changed_when: false
  tags: [gather, gpu]

- name: Check for NVIDIA GPU
  shell: lspci | grep -i nvidia
  register: system_info_nvidia_check
  changed_when: false
  failed_when: false
  tags: [gather, gpu]

- name: Gather NVIDIA GPU details (if available)
  shell: nvidia-smi --query-gpu=name,driver_version,memory.total,compute_cap --format=csv,noheader 2>/dev/null || echo "nvidia-smi not available"
  register: system_info_nvidia_smi_raw
  changed_when: false
  failed_when: false
  when: system_info_nvidia_check.rc == 0
  tags: [gather, gpu]

- name: Check for AMD GPU
  shell: lspci | grep -iE "AMD|ATI"
  register: system_info_amd_check
  changed_when: false
  failed_when: false
  tags: [gather, gpu]

- name: Gather AMD GPU details (if available)
  shell: |
    if command -v rocm-smi &> /dev/null; then
      rocm-smi --showproductname --showdriverversion
    else
      echo "rocm-smi not available"
    fi
  register: system_info_amd_rocm_raw
  changed_when: false
  failed_when: false
  when: system_info_amd_check.rc == 0
  tags: [gather, gpu]

- name: Check for Intel GPU
  shell: lspci | grep -i "intel.*graphics"
  register: system_info_intel_gpu_check
  changed_when: false
  failed_when: false
  tags: [gather, gpu]

- name: Parse GPU information
  set_fact:
    system_info_gpu_detected: "{{ system_info_gpu_lspci_raw.stdout != 'No GPU detected' }}"
    system_info_gpu_list: "{{ system_info_gpu_lspci_raw.stdout_lines | default([]) }}"
  tags: [gather, gpu]

- name: Build GPU details structure
  set_fact:
    system_info_gpu:
      detected: "{{ system_info_gpu_detected }}"
      devices: "{{ system_info_gpu_list }}"
      nvidia:
        present: "{{ system_info_nvidia_check.rc == 0 }}"
        details: "{{ system_info_nvidia_smi_raw.stdout_lines | default([]) if system_info_nvidia_check.rc == 0 else [] }}"
      amd:
        present: "{{ system_info_amd_check.rc == 0 }}"
        details: "{{ system_info_amd_rocm_raw.stdout_lines | default([]) if system_info_amd_check.rc == 0 else [] }}"
      intel:
        present: "{{ system_info_intel_gpu_check.rc == 0 }}"
      detailed_info: "{{ system_info_gpu_detailed_raw.stdout_lines | default([]) }}"
  tags: [gather, gpu]

- name: Check for GPU passthrough support (IOMMU)
  shell: |
    if dmesg | grep -iE "IOMMU|AMD-Vi|Intel VT-d" | grep -i enabled; then
      echo "IOMMU enabled"
    else
      echo "IOMMU disabled or not available"
    fi
  register: system_info_iommu_status_raw
  changed_when: false
  become: true
  failed_when: false
  tags: [gather, gpu]

- name: Add IOMMU status to GPU info
  set_fact:
    system_info_gpu: "{{ system_info_gpu | combine({'iommu_status': system_info_iommu_status_raw.stdout | default('Unknown')}) }}"
  tags: [gather, gpu]