- name: Require pinned single-user TransformerLab install mode
  assert:
    that:
      - courseware_transformerlab_install_mode == "single-user-pinned"
    fail_msg: "This courseware only supports the pinned single-user TransformerLab install mode."

- name: Check pinned TransformerLab source tree
  stat:
    path: "{{ courseware_transformerlab_home }}/src/install.sh"
  register: courseware_transformerlab_src_install

- name: Check pinned TransformerLab version file
  stat:
    path: "{{ courseware_transformerlab_home }}/src/LATEST_VERSION"
  register: courseware_transformerlab_src_version

- name: Check pinned TransformerLab web app
  stat:
    path: "{{ courseware_transformerlab_home }}/webapp/index.html"
  register: courseware_transformerlab_web_index

- name: Check pinned TransformerLab conda runtime
  stat:
    path: "{{ courseware_transformerlab_home }}/miniforge3/bin/conda"
  register: courseware_transformerlab_conda

- name: Check pinned TransformerLab Python runtime
  stat:
    path: "{{ courseware_transformerlab_home }}/envs/transformerlab/bin/python"
  register: courseware_transformerlab_env_python

- name: Check pinned TransformerLab install marker
  stat:
    path: "{{ courseware_transformerlab_home }}/.install_complete"
  register: courseware_transformerlab_install_marker

- name: Read pinned TransformerLab version file
  slurp:
    src: "{{ courseware_transformerlab_home }}/src/LATEST_VERSION"
  register: courseware_transformerlab_version_contents
  when: courseware_transformerlab_src_version.stat.exists

- name: Determine pinned TransformerLab drift state
  set_fact:
    courseware_transformerlab_installed_version: "{{ (courseware_transformerlab_version_contents.content | b64decode | trim) if courseware_transformerlab_src_version.stat.exists else '' }}"
    courseware_transformerlab_refresh_required: >-
      {{
        (not courseware_transformerlab_src_install.stat.exists) or
        (not courseware_transformerlab_src_version.stat.exists) or
        (not courseware_transformerlab_web_index.stat.exists) or
        ((courseware_transformerlab_version_contents.content | b64decode | trim) != courseware_transformerlab_version)
      }}
    courseware_transformerlab_env_refresh_required: >-
      {{
        (not courseware_transformerlab_conda.stat.exists) or
        (not courseware_transformerlab_env_python.stat.exists) or
        (not courseware_transformerlab_install_marker.stat.exists)
      }}

- name: Download pinned TransformerLab source archive
  get_url:
    url: "https://github.com/transformerlab/transformerlab-app/archive/refs/tags/{{ courseware_transformerlab_version }}.tar.gz"
    dest: "{{ courseware_transformerlab_source_archive }}"
    mode: "0644"
  when: courseware_transformerlab_refresh_required

- name: Download pinned TransformerLab web archive
  get_url:
    url: "https://github.com/transformerlab/transformerlab-app/releases/download/{{ courseware_transformerlab_version }}/transformerlab_web.tar.gz"
    dest: "{{ courseware_transformerlab_web_archive }}"
    mode: "0644"
  when: courseware_transformerlab_refresh_required

- name: Refresh pinned TransformerLab release layout
  shell: |
    set -euo pipefail
    home_dir="{{ courseware_transformerlab_home }}"
    source_archive="{{ courseware_transformerlab_source_archive }}"
    web_archive="{{ courseware_transformerlab_web_archive }}"
    version="{{ courseware_transformerlab_version }}"
    version_dir="{{ courseware_transformerlab_version_dir }}"
    extract_dir="${home_dir}/transformerlab-app-${version_dir}"

    rm -rf "${extract_dir}"
    tar -xzf "${source_archive}" -C "${home_dir}"

    if [ ! -d "${extract_dir}/api" ]; then
      echo "Pinned TransformerLab source archive did not contain api/." >&2
      exit 1
    fi

    rm -rf "${home_dir}/src" "${home_dir}/lab-sdk" "${home_dir}/webapp"
    mv "${extract_dir}/api" "${home_dir}/src"

    if [ -d "${extract_dir}/lab-sdk" ]; then
      mv "${extract_dir}/lab-sdk" "${home_dir}/lab-sdk"
    fi

    printf '%s\n' "${version}" > "${home_dir}/src/LATEST_VERSION"

    mkdir -p "${home_dir}/webapp"
    tar -xzf "${web_archive}" -C "${home_dir}/webapp"
    if [ -d "${home_dir}/webapp/transformerlab_web" ]; then
      shopt -s dotglob nullglob
      mv "${home_dir}/webapp/transformerlab_web/"* "${home_dir}/webapp/"
      rmdir "${home_dir}/webapp/transformerlab_web"
    fi

    rm -rf "${extract_dir}"
    rm -f "${home_dir}/.install_complete"
  args:
    executable: /bin/bash
  when: courseware_transformerlab_refresh_required

