---
- name: Provision WireGuard VPN for new client
  hosts: vpn_servers
  become: yes
  vars:
    client_dir: /etc/wireguard/clients
    test_client_dir: /etc/wireguard/test_clients
    wg_interface: wg0
    server_dir: /etc/wireguard
    server_ip: 10.8.0.1/24
    server_port: 51820
    server_endpoint: "{{ ansible_host | default(inventory_hostname) }}"
    is_test: false  # Default to production mode
    test_duration_minutes: 30  # Default test duration

  pre_tasks:
    - name: Check if WireGuard is installed
      package_facts:
        manager: auto

    - name: Install WireGuard (Debian/Ubuntu)
      apt:
        name: 
          - wireguard
          - wireguard-tools
        state: present
        update_cache: yes
      when: 
        - ansible_facts['os_family'] == "Debian"
        - "'wireguard' not in ansible_facts.packages"

    - name: Install WireGuard (RHEL/CentOS)
      dnf:
        name: 
          - wireguard-tools
          - wireguard-dkms
        state: present
      when: 
        - ansible_facts['os_family'] == "RedHat"
        - "'wireguard-tools' not in ansible_facts.packages"

    - name: Ensure WireGuard kernel module is loaded
      modprobe:
        name: wireguard
        state: present

    - name: Verify WireGuard installation
      command: which wg
      register: wg_check
      failed_when: wg_check.rc != 0
      changed_when: false

  tasks:
    - name: Debug invoice ID and test status
      debug:
        msg: 
          - "Processing invoice ID: {{ invoice_id }}"
          - "Test mode: {{ is_test }}"
          - "Test duration: {{ test_duration_minutes if is_test else 'N/A' }}"

    - name: Create required directories
      file:
        path: "{{ item }}"
        state: directory
        mode: '0700'
      with_items:
        - "{{ client_dir }}"
        - "{{ test_client_dir }}"
        - "{{ server_dir }}"

    - name: Set working directory based on mode
      set_fact:
        working_client_dir: "{{ test_client_dir if is_test else client_dir }}"
        
    - name: Check if server keys exist
      stat:
        path: "{{ server_dir }}/{{ wg_interface }}.conf"
      register: server_config
        
    - name: Generate server private key if not exists
      shell: wg genkey
      register: server_private_key
      when: not server_config.stat.exists
      
    - name: Save server private key
      copy:
        content: "{{ server_private_key.stdout }}"
        dest: "{{ server_dir }}/private.key"
        mode: '0600'
      when: not server_config.stat.exists
      
    - name: Generate server public key
      shell: "cat {{ server_dir }}/private.key | wg pubkey"
      register: server_public_key
      when: not server_config.stat.exists
      
    - name: Save server public key
      copy:
        content: "{{ server_public_key.stdout }}"
        dest: "{{ server_dir }}/public.key"
        mode: '0644'
      when: not server_config.stat.exists

    - name: Create initial server config
      template:
        src: templates/server.conf.j2
        dest: "{{ server_dir }}/{{ wg_interface }}.conf"
        mode: '0600'
      when: not server_config.stat.exists
      notify: restart wireguard
      
    - name: Ensure client directory exists
      file:
        path: "{{ working_client_dir }}/{{ invoice_id }}"
        state: directory
        mode: '0700'

    # Generate keys - no longer differentiating between test and production
    - name: Generate private key
      shell: wg genkey
      register: private_key
      changed_when: false

    - name: Generate public key
      shell: echo "{{ private_key.stdout }}" | wg pubkey
      register: public_key
      changed_when: false

    - name: Save private key
      copy:
        content: "{{ private_key.stdout }}"
        dest: "{{ working_client_dir }}/{{ invoice_id }}/private.key"
        mode: '0600'

    - name: Save public key
      copy:
        content: "{{ public_key.stdout }}"
        dest: "{{ working_client_dir }}/{{ invoice_id }}/public.key"
        mode: '0644'

    - name: Read server public key
      shell: "cat {{ server_dir }}/public.key"
      register: server_public_key_read
      changed_when: false
        
    - name: Get next available IP
      shell: |
        last_ip=$(grep -h '^Address' {{ working_client_dir }}/*/wg0.conf 2>/dev/null | tail -n1 | grep -oE '[0-9]+$' || echo 1)
        echo $((last_ip + 1))
      register: next_ip
      
    - name: Generate client config
      template:
        src: templates/client.conf.j2
        dest: "{{ working_client_dir }}/{{ invoice_id }}/wg0.conf"
        mode: '0600'
      vars:
        client_ip: "10.8.0.{{ next_ip.stdout }}"
        server_pubkey: "{{ server_public_key_read.stdout }}"
        client_private_key: "{{ private_key.stdout }}"
        
    - name: Add client to server config
      blockinfile:
        path: "{{ server_dir }}/{{ wg_interface }}.conf"
        marker: "# {mark} ANSIBLE MANAGED BLOCK FOR {{ invoice_id }}"
        block: |
          [Peer]
          PublicKey = {{ public_key.stdout }}
          AllowedIPs = 10.8.0.{{ next_ip.stdout }}/32
          {% if is_test %}# Test config expires: {{ ansible_date_time.iso8601 }}{% endif %}
      notify: restart wireguard

    # Calculate cleanup time for test configurations
    - name: Calculate cleanup time
      when: is_test
      set_fact:
        cleanup_minute: "{{ (ansible_date_time.minute | int + (test_duration_minutes | int)) % 60 }}"
        cleanup_hour: "{{ (ansible_date_time.hour | int + ((ansible_date_time.minute | int + (test_duration_minutes | int)) // 60)) % 24 }}"

    - name: Add cleanup cronjob for test configs
      when: is_test
      cron:
        name: "cleanup_test_vpn_{{ invoice_id }}"
        minute: "{{ cleanup_minute }}"
        hour: "{{ cleanup_hour }}"
        job: "ansible-playbook {{ playbook_dir }}/vpn_cleanup.yml -e 'invoice_id={{ invoice_id }} is_test=true'"
        state: present

    - name: Log provision completion
      shell: |
        logger -t vpn-provision "Provisioned VPN for {{ invoice_id }} ({{ 'test' if is_test else 'production' }}){% if is_test %} - expires in {{ test_duration_minutes }} minutes{% endif %}"

  handlers:
    - name: restart wireguard
      service:
        name: "wg-quick@{{ wg_interface }}"
        state: restarted