---
# =============================================================================
# proxmox_upgrade — drain.yml
# Migrate all VMs/LXCs off a node before upgrading it
# Uses Proxmox API — runs delegate_to: localhost
# =============================================================================

# ── Get all guests on this node ───────────────────────────────────────────────
- name: Drain | Get all VMs on node {{ current_node }}
  ansible.builtin.uri:
    url: "https://{{ api_host }}:{{ api_port }}/api2/json/nodes/{{ current_node }}/qemu"
    method: GET
    headers:
      Authorization: "PVEAPIToken={{ api_token_id }}={{ api_token_secret }}"
    validate_certs: false
  register: node_vms
  delegate_to: localhost

- name: Drain | Get all LXCs on node {{ current_node }}
  ansible.builtin.uri:
    url: "https://{{ api_host }}:{{ api_port }}/api2/json/nodes/{{ current_node }}/lxc"
    method: GET
    headers:
      Authorization: "PVEAPIToken={{ api_token_id }}={{ api_token_secret }}"
    validate_certs: false
  register: node_lxcs
  delegate_to: localhost

- name: Drain | Get available target nodes
  ansible.builtin.uri:
    url: "https://{{ api_host }}:{{ api_port }}/api2/json/nodes"
    method: GET
    headers:
      Authorization: "PVEAPIToken={{ api_token_id }}={{ api_token_secret }}"
    validate_certs: false
  register: all_nodes
  delegate_to: localhost

- name: Drain | Build target node list (exclude current node)
  ansible.builtin.set_fact:
    migration_targets: >-
      {{ all_nodes.json.data
         | selectattr('status', 'equalto', 'online')
         | rejectattr('node', 'equalto', current_node)
         | map(attribute='node')
         | list }}
  delegate_to: localhost

- name: Drain | Fail if no migration targets available
  ansible.builtin.fail:
    msg: "No online nodes available to migrate guests to. Cannot drain {{ current_node }}."
  when: migration_targets | length == 0
  delegate_to: localhost

# ── Classify VMs — live migratable vs needs fallback ─────────────────────────
- name: Drain | Get VM configs to check migratability
  ansible.builtin.uri:
    url: "https://{{ api_host }}:{{ api_port }}/api2/json/nodes/{{ current_node }}/qemu/{{ item.vmid }}/config"
    method: GET
    headers:
      Authorization: "PVEAPIToken={{ api_token_id }}={{ api_token_secret }}"
    validate_certs: false
  register: vm_configs
  loop: "{{ node_vms.json.data }}"
  delegate_to: localhost

- name: Drain | Build guest migration plan
  ansible.builtin.set_fact:
    migration_plan: >-
      {%- set plan = [] -%}
      {%- for vm in node_vms.json.data -%}
        {%- set cfg = vm_configs.results[loop.index0].json.data -%}
        {%- set tags = (vm.tags | default('')) .split(',') | map('trim') | list -%}
        {%- set excluded = tags | select('in', migrate_exclude_tags) | list | length > 0 -%}
        {%- set has_passthrough = 'hostpci0' in cfg or 'usb0' in cfg -%}
        {%- set has_local_disk = shared_storage == false -%}
        {%- set has_local_cdrom = cfg.values() | select('string') | select('match', '.*local.*\\.iso.*') | list | length > 0 -%}
        {%- set needs_fallback = has_passthrough or has_local_disk or has_local_cdrom -%}
        {%- if not excluded -%}
          {%- set _ = plan.append({
            'vmid': vm.vmid,
            'name': vm.name,
            'type': 'qemu',
            'status': vm.status,
            'needs_fallback': needs_fallback,
            'fallback_reason': ('passthrough' if has_passthrough else ('local_disk' if has_local_disk else ('local_cdrom' if has_local_cdrom else '')))
          }) -%}
        {%- endif -%}
      {%- endfor -%}
      {%- for lxc in node_lxcs.json.data -%}
        {%- set tags = (lxc.tags | default('')) .split(',') | map('trim') | list -%}
        {%- set excluded = tags | select('in', migrate_exclude_tags) | list | length > 0 -%}
        {%- if not excluded -%}
          {%- set _ = plan.append({
            'vmid': lxc.vmid,
            'name': lxc.name,
            'type': 'lxc',
            'status': lxc.status,
            'needs_fallback': false,
            'fallback_reason': ''
          }) -%}
        {%- endif -%}
      {%- endfor -%}
      {{ plan }}
  delegate_to: localhost

