---
- name: Validate WireGuard variables
  ansible.builtin.import_tasks: validate.yaml

- name: Install WireGuard
  ansible.builtin.apt:
    name:
      - wireguard
      - wireguard-tools
    state: present
    update_cache: true

- name: Install UFW (if needed)
  ansible.builtin.apt:
    name: ufw
    state: present
  when: wireguard_enable_ufw

- name: Create WireGuard directories
  ansible.builtin.file:
    path: "{{ item }}"
    state: directory
    mode: '0700'
    owner: root
    group: root
  loop:
    - "{{ wireguard_config_dir }}"
    - "{{ wireguard_keys_dir }}"
    - "{{ wireguard_clients_dir }}"

- name: Check if server private key exists
  ansible.builtin.stat:
    path: "{{ wireguard_keys_dir }}/server_private.key"
  register: server_private_key_stat

- name: Generate server private key
  ansible.builtin.shell: |
    wg genkey > {{ wireguard_keys_dir }}/server_private.key
    chmod 600 {{ wireguard_keys_dir }}/server_private.key
  when: not server_private_key_stat.stat.exists

- name: Generate server public key
  ansible.builtin.shell: |
    wg pubkey < {{ wireguard_keys_dir }}/server_private.key > {{ wireguard_keys_dir }}/server_public.key
    chmod 600 {{ wireguard_keys_dir }}/server_public.key
  when: not server_private_key_stat.stat.exists

- name: Read server private key
  ansible.builtin.slurp:
    src: "{{ wireguard_keys_dir }}/server_private.key"
  register: server_private_key_content

- name: Read server public key
  ansible.builtin.slurp:
    src: "{{ wireguard_keys_dir }}/server_public.key"
  register: server_public_key_content

- name: Set server key facts
  ansible.builtin.set_fact:
    wireguard_server_private_key: "{{ server_private_key_content.content | b64decode | trim }}"
    wireguard_server_public_key: "{{ server_public_key_content.content | b64decode | trim }}"

- name: Generate client keys
  block:
    - name: Check if client private key exists
      ansible.builtin.stat:
        path: "{{ wireguard_keys_dir }}/{{ item.name }}_private.key"
      register: client_private_key_stats
      loop: "{{ wireguard_clients }}"

    - name: Generate client private keys
      ansible.builtin.shell: |
        wg genkey > {{ wireguard_keys_dir }}/{{ item.item.name }}_private.key
        chmod 600 {{ wireguard_keys_dir }}/{{ item.item.name }}_private.key
      loop: "{{ client_private_key_stats.results }}"
      when: not item.stat.exists

    - name: Check if client public key exists
      ansible.builtin.stat:
        path: "{{ wireguard_keys_dir }}/{{ item.name }}_public.key"
      register: client_public_key_stats
      loop: "{{ wireguard_clients }}"

    - name: Generate client public keys
      ansible.builtin.shell: |
        wg pubkey < {{ wireguard_keys_dir }}/{{ item.item.name }}_private.key > {{ wireguard_keys_dir }}/{{ item.item.name }}_public.key
        chmod 600 {{ wireguard_keys_dir }}/{{ item.item.name }}_public.key
      loop: "{{ client_public_key_stats.results }}"
      when: not item.stat.exists

    - name: Check if client preshared key exists
      ansible.builtin.stat:
        path: "{{ wireguard_keys_dir }}/{{ item.name }}_preshared.key"
      register: client_preshared_key_stats
      loop: "{{ wireguard_clients }}"
      when: item.preshared_key | default(false)

    - name: Generate client preshared keys (if enabled)
      ansible.builtin.shell: |
        wg genpsk > {{ wireguard_keys_dir }}/{{ item.item.name }}_preshared.key
        chmod 600 {{ wireguard_keys_dir }}/{{ item.item.name }}_preshared.key
      loop: "{{ client_preshared_key_stats.results | default([]) }}"
      when:
        - client_preshared_key_stats.results is defined
        - not item.stat.exists

    - name: Read client private keys
      ansible.builtin.slurp:
        src: "{{ wireguard_keys_dir }}/{{ item.name }}_private.key"
      register: client_private_keys
      loop: "{{ wireguard_clients }}"

    - name: Read client public keys
      ansible.builtin.slurp:
        src: "{{ wireguard_keys_dir }}/{{ item.name }}_public.key"
      register: client_public_keys
      loop: "{{ wireguard_clients }}"

    - name: Read client preshared keys
      ansible.builtin.slurp:
        src: "{{ wireguard_keys_dir }}/{{ item.name }}_preshared.key"
      register: client_preshared_keys
      loop: "{{ wireguard_clients }}"
      when: item.preshared_key | default(false)
      ignore_errors: true

  when: wireguard_clients | length > 0

