---
- name: "STRESS-SETUP: Manage full worker lifecycle based on inventory"
  hosts: workers
  gather_facts: no
  vars:
    # Default action
    action: "status" # Available actions: start, stop, status, start-auth, stop-auth, start-download, stop-download
    graceful_shutdown_timeout_seconds: 30

  tasks:
    - name: "Ensure profile_prefixes is a flat list of all prefixes from profile_pools"
      ansible.builtin.set_fact:
        profile_prefixes: "{{ profile_pools | map(attribute='prefixes') | flatten }}"
      when: profile_pools is defined

    - name: "Start all configured generators and simulators"
      when: action == "start"
      block:
        - name: "Set combined profile prefixes string"
          ansible.builtin.set_fact:
            combined_prefixes: "{{ profile_prefixes | default([]) | join(',') }}"
          when: profile_prefixes is defined and profile_prefixes | length > 0

        - name: "Start auth generator(s)"
          when: profile_prefixes is defined and profile_prefixes | length > 0
          block:
            - name: "Start single auth generator for all profiles: {{ combined_prefixes | default('none') }}"
              ansible.builtin.command: >-
                ansible-playbook {{ playbook_dir }}/playbook-stress-auth-generator.yml
                -i {{ inventory_file }}
                --limit {{ inventory_hostname }}
                -e "start_generator=true"
                -e "profile_prefix={{ combined_prefixes }}"
                -e "dummy_batch={{ dummy_batch | default(false) }}"
                {% if auth_min_seconds is defined %}-e "auth_min_seconds={{ auth_min_seconds }}"{% endif %}
                {% if auth_max_seconds is defined %}-e "auth_max_seconds={{ auth_max_seconds }}"{% endif %}
                {% if batch_size is defined %}-e "batch_size={{ batch_size }}"{% endif %}
                {% if create_download_tasks is defined %}-e "create_download_tasks={{ create_download_tasks }}"{% endif %}
                {% if formats_to_download is defined %}-e "formats_to_download={{ formats_to_download }}"{% endif %}
              delegate_to: localhost
              changed_when: true
              when: (auth_workers_per_profile | default(0) | int == 0) and (auth_workers_total | default(0) | int > 0)

            - name: "Start parallel auth generators for each profile"
              ansible.builtin.command: >-
                ansible-playbook {{ playbook_dir }}/playbook-stress-auth-generator.yml
                -i {{ inventory_file }}
                --limit {{ inventory_hostname }}
                -e "start_generator=true"
                -e "profile_prefix={{ item }}"
                -e "dummy_batch={{ dummy_batch | default(false) }}"
                {% if auth_min_seconds is defined %}-e "auth_min_seconds={{ auth_min_seconds }}"{% endif %}
                {% if auth_max_seconds is defined %}-e "auth_max_seconds={{ auth_max_seconds }}"{% endif %}
                {% if batch_size is defined %}-e "batch_size={{ batch_size }}"{% endif %}
                {% if create_download_tasks is defined %}-e "create_download_tasks={{ create_download_tasks }}"{% endif %}
                {% if formats_to_download is defined %}-e "formats_to_download={{ formats_to_download }}"{% endif %}
              delegate_to: localhost
              changed_when: true
              loop: "{{ profile_prefixes }}"
              loop_control:
                loop_var: item
                label: "profile: {{ item }}"
              when: auth_workers_per_profile | default(0) | int > 0

        - name: "WORKAROUND: Align download worker config with auth worker config to bypass inventory bug"
          ansible.builtin.set_fact:
            download_workers_total: "{{ auth_workers_total | default(0) }}"
            download_workers_per_profile: "{{ auth_workers_per_profile | default(0) }}"

        - name: "Start download simulator(s)"
          ansible.builtin.include_tasks: tasks/start-download-simulators.yml
          when: profile_prefixes is defined and profile_prefixes | length > 0

    - name: "Start only auth generators on workers"
      when: action == "start-auth"
      block:
        - name: "Set combined profile prefixes string"
          ansible.builtin.set_fact:
            combined_prefixes: "{{ profile_prefixes | default([]) | join(',') }}"
          when: profile_prefixes is defined and profile_prefixes | length > 0

        - name: "Start auth generator(s)"
          when: profile_prefixes is defined and profile_prefixes | length > 0
          block:
            - name: "Start single auth generator for all profiles: {{ combined_prefixes | default('none') }}"
              ansible.builtin.command: >-
                ansible-playbook {{ playbook_dir }}/playbook-stress-auth-generator.yml
                -i {{ inventory_file }}
                --limit {{ inventory_hostname }}
                -e "start_generator=true"
                -e "profile_prefix={{ combined_prefixes }}"
                -e "dummy_batch={{ dummy_batch | default(false) }}"
                {% if auth_min_seconds is defined %}-e "auth_min_seconds={{ auth_min_seconds }}"{% endif %}
                {% if auth_max_seconds is defined %}-e "auth_max_seconds={{ auth_max_seconds }}"{% endif %}
                {% if batch_size is defined %}-e "batch_size={{ batch_size }}"{% endif %}
                {% if create_download_tasks is defined %}-e "create_download_tasks={{ create_download_tasks }}"{% endif %}
                {% if formats_to_download is defined %}-e "formats_to_download={{ formats_to_download }}"{% endif %}
              delegate_to: localhost
              changed_when: true
              when: (auth_workers_per_profile | default(0) | int == 0) and (auth_workers_total | default(0) | int > 0)

            - name: "Start parallel auth generators for each profile"
              ansible.builtin.command: >-
                ansible-playbook {{ playbook_dir }}/playbook-stress-auth-generator.yml
                -i {{ inventory_file }}
                --limit {{ inventory_hostname }}
                -e "start_generator=true"
                -e "profile_prefix={{ item }}"
                -e "dummy_batch={{ dummy_batch | default(false) }}"
                {% if auth_min_seconds is defined %}-e "auth_min_seconds={{ auth_min_seconds }}"{% endif %}
                {% if auth_max_seconds is defined %}-e "auth_max_seconds={{ auth_max_seconds }}"{% endif %}
                {% if batch_size is defined %}-e "batch_size={{ batch_size }}"{% endif %}
                {% if create_download_tasks is defined %}-e "create_download_tasks={{ create_download_tasks }}"{% endif %}
                {% if formats_to_download is defined %}-e "formats_to_download={{ formats_to_download }}"{% endif %}
              delegate_to: localhost
              changed_when: true
              loop: "{{ profile_prefixes }}"
              loop_control:
                loop_var: item
                label: "profile: {{ item }}"
              when: auth_workers_per_profile | default(0) | int > 0

    - name: "Start only download simulators on workers"
      when: action == "start-download"
      block:
        - name: "WORKAROUND: Align download worker config with auth worker config to bypass inventory bug"
          ansible.builtin.set_fact:
            download_workers_total: "{{ auth_workers_total | default(0) }}"
            download_workers_per_profile: "{{ auth_workers_per_profile | default(0) }}"

        - name: "Set combined profile prefixes string"
          ansible.builtin.set_fact:
            combined_prefixes: "{{ profile_prefixes | default([]) | join(',') }}"
          when: profile_prefixes is defined and profile_prefixes | length > 0

        - name: "Start download simulator(s)"
          ansible.builtin.include_tasks: tasks/start-download-simulators.yml
          when: profile_prefixes is defined and profile_prefixes | length > 0


    - name: "Stop only auth generators on workers"
      when: action == "stop-auth"
      block:
        - name: "Set combined profile prefixes string"
          ansible.builtin.set_fact:
            combined_prefixes: "{{ profile_prefixes | default([]) | join(',') }}"
          when: profile_prefixes is defined and profile_prefixes | length > 0

        - name: "Gracefully stop auth generator(s) via playbook call"
          when: profile_prefixes is defined and profile_prefixes | length > 0
          block:
            - name: "Stop single auth generator for all profiles: {{ combined_prefixes | default('none') }}"
              ansible.builtin.command: >-
                ansible-playbook {{ playbook_dir }}/playbook-stress-auth-generator.yml
                -i {{ inventory_file }}
                --limit {{ inventory_hostname }}
                -e "stop_generator=true"
                -e "profile_prefix={{ combined_prefixes }}"
              delegate_to: localhost
              changed_when: true
              when: (auth_workers_per_profile | default(0) | int == 0) and (auth_workers_total | default(0) | int > 0)

            - name: "Stop parallel auth generators for each profile"
              ansible.builtin.command: >-
                ansible-playbook {{ playbook_dir }}/playbook-stress-auth-generator.yml
                -i {{ inventory_file }}
                --limit {{ inventory_hostname }}
                -e "stop_generator=true"
                -e "profile_prefix={{ item }}"
              delegate_to: localhost
              changed_when: true
              loop: "{{ profile_prefixes }}"
              loop_control:
                loop_var: item
                label: "profile: {{ item }}"
              when: auth_workers_per_profile | default(0) | int > 0

    - name: "Stop only download simulators on workers"
      when: action == "stop-download"
      block:
        - name: "WORKAROUND: Align download worker config with auth worker config to bypass inventory bug"
          ansible.builtin.set_fact:
            download_workers_total: "{{ auth_workers_total | default(0) }}"
            download_workers_per_profile: "{{ auth_workers_per_profile | default(0) }}"

        - name: "Set combined profile prefixes string"
          ansible.builtin.set_fact:
            combined_prefixes: "{{ profile_prefixes | default([]) | join(',') }}"
          when: profile_prefixes is defined and profile_prefixes | length > 0

        - name: "Stop single download simulator group"
          ansible.builtin.shell:
            cmd: |
              for session in $(tmux list-sessions -F "#{session_name}" 2>/dev/null | grep -E '^stress-download-worker-[0-9]+$'); do
                tmux kill-session -t "$session"
              done || true
          when: download_workers_total | default(0) | int > 0
          changed_when: true
          ignore_errors: yes

        - name: "Stop parallel download simulators for each profile"
          ansible.builtin.command: "tmux kill-session -t stress-download-{{ item }}"
          loop: "{{ profile_prefixes }}"
          loop_control:
            loop_var: item
            label: "profile: {{ item }}"
          when: (download_workers_total | default(0) | int == 0) and (download_workers_per_profile | default(0) | int > 0)
          changed_when: true
          ignore_errors: yes

    - name: "Stop all worker generators and simulators"
      when: action == "stop"
      block:
        - name: Kill all tmux sessions starting with 'stress-' on this worker
          ansible.builtin.shell:
            cmd: |
              for session in $(tmux list-sessions -F "#{session_name}" 2>/dev/null | grep -E "^stress-"); do
                tmux kill-session -t "$session"
              done || true
          ignore_errors: yes
          changed_when: false

        - name: Kill all ytops-client processes on this worker
          ansible.builtin.shell:
            cmd: |
              # Gracefully terminate
              ps aux | grep "[y]tops-client.*stress-policy" | awk '{print $2}' | xargs kill >/dev/null 2>&1 || true
              sleep 0.5
              # Force kill
              ps aux | grep "[y]tops-client.*stress-policy" | awk '{print $2}' | xargs kill -9 >/dev/null 2>&1 || true
          ignore_errors: yes
          changed_when: false

    - name: "Check status of all configured generators and simulators"
      when: action == "status"
      block:
        - name: "Set combined profile prefixes string"
          ansible.builtin.set_fact:
            combined_prefixes: "{{ profile_prefixes | default([]) | join(',') }}"
          when: profile_prefixes is defined and profile_prefixes | length > 0

        - name: "Check auth generator status"
          when: profile_prefixes is defined and profile_prefixes | length > 0
          block:
            - name: "Check single auth generator for all profiles: {{ combined_prefixes | default('none') }}"
              ansible.builtin.command: >-
                ansible-playbook {{ playbook_dir }}/playbook-stress-auth-generator.yml
                -i {{ inventory_file }}
                --limit {{ inventory_hostname }}
                -e "check_status=true"
                -e "profile_prefix={{ combined_prefixes }}"
              delegate_to: localhost
              changed_when: false
              register: auth_status_check_combined
              when: (auth_workers_per_profile | default(0) | int == 0) and (auth_workers_total | default(0) | int > 0)

            - name: "Display combined auth generator status for {{ inventory_hostname }}"
              ansible.builtin.debug:
                var: auth_status_check_combined.stdout_lines
              when: auth_status_check_combined is defined and auth_status_check_combined.stdout_lines is defined

            - name: "Check parallel auth generators for each profile"
              ansible.builtin.command: >-
                ansible-playbook {{ playbook_dir }}/playbook-stress-auth-generator.yml
                -i {{ inventory_file }}
                --limit {{ inventory_hostname }}
                -e "check_status=true"
                -e "profile_prefix={{ item }}"
              delegate_to: localhost
              changed_when: false
              loop: "{{ profile_prefixes }}"
              loop_control:
                loop_var: item
                label: "profile: {{ item }}"
              register: auth_status_check_parallel
              when: auth_workers_per_profile | default(0) | int > 0

            - name: "Display parallel auth generator status for {{ inventory_hostname }}"
              ansible.builtin.debug:
                msg: "{{ item.stdout_lines }}"
              loop: "{{ auth_status_check_parallel.results | default([]) }}"
              loop_control:
                label: "status for profile: {{ item.item }}"
              when: auth_status_check_parallel is defined and auth_status_check_parallel.results is defined

        - name: "Check download simulator status"
          when: profile_prefixes is defined and profile_prefixes | length > 0
          block:
            - name: "Check single download simulator for all profiles: {{ combined_prefixes | default('none') }}"
              ansible.builtin.command: >-
                ansible-playbook {{ playbook_dir }}/playbook-stress-download-simulation.yml
                -i {{ inventory_file }}
                --limit {{ inventory_hostname }}
                -e "check_status=true"
                -e "profile_prefix={{ combined_prefixes }}"
              delegate_to: localhost
              changed_when: false
              register: download_status_check_combined
              when: (download_workers_per_profile | default(0) | int == 0) and (download_workers_total | default(0) | int > 0)

            - name: "Display combined download simulator status for {{ inventory_hostname }}"
              ansible.builtin.debug:
                var: download_status_check_combined.stdout_lines
              when: download_status_check_combined is defined and download_status_check_combined.stdout_lines is defined

            - name: "Check parallel download simulators for each profile"
              ansible.builtin.command: >-
                ansible-playbook {{ playbook_dir }}/playbook-stress-download-simulation.yml
                -i {{ inventory_file }}
                --limit {{ inventory_hostname }}
                -e "check_status=true"
                -e "profile_prefix={{ item }}"
              delegate_to: localhost
              changed_when: false
              loop: "{{ profile_prefixes }}"
              loop_control:
                loop_var: item
                label: "profile: {{ item }}"
              register: download_status_check_parallel
              when: download_workers_per_profile | default(0) | int > 0

            - name: "Display parallel download simulator status for {{ inventory_hostname }}"
              ansible.builtin.debug:
                msg: "{{ item.stdout_lines }}"
              loop: "{{ download_status_check_parallel.results | default([]) }}"
              loop_control:
                label: "status for profile: {{ item.item }}"
              when: download_status_check_parallel is defined and download_status_check_parallel.results is defined