- name: Setup AI tools with Podman (Ollama)
  hosts: ai

  tasks:
    - name: Install Podman and its dependencies
      become: true
      ansible.builtin.dnf:
        name:
          - podman
          - nvidia-container-toolkit
        state: present

    - name: Create group and user for AI workloads
      become: true
      ansible.builtin.group:
        name: "{{ ai_group }}"
        system: true

    - name: Create user and group for AI containers
      become: true
      ansible.builtin.user:
        name: "{{ ai_user }}"
        groups: "{{ ai_group }}"
        shell: /bin/bash
        home: "{{ ai_home }}"
        group: "{{ ai_group }}"

    - name: Persistently set SELinux context for home directories on AI home
      become: true
      community.general.sefcontext:
        target: "{{ ai_home }}(/.*)?"
        setype: user_home_dir_t
        state: present

    - name: Apply SELinux context changes to AI home directories
      become: true
      ansible.builtin.command: restorecon -Rv {{ ai_home }}
      changed_when: false

    - name: Disable password login for AI user
      ansible.builtin.command: passwd -d "{{ ai_user }}"
      become: true
      changed_when: false

    - name: Create data folders for AI tools
      become: true
      ansible.builtin.file:
        path: "{{ ai_home }}"
        state: directory
        owner: "{{ ai_user }}"
        group: "{{ ai_group }}"
        mode: "0750"

    - name: Create build folders for AI tools
      become: true
      ansible.builtin.file:
        path: "{{ ai_home }}/build"
        state: directory
        owner: "{{ ai_user }}"
        group: "{{ ai_group }}"
        mode: "0750"

    - name: Create data folders for AI tools
      become: true
      ansible.builtin.file:
        path: "{{ ai_home }}/{{ item.name }}"
        state: directory
        owner: "{{ ai_user }}"
        group: "{{ ai_group }}"
        mode: "0750"
      loop: "{{ ai_workloads }}"

    - name: Create systemd service file for nvidia-cdi-generator
      become: true
      ansible.builtin.copy:
        src: services/nvidia-cdi-generator.service
        dest: /etc/systemd/system/nvidia-cdi-generator.service
        owner: root
        group: root
        mode: "0644"

    - name: Reload systemd daemon
      become: true
      ansible.builtin.systemd:
        daemon_reload: true
        name: nvidia-cdi-generator

    - name: Enable nvidia-cdi-generator service
      become: true
      ansible.builtin.systemd:
        name: nvidia-cdi-generator
        enabled: true
        state: started

    - name: Enable sudo access to ai user (temporary)
      become: true
      ansible.builtin.lineinfile:
        path: /etc/sudoers.d/ansible_ai_tmp
        create: true
        mode: '0440'
        line: "{{ ansible_user_id }} ALL=(ai) NOPASSWD: ALL"
        validate: 'visudo -cf %s'

    - name: Enable lingering for ai user
      become: true
      ansible.builtin.command: loginctl enable-linger {{ ai_user }}
      register: linger_status
      changed_when: "'created' in linger_status.stdout or (linger_status.rc == 0 and not ('linger file already exists' in linger_status.stderr or 'linger file does not exist' in linger_status.stderr))" # Heuristic for change
      failed_when: linger_status.rc != 0 and not ('linger file already exists' in linger_status.stderr)

    - name: Get current user's UID
      ansible.builtin.command: id -u {{ ai_user }}
      changed_when: false
      register: uid

    - name: Get current user's GID
      ansible.builtin.command: id -g {{ ai_user }}
      changed_when: false
      register: gid

    - name: SHowing UID and GID of ai user
      ansible.builtin.debug:
        msg: "UID: {{ uid.stdout }}, GID: {{ gid.stdout }}"

    - name: Run tasks as ai user
      become: true
      become_user: "{{ ai_user }}"
      block:
        - name: Create a custom Podman network
          containers.podman.podman_network:
            name: "{{ ai_network_name }}"
            state: present

        - name: Build containers images
          containers.podman.podman_image:
            name: "ai/{{ item.name }}"
            path: "{{ ai_home }}/build"
            build:
              cache: false
              force_rm: true
              container_file: "{{ item.build.dockerfile }}"
              extra_args: "--build-arg UID={{ uid.stdout }} --build-arg GID={{ gid.stdout }} --security-opt=label=disable {{ item.build.args | default('') }}"
            tag: "latest"
            state: present
          loop: "{{ ai_workloads }}"
          when: item.build is defined

        - name: Pull the requested images
          containers.podman.podman_image:
            name: "{{ item.repo.image }}"
            tag: "{{ item.repo.tag | default('latest') }}"
            state: present
          loop: "{{ ai_workloads }}"
          when: item.repo is defined

        - name: Create a container instance for each workload
          containers.podman.podman_container:
            name: "{{ item.name }}"
            image: "{{ item.repo.image | default('ai/' + item.name) }}"
            state: started
            device: "{{ item.device | default(omit) }}"
            env: "{{ item.env | default(omit) }}"
            network: "{{ ai_network_name }}"
            ports: "{{ item.ports | default(omit) }}"
            volumes: "{{ item.volumes | default(omit) }}"
            cmd_args:
              - "--userns=keep-id"
              - "--security-opt=label=disable"
          loop: "{{ ai_workloads }}"

    - name: Create systemd service file for containers
      become: true
      ansible.builtin.template:
        src: podman.service.j2
        dest: "/etc/systemd/system/podman-{{ item.name }}.service"
        owner: root
        group: root
        mode: "0644"
      loop: "{{ ai_workloads }}"

    - name: Reload systemd daemon
      become: true
      ansible.builtin.systemd:
        daemon_reload: true

    - name: Enable containers at boot
      become: true
      ansible.builtin.systemd:
        name: "podman-{{ item.name }}"
        enabled: true
        state: started
      loop: "{{ ai_workloads }}"

    - name: Remove temporary sudo access
      become: true
      ansible.builtin.file:
        path: /etc/sudoers.d/ansible_ai_tmp
        state: absent