---
- name: Get details about rootfs disks
  community.general.parted:
    device: "{{ item }}"
    unit: MiB
  register: base_root_disks_info
  loop: "{{ base_root_disks | list }}"

- name: Calculate maximum usable disk space
  ansible.builtin.set_fact:
    base_root_usable_mib: "{{ (base_root_disks_info.results | map(attribute='disk.size') | min | int) - 1 }}"

- debug: var=base_root_usable_mib

- name: Calculate disk utilization percentage
  ansible.builtin.set_fact:
    base_root_usable_mib: "{{ base_root_usable_mib|int -  ((base_root_usable_mib|float) * (0.01 * (base_root_free_percent|float))) | round(method='floor') | int }}"

- name: Calculate zfs volume size
  set_fact:
    base_root_zpool_mib: "{{ base_root_usable_mib|int - fixed_size_partitions|int }}"
  vars:
    fixed_size_partitions: "{{ base_root_swap_mib|int + base_root_efi_mib|int }}"
#
- name: Calculate partition layouts
  set_fact:
    partition_ranges: >-
      {{
      (partition_ranges|d([])) + [{
        'begin': partitions[:(loop_index|int)] | map(attribute='size_mib') | map('int') | sum(),
        'end': partitions[:(loop_index|int+1)] | map(attribute='size_mib') | map('int') | sum(),
        }]
      }}
  with_sequence: start=0 end="{{ partitions | length - 1}}"
  vars:
    loop_index: "{{ item }}"
    partitions:
      - name: efi
        size_mib: "{{ base_root_efi_mib }}"
      - name: root
        size_mib: "{{ base_root_zpool_mib }}"
      - name: swap
        size_mib: "{{ base_root_swap_mib }}"

- debug:
    var: partition_ranges

- name: Initialize partition facts
  set_fact:
    base_efi_partitions: []
    base_root_partitions: []
    base_swap_partitions: []

#### efi/boot
- name: Create EFI partition
  become: true
  vars:
    part_index: 0
  parted:
    label: gpt
    unit: MiB
    name: EFI Boot
    device: "{{ item.disk.dev }}"
    number: "{{ part_index + 1 }}"
    flags: [ boot, esp ]
    part_start: "{{ [partition_ranges[part_index].begin, 1]|max }}MiB"
    part_end: "{{ partition_ranges[part_index].end }}MiB"
    state: present
    fs_type: fat32
  loop: "{{ base_root_disks_info.results }}"
  register: parted_create

- name: Collect EFI partition devices
  command: "lsblk -r --noheadings -o PATH {{ item }}"
  register: lsblk
  loop: "{{ parted_create.results | map(attribute='disk.dev') }}"

- name: Store EFI partition devices
  set_fact:
    base_efi_partitions: >-
      {{ base_efi_partitions|d([]) + [
          item[part_index+1]
        ]
      }}
  loop: "{{ lsblk.results | map(attribute='stdout_lines') | sort }}"
  vars:
    part_index: 0

- debug: var=base_efi_partitions

#### zroot
- name: Create root zvol partition
  become: true
  vars:
    part_index: 1
  parted:
    label: gpt
    unit: MiB
    name: ArchLinux ZFS Root
    device: "{{ item.disk.dev }}"
    number: "{{ part_index + 1 }}"
    part_start: "{{ partition_ranges[part_index].begin }}MiB"
    part_end: "{{ partition_ranges[part_index].end }}MiB"
    state: present
  loop: "{{ parted_create.results }}"
  register: parted_create

- name: Collect root zvol partition devices
  command: "lsblk -r --noheadings -o PATH {{ item }}"
  register: lsblk
  loop: "{{ parted_create.results | map(attribute='disk.dev') }}"

- name: Store root zvol partition devices
  set_fact:
    base_root_partitions: >-
      {{ base_root_partitions|d([]) + [
          item[part_index+1]
        ]
      }}
  loop: "{{ lsblk.results | map(attribute='stdout_lines') | sort }}"
  vars:
    part_index: 1

- debug: var=base_root_partitions

#### swap
- name: Create swap partition
  become: true
  vars:
    part_index: 2
  parted:
    label: gpt
    unit: MiB
    device: "{{ item.disk.dev }}"
    number: "{{ part_index + 1 }}"
    part_start: "{{ partition_ranges[part_index].begin }}MiB"
    part_end: "{{ partition_ranges[part_index].end }}MiB"
    state: present
    fs_type: linux-swap
  loop: "{{ parted_create.results }}"
  register: parted_create
  when: base_root_swap_mib is defined and base_root_swap_mib|int > 0

- name: Collect swap partition devices
  command: "lsblk -r --noheadings -o PATH {{ item }}"
  register: lsblk
  loop: "{{ parted_create.results | map(attribute='disk.dev') }}"

- name: Store swap partition devices
  set_fact:
    base_swap_partitions: >-
      {{ base_swap_partitions|d([]) + [
          item[part_index+1]
        ]
      }}
  loop: "{{ lsblk.results | map(attribute='stdout_lines') | sort }}"
  vars:
    part_index: 2

- name: Analyze resulting partition layouts
  parted:
    unit: MiB
    device: "{{ item }}"
  register: base_root_disks_info
  loop: "{{ base_root_disks | list }}"

- name: Collect disk device identifiers
  shell: "for x in /dev/disk/by-id/*; do echo $x $(realpath $x); done"
  register: disk_realpaths

- name: Collect disk device identifiers into a base_partitions_by_id dictionary
  set_fact:
    base_partitions_by_id: >-
      {{
        dict(
          disk_realpaths.stdout_lines
          | map('split', ' ')
          | map('list')
        )
      }}

- debug: var=base_efi_partitions
- debug: var=base_root_partitions
- debug: var=base_swap_partitions