- name: Drain | Log migration plan
  ansible.builtin.debug:
    msg: >-
      Migration plan for {{ current_node }}:
      {% for g in migration_plan %}
      - {{ g.type | upper }} {{ g.vmid }} ({{ g.name }}) [{{ g.status }}]
        {% if g.needs_fallback %} ⚠ needs fallback ({{ g.fallback_reason }}) — action: {{ live_migrate_fallback }}{% endif %}
      {% endfor %}
  delegate_to: localhost

# ── Abort if any guests need fallback and live_migrate_fallback is 'migrate' ──
- name: Drain | Warn about non-migratable guests
  ansible.builtin.debug:
    msg: >-
      WARNING — {{ item.type | upper }} {{ item.vmid }} ({{ item.name }})
      cannot be live migrated ({{ item.fallback_reason }}).
      live_migrate_fallback={{ live_migrate_fallback }} —
      {% if live_migrate_fallback == 'skip' %}
      THIS VM WILL GO DOWN DURING NODE REBOOT.
      {% elif live_migrate_fallback == 'shutdown' %}
      Will be shut down, cold migrated, and restarted.
      {% else %}
      Will attempt live migrate anyway (may fail).
      {% endif %}
  loop: "{{ migration_plan | selectattr('needs_fallback') | list }}"
  delegate_to: localhost

# ── Perform migrations ────────────────────────────────────────────────────────
- name: Drain | Migrate guests (sequential)
  when: not migration_bulk | bool
  include_tasks: migrate_guest.yml
  loop: "{{ migration_plan | rejectattr('needs_fallback') | list + migration_plan | selectattr('needs_fallback') | rejectattr('needs_fallback' if live_migrate_fallback == 'skip' else 'nonexistent') | list }}"
  loop_control:
    loop_var: guest

- name: Drain | Migrate guests (bulk — fire all at once)
  when: migration_bulk | bool
  block:
    - name: Drain | Bulk | Trigger all live migrations simultaneously
      ansible.builtin.uri:
        url: "https://{{ api_host }}:{{ api_port }}/api2/json/nodes/{{ current_node }}/{{ 'qemu' if guest.type == 'qemu' else 'lxc' }}/{{ guest.vmid }}/migrate"
        method: POST
        headers:
          Authorization: "PVEAPIToken={{ api_token_id }}={{ api_token_secret }}"
        body_format: json
        body:
          target: "{{ migration_targets | first }}"
          online: "{{ 1 if not guest.needs_fallback else 0 }}"
        validate_certs: false
      register: bulk_migration_tasks
      loop: "{{ migration_plan | rejectattr('needs_fallback') | list }}"
      loop_control:
        loop_var: guest
      delegate_to: localhost

    - name: Drain | Bulk | Wait for all migrations to complete
      ansible.builtin.uri:
        url: "https://{{ api_host }}:{{ api_port }}/api2/json/nodes/{{ current_node }}/tasks/{{ item.json.data }}/status"
        method: GET
        headers:
          Authorization: "PVEAPIToken={{ api_token_id }}={{ api_token_secret }}"
        validate_certs: false
      register: task_status
      until: task_status.json.data.status == 'stopped'
      retries: 60
      delay: 10
      loop: "{{ bulk_migration_tasks.results }}"
      delegate_to: localhost

    - name: Drain | Bulk | Check all migrations succeeded
      ansible.builtin.fail:
        msg: "Migration task failed for VMID — exitstatus: {{ item.json.data.exitstatus }}"
      loop: "{{ task_status.results }}"
      when: item.json.data.exitstatus != 'OK'
      delegate_to: localhost

    - name: Drain | Bulk | Handle fallback guests sequentially
      include_tasks: migrate_guest.yml
      loop: "{{ migration_plan | selectattr('needs_fallback') | list }}"
      loop_control:
        loop_var: guest
      when: live_migrate_fallback != 'skip'