- name: Setup ollama
  hosts: ollama
  vars:
    podman_user: "{{ ollama_user | default('ollama') }}"
    podman_group: "{{ ollama_group | default('ollama') }}"
    podman_extra_groups: "users,video"
    podman_user_home: "{{ ollama_home | default('/home/ollama') }}"
    podman_user_folders:
      - data
    __ollama_models__: "{{ ollama_models | default([]) }}"
    __ollama_model_files__: "{{ ollama_model_files | default([]) }}"
    pods:
      - name: ollama-rocm
        build:
          dockerfile: "{{ lookup('file', 'dockerfiles/Dockerfile.ollama.rocm') }}"
        env:
          OLLAMA_NUM_THREADS: 10
          OLLAMA_NUM_PARALLEL: 2
          OLLAMA_FLASH_ATTENTION: 1
          OLLAMA_KV_CACHE_TYPE: q8_0
          OLLAMA_KV_CACHE: "true"
          HSA_OVERRIDE_GFX_VERSION: "11.0.0"
        ports:
          - "0.0.0.0:{{ ollama_port | default(11434) }}:11434/tcp"
        volumes:
          - "{{ ollama_home }}/data:/home/worker/.ollama:Z"
        device:
          - "/dev/kfd:/dev/kfd:rw"
          - "/dev/dri:/dev/dri:rw"

  tasks:
    - name: Setup Pods
      ansible.builtin.include_tasks:
        file: tasks/podman.yml

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

    - name: Create staging directory for model files
      become: true
      ansible.builtin.file:
        path: "{{ ollama_home }}/data/.modelfiles"
        state: directory
        owner: "{{ podman_user }}"
        group: "{{ podman_group }}"
        mode: "0755"
      when: __ollama_model_files__ | length > 0

    - name: Copy model files to remote host staging directory
      become: true
      ansible.builtin.copy:
        src: "{{ item.file }}"
        dest: "{{ ollama_home }}/data/.modelfiles/{{ item.file | basename }}"
        owner: "{{ podman_user }}"
        group: "{{ podman_group }}"
        mode: "0644"
      loop: "{{ __ollama_model_files__ }}"
      when: __ollama_model_files__ | length > 0

    - name: Check existing ollama models
      become: true
      become_user: "{{ podman_user }}"
      containers.podman.podman_container_exec:
        name: ollama-rocm
        command: "ollama ls"
      register: ollama_existing_models
      changed_when: false

    - name: Set fact for existing model names
      ansible.builtin.set_fact:
        __ollama_existing_model_names__: >-
          {{ ollama_existing_models.stdout_lines | default([])
             | reject('match', '^\\s*$|^(NAME|\\-+)')
             | map('regex_replace', '^\\s*(\\S+)\\s.*', '\\1')
             | list }}
      when: ollama_existing_models.stdout is defined and ollama_existing_models.stdout | length > 0
      changed_when: false

    - name: Pull ollama models that do not exist
      become: true
      become_user: "{{ podman_user }}"
      containers.podman.podman_container_exec:
        name: ollama-rocm
        command: "ollama pull {{ item }}"
      loop: "{{ __ollama_models__ }}"
      when:
        - __ollama_models__ | length > 0
        - item not in (__ollama_existing_model_names__ | default([]))

    - name: Create ollama models that do not exist
      become: true
      become_user: "{{ podman_user }}"
      containers.podman.podman_container_exec:
        name: ollama-rocm
        command: "ollama create {{ item.name }} -f /home/worker/.ollama/.modelfiles/{{ item.file | basename }}"
      loop: "{{ __ollama_model_files__ }}"
      when:
        - __ollama_model_files__ | length > 0
        - item.name not in (__ollama_existing_model_names__ | default([]))

    - name: Remove temporary sudo access
      become: true
      ansible.builtin.file:
        path: "/etc/sudoers.d/ansible_{{ podman_user }}_tmp"
        state: absent