- name: Ensure pinned TransformerLab environment file exists
  file:
    path: "{{ courseware_transformerlab_home }}/.env"
    state: touch
    mode: "0600"

- name: Ensure pinned TransformerLab JWT secrets exist
  shell: |
    set -euo pipefail
    env_file="{{ courseware_transformerlab_home }}/.env"

    if ! grep -q '^TRANSFORMERLAB_JWT_SECRET=' "$env_file"; then
      printf 'TRANSFORMERLAB_JWT_SECRET=%s\n' "$(od -An -N 64 -tx1 /dev/urandom | tr -d ' \n')" >> "$env_file"
    fi

    if ! grep -q '^TRANSFORMERLAB_REFRESH_SECRET=' "$env_file"; then
      printf 'TRANSFORMERLAB_REFRESH_SECRET=%s\n' "$(od -An -N 64 -tx1 /dev/urandom | tr -d ' \n')" >> "$env_file"
    fi
  args:
    executable: /bin/bash

- name: Remove stale TransformerLab general-uv environment
  file:
    path: "{{ courseware_transformerlab_home }}/envs/general-uv"
    state: absent

- name: Remove stale TransformerLab multiuser marker
  file:
    path: "{{ courseware_transformerlab_home }}/.multiuser_setup_complete"
    state: absent

- name: Determine Miniforge platform suffix
  set_fact:
    courseware_transformerlab_miniforge_platform: "{{ 'Linux-x86_64' if ansible_system == 'Linux' and ansible_architecture == 'x86_64' else 'Linux-aarch64' if ansible_system == 'Linux' and ansible_architecture in ['aarch64', 'arm64'] else 'MacOSX-arm64' if ansible_system == 'Darwin' and ansible_architecture == 'arm64' else 'MacOSX-x86_64' if ansible_system == 'Darwin' and ansible_architecture == 'x86_64' else 'unsupported' }}"

- name: Fail for unsupported Miniforge platform
  fail:
    msg: "Unsupported TransformerLab Miniforge platform: {{ ansible_system }} {{ ansible_architecture }}"
  when: courseware_transformerlab_miniforge_platform == "unsupported"

- name: Download Miniforge installer for pinned TransformerLab
  get_url:
    url: "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-{{ courseware_transformerlab_miniforge_platform }}.sh"
    dest: "{{ courseware_transformerlab_miniforge_installer }}"
    mode: "0755"
  when: not courseware_transformerlab_conda.stat.exists

- name: Install Miniforge for pinned TransformerLab
  shell: |
    set -euo pipefail
    rm -rf "{{ courseware_transformerlab_home }}/miniforge3"
    bash "{{ courseware_transformerlab_miniforge_installer }}" -b -p "{{ courseware_transformerlab_home }}/miniforge3"
  args:
    executable: /bin/bash
  when: not courseware_transformerlab_conda.stat.exists

- name: Create pinned TransformerLab environment
  shell: |
    set -euo pipefail
    "{{ courseware_transformerlab_home }}/miniforge3/bin/conda" create -y -k --prefix "{{ courseware_transformerlab_home }}/envs/transformerlab" python=3.11
  args:
    executable: /bin/bash
  when: not courseware_transformerlab_env_python.stat.exists