- name: Process client keys into structured data
  ansible.builtin.set_fact:
    wireguard_client_configs: "{{ wireguard_client_configs | default([]) + [client_config] }}"
  loop: "{{ wireguard_clients }}"
  loop_control:
    index_var: client_index
  vars:
    private_key: "{{ client_private_keys.results[client_index].content | b64decode | trim }}"
    public_key: "{{ client_public_keys.results[client_index].content | b64decode | trim }}"
    preshared_key: "{{ client_preshared_keys.results[client_index].content | b64decode | trim if (item.preshared_key | default(false) and client_preshared_keys.results[client_index] is defined and not client_preshared_keys.results[client_index].failed | default(false)) else '' }}"
    client_config:
      name: "{{ item.name }}"
      ipv4: "{{ item.ipv4 }}"
      ipv6: "{{ item.ipv6 }}"
      private_key: "{{ private_key }}"
      public_key: "{{ public_key }}"
      preshared_key: "{{ preshared_key }}"
      persistent_keepalive: "{{ item.persistent_keepalive | default(25) }}"
  when: wireguard_clients | length > 0

- name: Enable IPv4 forwarding
  ansible.posix.sysctl:
    name: net.ipv4.ip_forward
    value: '1'
    sysctl_set: true
    state: present
    reload: true
  when: wireguard_enable_forwarding

- name: Enable IPv6 forwarding
  ansible.posix.sysctl:
    name: net.ipv6.conf.all.forwarding
    value: '1'
    sysctl_set: true
    state: present
    reload: true
  when: wireguard_enable_forwarding

- name: Get default network interface
  ansible.builtin.shell: ip route | grep default | awk '{print $5}' | head -n1
  register: default_interface
  changed_when: false

- name: Set default interface fact
  ansible.builtin.set_fact:
    wireguard_default_interface: "{{ default_interface.stdout }}"

- name: Deploy WireGuard server configuration
  ansible.builtin.template:
    src: wg0.conf.j2
    dest: "{{ wireguard_config_dir }}/{{ wireguard_interface }}.conf"
    mode: '0600'
    owner: root
    group: root
  notify: restart wireguard

- name: Configure UFW for WireGuard
  block:
    - name: Allow WireGuard port through UFW
      community.general.ufw:
        rule: allow
        port: "{{ wireguard_port }}"
        proto: udp

    - name: Allow forwarding on WireGuard interface
      community.general.ufw:
        rule: allow
        route: true
        interface_in: "{{ wireguard_interface }}"
  when: wireguard_enable_ufw

- name: Enable and start WireGuard service
  ansible.builtin.systemd:
    name: "wg-quick@{{ wireguard_interface }}"
    enabled: true
    state: started

- name: Generate client configuration files
  ansible.builtin.template:
    src: client.conf.j2
    dest: "{{ wireguard_clients_dir }}/{{ item.name }}.conf"
    mode: '0600'
    owner: root
    group: root
  loop: "{{ wireguard_client_configs | default([]) }}"
  when: wireguard_clients | length > 0

- name: Fetch client configuration files to local machine
  ansible.builtin.fetch:
    src: "{{ wireguard_clients_dir }}/{{ item.name }}.conf"
    dest: "./wireguard-configs/{{ inventory_hostname }}/{{ item.name }}.conf"
    flat: true
  loop: "{{ wireguard_clients }}"
  when: wireguard_clients | length > 0