---
- name: Validate QoS settings
  ansible.builtin.assert:
    that:
      - item.name is defined
    fail_msg: "QoS entry missing 'name': {{ item }}"
  loop: "{{ slurm_qos_settings }}"
  when: slurm_qos_assert_vars | bool

- name: Validate account settings
  ansible.builtin.assert:
    that:
      - item.name is defined
    fail_msg: "Account entry missing 'name': {{ item }}"
  loop: "{{ slurm_qos_accounts }}"
  when: slurm_qos_assert_vars | bool

- name: Validate user settings
  ansible.builtin.assert:
    that:
      - item.username is defined
    fail_msg: "User entry missing 'username': {{ item }}"
  loop: "{{ slurm_qos_users }}"
  when: slurm_qos_assert_vars | bool

- name: Ensure Slurm QoS configured
  when: slurm_qos_settings | default([]) | length > 0
  block:
    - name: Check existing Slurm QoS
      ansible.builtin.command: "sacctmgr -Pn show qos format=name"
      register: qos_list
      changed_when: false
      failed_when: false

    - name: Add missing Slurm QoS
      ansible.builtin.command: "sacctmgr -i add qos {{ item.name }}"
      loop: "{{ slurm_qos_settings }}"
      loop_control:
        label: "{{ item.name }}"
      when: qos_list.stdout is not search(item.name)
      register: qos_add
      changed_when: qos_add.rc == 0

    - name: Configure Slurm QoS
      ansible.builtin.command: >
        sacctmgr -i modify qos where name={{ item.name }} set
        {% if item.max_tres_per_user is defined and item.max_tres_per_user is not none %}MaxTRESPU={{ item.max_tres_per_user }}{% endif %}
        {% if item.max_submit_jobs is defined and item.max_submit_jobs is not none %}MaxSubmitJobs={{ item.max_submit_jobs }}{% endif %}
        {% if item.max_jobs_per_user is defined and item.max_jobs_per_user is not none %}MaxJobsPerUser={{ item.max_jobs_per_user }}{% endif %}
        {% if item.grace_time is defined and item.grace_time is not none %}GraceTime={{ item.grace_time }}{% endif %}
        {% if item.preempt is defined and item.preempt is not none %}Preempt={{ item.preempt }}{% endif %}
      loop: "{{ slurm_qos_settings }}"
      loop_control:
        label: "{{ item.name }}"
      register: qos_modify
      changed_when: qos_modify.rc == 0
      failed_when: qos_modify.rc != 0 and "Nothing modified" not in qos_modify.stdout

- name: Ensure Slurm accounts exist
  when: slurm_qos_accounts | default([]) | length > 0
  block:
    - name: Get Slurm cluster list
      ansible.builtin.command: "sacctmgr -Pn show cluster format=cluster"
      register: cluster_list
      changed_when: false
      failed_when: false

    - name: Set cluster fact if only one cluster is found
      ansible.builtin.set_fact:
        slurm_cluster_name: "{{ cluster_list.stdout_lines[0] }}"
      when: cluster_list.stdout_lines | length == 1

    - name: Fail if multiple clusters detected without explicit config
      ansible.builtin.fail:
        msg: "Multiple clusters found. Set 'slurm_cluster_name' explicitly."
      when: cluster_list.stdout_lines | length > 1 and slurm_cluster_name is not defined

    - name: Check if Slurm account exists
      ansible.builtin.command: "sacctmgr -Pn list account name={{ item.name }} format=account"
      register: account_check
      changed_when: false
      failed_when: false
      loop: "{{ slurm_qos_accounts }}"
      loop_control:
        label: "{{ item.name }}"

    - name: Add Slurm account
      ansible.builtin.command: >
        sacctmgr -i add account name={{ item.item.name }}
        cluster={{ slurm_cluster_name }}
        {% if item.item.organization is defined %}Organization={{ item.item.organization }}{% endif %}
        {% if item.item.description is defined %}Description="{{ item.item.description }}"{% endif %}
        {% if item.item.parent is defined %}Parent={{ item.item.parent }}{% endif %}
      loop: "{{ account_check.results }}"
      loop_control:
        label: "{{ item.item.name }}"
      when: item.stdout is not search(item.item.name)
      changed_when: true