- name: Detect pinned TransformerLab accelerator profile
  shell: |
    set -euo pipefail
    if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi --query-gpu=name --format=csv,noheader,nounits >/dev/null 2>&1; then
      printf '%s\n' nvidia
    elif command -v rocminfo >/dev/null 2>&1; then
      printf '%s\n' rocm
    else
      printf '%s\n' cpu
    fi
  args:
    executable: /bin/bash
  register: courseware_transformerlab_accelerator
  changed_when: false

- name: Install pinned TransformerLab Python dependencies
  shell: |
    set -euo pipefail
    env_python="{{ courseware_transformerlab_home }}/envs/transformerlab/bin/python"
    conda_bin="{{ courseware_transformerlab_home }}/miniforge3/bin/conda"
    accelerator="{{ courseware_transformerlab_accelerator.stdout | trim }}"
    project_dir="{{ courseware_transformerlab_home }}/src"
    wheel_args=()
    extra="cpu"

    "$env_python" -m pip install --upgrade pip
    "$env_python" -m pip install --upgrade uv

    if [ "$accelerator" = "nvidia" ]; then
      "$conda_bin" install -y --prefix "{{ courseware_transformerlab_home }}/envs/transformerlab" cuda==12.8.1 --force-reinstall -c nvidia/label/cuda-12.8.1
      extra="nvidia"
    elif [ "$accelerator" = "rocm" ]; then
      wheel_args+=(--index https://download.pytorch.org/whl/rocm6.4 --index-strategy unsafe-best-match)
      extra="rocm"
    elif [ "{{ ansible_system }}" != "Darwin" ]; then
      wheel_args+=(--index https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match)
    fi

    cd "$project_dir"
    "$env_python" -m uv pip install --python "$env_python" "${wheel_args[@]}" ".[${extra}]"

    if [ -d "{{ courseware_transformerlab_home }}/lab-sdk" ]; then
      "$env_python" -m uv pip install --python "$env_python" "{{ courseware_transformerlab_home }}/lab-sdk"
    fi

    "$env_python" -c "import uvicorn"
    touch "{{ courseware_transformerlab_home }}/.install_complete"
  args:
    executable: /bin/bash
  when: courseware_transformerlab_refresh_required or courseware_transformerlab_env_refresh_required

- name: Install TransformerLab plugin repair helper
  copy:
    src: "{{ playbook_dir }}/../../scripts/repair_transformerlab_plugin_supports.py"
    dest: "{{ courseware_state_dir }}/repair_transformerlab_plugin_supports.py"
    mode: "0755"

- name: Install TransformerLab default-user helper
  copy:
    src: "{{ playbook_dir }}/../../scripts/ensure_transformerlab_user.py"
    dest: "{{ courseware_state_dir }}/ensure_transformerlab_user.py"
    mode: "0755"

- name: Repair installed Fastchat plugin supports
  command:
    argv:
      - python3
      - "{{ courseware_state_dir }}/repair_transformerlab_plugin_supports.py"
      - --transformerlab-dir
      - "{{ courseware_transformerlab_home }}"
      - --plugin
      - fastchat_server
      - --required-support
      - chat
      - --required-support
      - completion
      - --required-support
      - visualize_model
      - --required-support
      - model_layers
      - --required-support
      - rag
      - --required-support
      - tools
      - --required-support
      - template
      - --required-support
      - embeddings
      - --required-support
      - tokenize
      - --required-support
      - logprobs
      - --required-support
      - batched
  changed_when: false

- name: Ensure pinned TransformerLab default user when database already exists
  command:
    argv:
      - "{{ courseware_transformerlab_home }}/envs/transformerlab/bin/python"
      - "{{ courseware_state_dir }}/ensure_transformerlab_user.py"
      - --transformerlab-dir
      - "{{ courseware_transformerlab_home }}"
      - --email
      - "{{ courseware_transformerlab_default_user_email }}"
      - --password
      - "{{ courseware_transformerlab_default_user_password }}"
      - --first-name
      - "{{ courseware_transformerlab_default_user_first_name }}"
      - --last-name
      - "{{ courseware_transformerlab_default_user_last_name }}"
  changed_when: false