- name: Ensure Slurm users and partitions configured
  when: slurm_qos_users | default([]) | length > 0
  block:
    - name: Check existing Slurm users
      ansible.builtin.command: "sacctmgr -Pn list users format=user"
      register: user_list
      changed_when: false
      failed_when: false

    - name: Add new Slurm users if missing with default preempt partition
      ansible.builtin.command: >
        sacctmgr -i add user name={{ item.username }}
        cluster={{ slurm_cluster_name }}
        account={{ item.sponsor | default('orcd') }}
        partition=preempt
        QOS=preempt_qos
        DefaultQOS={{ item.qos_default | default('preempt_qos') }}
        {% if item.comment is defined %}Comment="{{ item.comment }}"{% endif %}
      loop: "{{ slurm_qos_users }}"
      loop_control:
        label: "{{ item.username }}"
      when: user_list.stdout is not search(item.username)
      register: user_add
      changed_when: user_add.rc == 0

    - name: Configure Slurm user attributes
      ansible.builtin.command: >
        sacctmgr -i modify user where name={{ item.username }}
        cluster={{ slurm_cluster_name }}
        account={{ item.sponsor | default('orcd') }} set
        {% if item.comment is defined %}Comment="{{ item.comment }}"{% endif %}
      loop: "{{ slurm_qos_users }}"
      loop_control:
        label: "{{ item.username }}"
      register: user_modify
      changed_when: user_modify.rc == 0
      failed_when: user_modify.rc != 0 and "Nothing modified" not in user_modify.stdout

    - name: Check existing partition associations for users
      ansible.builtin.command: >
        sacctmgr -Pn show assoc where user={{ item.username }}
        cluster={{ slurm_cluster_name }}
        account={{ item.sponsor | default('orcd') }}
        format=partition
      loop: "{{ slurm_qos_users }}"
      loop_control:
        label: "{{ item.username }}"
      register: assoc_check
      changed_when: false
      failed_when: false

    - name: Add missing partition associations for users
      ansible.builtin.command: >
        sacctmgr -i add user name={{ item.0.username }}
        cluster={{ slurm_cluster_name }}
        account={{ item.0.sponsor | default('orcd') }}
        partition={{ item.1.name }}
        {% if item.1.qos is defined %}
          QOS={{ item.1.qos }}
          DefaultQOS={{ item.1.qos }}
        {% else %}
          QOS={{ slurm_partition_qos_defaults[item.1.name] | default(item.0.qos_default | default('preempt_qos')) }}
          DefaultQOS={{ slurm_partition_qos_defaults[item.1.name] | default(item.0.qos_default | default('preempt_qos')) }}
        {% endif %}
      loop: "{{ slurm_qos_users | subelements('partitions', skip_missing=True) }}"
      loop_control:
        label: "{{ item.0.username }} - {{ item.1.name }}"
      when: >
        assoc_check.results | selectattr('item.username', 'equalto', item.0.username) | map(attribute='stdout') | join('') is not search(item.1.name)
      register: assoc_add
      changed_when: assoc_add.rc == 0

    - name: Ensure partition associations for users (defined partitions)
      ansible.builtin.command: >
        sacctmgr -i modify user where
        name={{ item.0.username }}
        cluster={{ slurm_cluster_name }}
        account={{ item.sponsor | default('orcd') }}
        partition={{ item.1.name }} set
        {% if item.1.qos is defined %}
          QOS={{ item.1.qos }}
          DefaultQOS={{ item.1.qos }}
        {% else %}
          QOS={{ slurm_partition_qos_defaults[item.1.name] | default(item.0.qos_default | default('preempt_qos')) }}
          DefaultQOS={{ slurm_partition_qos_defaults[item.1.name] | default(item.0.qos_default | default('preempt_qos')) }}
        {% endif %}
      loop: "{{ slurm_qos_users | subelements('partitions', skip_missing=True) }}"
      loop_control:
        label: "{{ item.0.username }} - {{ item.1.name }}"
      register: assoc_modify
      changed_when: assoc_modify.rc == 0
      failed_when: assoc_modify.rc != 0 and "Nothing modified" not in assoc_modify.stdout

    - name: Ensure default preempt partition association
      ansible.builtin.command: >
        sacctmgr -i modify user where
        name={{ item.username }}
        cluster={{ slurm_cluster_name }}
        account={{ item.sponsor | default('orcd') }}
        partition=preempt set
        QOS=preempt_qos
        DefaultQOS={{ item.qos_default | default('preempt_qos') }}
      loop: "{{ slurm_qos_users }}"
      loop_control:
        label: "{{ item.username }}"
      when: item.partitions is not defined or item.partitions | length == 0
      register: default_assoc_modify
      changed_when: default_assoc_modify.rc == 0
      failed_when: default_assoc_modify.rc != 0 and "Nothing modified" not in default_assoc_modify.stdout