#!/usr/bin/python

# Copyright (c) 2013, Alexander Bulimov <lazywolf0@gmail.com>
# Based on lvol module by Jeroen Hoekx <jeroen.hoekx@dsquare.be>
# GNU General Public License v3.0+ (see LICENSES/GPL-3.0-or-later.txt or https://www.gnu.org/licenses/gpl-3.0.txt)
# SPDX-License-Identifier: GPL-3.0-or-later

from __future__ import annotations

DOCUMENTATION = r"""
author:
  - Alexander Bulimov (@abulimov)
module: lvg
short_description: Configure LVM volume groups
description:
  - This module creates, removes or resizes volume groups.
extends_documentation_fragment:
  - community.general.attributes
attributes:
  check_mode:
    support: full
  diff_mode:
    support: none
options:
  vg:
    description:
      - The name of the volume group.
    type: str
    required: true
  pvs:
    description:
      - List of comma-separated devices to use as physical devices in this volume group.
      - Required when creating or resizing volume group.
      - The module runs C(pvcreate) if needed.
      - This parameter defines the B(desired state) of the physical volumes in the volume group.
        When the volume group already exists, physical volumes not listed here are removed from it by default.
        To add physical volumes without removing existing unlisted ones, set O(remove_extra_pvs=false).
    type: list
    elements: str
  pesize:
    description:
      - The size of the physical extent. O(pesize) must be a power of 2 of at least 1 sector (where the sector size is the
        largest sector size of the PVs currently used in the VG), or at least 128KiB.
      - O(pesize) can be optionally suffixed by a UNIT (k/K/m/M/g/G), default unit is megabyte.
    type: str
    default: "4"
  pv_options:
    description:
      - Additional options to pass to C(pvcreate) when creating the volume group.
    type: str
    default: ''
  pvresize:
    description:
      - If V(true), resize the physical volume to the maximum available size.
    type: bool
    default: false
    version_added: '0.2.0'
  vg_options:
    description:
      - Additional options to pass to C(vgcreate) when creating the volume group.
    type: str
    default: ''
  state:
    description:
      - Control if the volume group exists and its state.
      - The states V(active) and V(inactive) implies V(present) state. Added in 7.1.0.
      - If V(active) or V(inactive), the module manages the VG's logical volumes current state. The module also handles the
        VG's autoactivation state if supported unless when creating a volume group and the autoactivation option specified
        in O(vg_options).
    type: str
    choices: [absent, present, active, inactive]
    default: present
  force:
    description:
      - If V(true), allows to remove volume group with logical volumes.
    type: bool
    default: false
  reset_vg_uuid:
    description:
      - Whether the volume group's UUID is regenerated.
      - This is B(not idempotent). Specifying this parameter always results in a change.
    type: bool
    default: false
    version_added: 7.1.0
  reset_pv_uuid:
    description:
      - Whether the volume group's physical volumes' UUIDs are regenerated.
      - This is B(not idempotent). Specifying this parameter always results in a change.
    type: bool
    default: false
    version_added: 7.1.0
  remove_extra_pvs:
    description:
      - Remove physical volumes from the volume group which are not in O(pvs).
    type: bool
    default: true
    version_added: 10.4.0
seealso:
  - module: community.general.filesystem
  - module: community.general.lvol
  - module: community.general.parted
notes:
  - This module does not modify PE size for already present volume group.
"""

EXAMPLES = r"""
- name: Create a volume group on top of /dev/sda1 with physical extent size = 32MB
  community.general.lvg:
    vg: vg.services
    pvs: /dev/sda1
    pesize: 32

- name: Create a volume group on top of /dev/sdb with physical extent size = 128KiB
  community.general.lvg:
    vg: vg.services
    pvs: /dev/sdb
    pesize: 128K

# If, for example, we already have VG vg.services on top of /dev/sdb1,
# this VG will be extended by /dev/sdc5.  Or if vg.services was created on
# top of /dev/sda5, we first extend it with /dev/sdb1 and /dev/sdc5,
# and then reduce by /dev/sda5.
- name: Create or resize a volume group on top of /dev/sdb1 and /dev/sdc5.
  community.general.lvg:
    vg: vg.services
    pvs:
      - /dev/sdb1
      - /dev/sdc5

- name: Remove a volume group with name vg.services
  community.general.lvg:
    vg: vg.services
    state: absent

- name: Create a volume group on top of /dev/sda3 and resize the volume group /dev/sda3 to the maximum possible
  community.general.lvg:
    vg: resizableVG
    pvs: /dev/sda3
    pvresize: true

- name: Deactivate a volume group
  community.general.lvg:
    state: inactive
    vg: vg.services

- name: Activate a volume group
  community.general.lvg:
    state: active
    vg: vg.services

- name: Add new PVs to volume group without removing existing ones
  community.general.lvg:
    vg: vg.services
    pvs: /dev/sdb1,/dev/sdc1
    remove_extra_pvs: false
    state: present

- name: Reset a volume group UUID
  community.general.lvg:
    state: inactive
    vg: vg.services
    reset_vg_uuid: true

- name: Reset both volume group and pv UUID
  community.general.lvg:
    state: inactive
    vg: vg.services
    pvs:
      - /dev/sdb1
      - /dev/sdc5
    reset_vg_uuid: true
    reset_pv_uuid: true
"""

import itertools
import os

from ansible.module_utils.basic import AnsibleModule

VG_AUTOACTIVATION_OPT = "--setautoactivation"


def parse_vgs(data):
    vgs = []
    for line in data.splitlines():
        parts = line.strip().split(";")
        vgs.append(
            {
                "name": parts[0],
                "pv_count": int(parts[1]),
                "lv_count": int(parts[2]),
            }
        )
    return vgs


def find_mapper_device_name(module, dm_device):
    dmsetup_cmd = module.get_bin_path("dmsetup", True)
    mapper_prefix = "/dev/mapper/"
    rc, dm_name, err = module.run_command([dmsetup_cmd, "info", "-C", "--noheadings", "-o", "name", dm_device])
    if rc != 0:
        module.fail_json(msg="Failed executing dmsetup command.", rc=rc, err=err)
    mapper_device = mapper_prefix + dm_name.rstrip()
    return mapper_device


def parse_pvs(module, data):
    pvs = []
    dm_prefix = "/dev/dm-"
    for line in data.splitlines():
        parts = line.strip().split(";")
        if parts[0].startswith(dm_prefix):
            parts[0] = find_mapper_device_name(module, parts[0])
        pvs.append(
            {
                "name": parts[0],
                "vg_name": parts[1],
            }
        )
    return pvs


def find_vg(module, vg):
    if not vg:
        return None
    vgs_cmd = module.get_bin_path("vgs", True)
    dummy, current_vgs, dummy = module.run_command(
        [vgs_cmd, "--noheadings", "-o", "vg_name,pv_count,lv_count", "--separator", ";"], check_rc=True
    )

    vgs = parse_vgs(current_vgs)

    for test_vg in vgs:
        if test_vg["name"] == vg:
            this_vg = test_vg
            break
    else:
        this_vg = None

    return this_vg


def is_autoactivation_supported(module, vg_cmd):
    autoactivation_supported = False
    dummy, vgchange_opts, dummy = module.run_command([vg_cmd, "--help"], check_rc=True)

    if VG_AUTOACTIVATION_OPT in vgchange_opts:
        autoactivation_supported = True

    return autoactivation_supported


def activate_vg(module, vg, active):
    changed = False
    vgchange_cmd = module.get_bin_path("vgchange", True)
    vgs_cmd = module.get_bin_path("vgs", True)
    vgs_fields = ["lv_attr"]

    autoactivation_enabled = False
    autoactivation_supported = is_autoactivation_supported(module=module, vg_cmd=vgchange_cmd)

    if autoactivation_supported:
        vgs_fields.append("autoactivation")

    vgs_cmd_with_opts = [vgs_cmd, "--noheadings", "-o", ",".join(vgs_fields), "--separator", ";", vg]
    dummy, current_vg_lv_states, dummy = module.run_command(vgs_cmd_with_opts, check_rc=True)

    lv_active_count = 0
    lv_inactive_count = 0

    for line in current_vg_lv_states.splitlines():
        parts = line.strip().split(";")
        if parts[0][4] == "a":
            lv_active_count += 1
        else:
            lv_inactive_count += 1
        if autoactivation_supported:
            autoactivation_enabled = autoactivation_enabled or parts[1] == "enabled"

    activate_flag = None
    if active and lv_inactive_count > 0:
        activate_flag = "y"
    elif not active and lv_active_count > 0:
        activate_flag = "n"

    # Extra logic necessary because vgchange returns error when autoactivation is already set
    if autoactivation_supported:
        if active and not autoactivation_enabled:
            if module.check_mode:
                changed = True
            else:
                module.run_command([vgchange_cmd, VG_AUTOACTIVATION_OPT, "y", vg], check_rc=True)
                changed = True
        elif not active and autoactivation_enabled:
            if module.check_mode:
                changed = True
            else:
                module.run_command([vgchange_cmd, VG_AUTOACTIVATION_OPT, "n", vg], check_rc=True)
                changed = True

    if activate_flag is not None:
        if module.check_mode:
            changed = True
        else:
            module.run_command([vgchange_cmd, "--activate", activate_flag, vg], check_rc=True)
            changed = True

    return changed


def append_vgcreate_options(module, state, vgoptions):
    vgcreate_cmd = module.get_bin_path("vgcreate", True)

    autoactivation_supported = is_autoactivation_supported(module=module, vg_cmd=vgcreate_cmd)

    if autoactivation_supported and state in ["active", "inactive"]:
        if VG_AUTOACTIVATION_OPT not in vgoptions:
            if state == "active":
                vgoptions += [VG_AUTOACTIVATION_OPT, "y"]
            else:
                vgoptions += [VG_AUTOACTIVATION_OPT, "n"]


def get_pv_values_for_resize(module, device):
    pvdisplay_cmd = module.get_bin_path("pvdisplay", True)
    pvdisplay_ops = [
        "--units",
        "b",
        "--columns",
        "--noheadings",
        "--nosuffix",
        "--separator",
        ";",
        "-o",
        "dev_size,pv_size,pe_start,vg_extent_size",
    ]
    pvdisplay_cmd_device_options = [pvdisplay_cmd, device] + pvdisplay_ops

    dummy, pv_values, dummy = module.run_command(pvdisplay_cmd_device_options, check_rc=True)

    values = pv_values.strip().split(";")

    dev_size = int(values[0])
    pv_size = int(values[1])
    pe_start = int(values[2])
    vg_extent_size = int(values[3])

    return (dev_size, pv_size, pe_start, vg_extent_size)


def resize_pv(module, device):
    changed = False
    pvresize_cmd = module.get_bin_path("pvresize", True)

    dev_size, pv_size, pe_start, vg_extent_size = get_pv_values_for_resize(module=module, device=device)
    if (dev_size - (pe_start + pv_size)) > vg_extent_size:
        if module.check_mode:
            changed = True
        else:
            # If there is a missing pv on the machine, versions of pvresize rc indicates failure.
            rc, out, err = module.run_command([pvresize_cmd, device])
            dummy, new_pv_size, dummy, dummy = get_pv_values_for_resize(module=module, device=device)
            if pv_size == new_pv_size:
                module.fail_json(msg="Failed executing pvresize command.", rc=rc, err=err, out=out)
            else:
                changed = True

    return changed


def reset_uuid_pv(module, device):
    changed = False
    pvs_cmd = module.get_bin_path("pvs", True)
    pvs_cmd_with_opts = [pvs_cmd, "--noheadings", "-o", "uuid", device]
    pvchange_cmd = module.get_bin_path("pvchange", True)
    pvchange_cmd_with_opts = [pvchange_cmd, "-u", device]

    dummy, orig_uuid, dummy = module.run_command(pvs_cmd_with_opts, check_rc=True)

    if module.check_mode:
        changed = True
    else:
        # If there is a missing pv on the machine, pvchange rc indicates failure.
        pvchange_rc, pvchange_out, pvchange_err = module.run_command(pvchange_cmd_with_opts)
        dummy, new_uuid, dummy = module.run_command(pvs_cmd_with_opts, check_rc=True)
        if orig_uuid.strip() == new_uuid.strip():
            module.fail_json(
                msg=f"PV ({device}) UUID change failed", rc=pvchange_rc, err=pvchange_err, out=pvchange_out
            )
        else:
            changed = True

    return changed


def reset_uuid_vg(module, vg):
    changed = False
    vgchange_cmd = module.get_bin_path("vgchange", True)
    vgchange_cmd_with_opts = [vgchange_cmd, "-u", vg]
    if module.check_mode:
        changed = True
    else:
        module.run_command(vgchange_cmd_with_opts, check_rc=True)
        changed = True

    return changed


def main():
    module = AnsibleModule(
        argument_spec=dict(
            vg=dict(type="str", required=True),
            pvs=dict(type="list", elements="str"),
            pesize=dict(type="str", default="4"),
            pv_options=dict(type="str", default=""),
            pvresize=dict(type="bool", default=False),
            vg_options=dict(type="str", default=""),
            state=dict(type="str", default="present", choices=["absent", "present", "active", "inactive"]),
            force=dict(type="bool", default=False),
            reset_vg_uuid=dict(type="bool", default=False),
            reset_pv_uuid=dict(type="bool", default=False),
            remove_extra_pvs=dict(type="bool", default=True),
        ),
        required_if=[
            ["reset_pv_uuid", True, ["pvs"]],
        ],
        supports_check_mode=True,
    )
    module.run_command_environ_update = {"LANGUAGE": "C", "LC_ALL": "C"}

    vg = module.params["vg"]
    state = module.params["state"]
    force = module.boolean(module.params["force"])
    pvresize = module.boolean(module.params["pvresize"])
    pesize = module.params["pesize"]
    pvoptions = module.params["pv_options"].split()
    vgoptions = module.params["vg_options"].split()
    reset_vg_uuid = module.boolean(module.params["reset_vg_uuid"])
    reset_pv_uuid = module.boolean(module.params["reset_pv_uuid"])
    remove_extra_pvs = module.boolean(module.params["remove_extra_pvs"])

    this_vg = find_vg(module=module, vg=vg)
    present_state = state in ["present", "active", "inactive"]
    pvs_required = present_state and this_vg is None
    changed = False

    dev_list = []
    if module.params["pvs"]:
        dev_list = list(module.params["pvs"])
    elif pvs_required:
        module.fail_json(msg="No physical volumes given.")

    # LVM always uses real paths not symlinks so replace symlinks with actual path
    for idx, dev in enumerate(dev_list):
        dev_list[idx] = os.path.realpath(dev)

    if present_state:
        # check given devices
        for test_dev in dev_list:
            if not os.path.exists(test_dev):
                module.fail_json(msg=f"Device {test_dev} not found.")

        # get pv list
        pvs_cmd = module.get_bin_path("pvs", True)
        if dev_list:
            pvs_filter_pv_name = " || ".join(f"pv_name = {x}" for x in itertools.chain(dev_list, module.params["pvs"]))
            pvs_filter_vg_name = f"vg_name = {vg}"
            pvs_filter = ["--select", f"{pvs_filter_pv_name} || {pvs_filter_vg_name}"]
        else:
            pvs_filter = []
        rc, current_pvs, err = module.run_command(
            [pvs_cmd, "--noheadings", "-o", "pv_name,vg_name", "--separator", ";"] + pvs_filter
        )
        if rc != 0:
            module.fail_json(msg="Failed executing pvs command.", rc=rc, err=err)

        # check pv for devices
        pvs = parse_pvs(module, current_pvs)
        used_pvs = [pv for pv in pvs if pv["name"] in dev_list and pv["vg_name"] and pv["vg_name"] != vg]
        if used_pvs:
            module.fail_json(msg=f"Device {used_pvs[0]['name']} is already in {used_pvs[0]['vg_name']} volume group.")

    if this_vg is None:
        if present_state:
            append_vgcreate_options(module=module, state=state, vgoptions=vgoptions)
            # create VG
            if module.check_mode:
                changed = True
            else:
                # create PV
                pvcreate_cmd = module.get_bin_path("pvcreate", True)
                for current_dev in dev_list:
                    rc, dummy, err = module.run_command([pvcreate_cmd] + pvoptions + ["-f", str(current_dev)])
                    if rc == 0:
                        changed = True
                    else:
                        module.fail_json(msg=f"Creating physical volume '{current_dev}' failed", rc=rc, err=err)
                vgcreate_cmd = module.get_bin_path("vgcreate")
                rc, dummy, err = module.run_command([vgcreate_cmd] + vgoptions + ["-s", pesize, vg] + dev_list)
                if rc == 0:
                    changed = True
                else:
                    module.fail_json(msg=f"Creating volume group '{vg}' failed", rc=rc, err=err)
    else:
        if state == "absent":
            if module.check_mode:
                module.exit_json(changed=True)
            else:
                if this_vg["lv_count"] == 0 or force:
                    # remove VG
                    vgremove_cmd = module.get_bin_path("vgremove", True)
                    rc, dummy, err = module.run_command([vgremove_cmd, "--force", vg])
                    if rc == 0:
                        module.exit_json(changed=True)
                    else:
                        module.fail_json(msg=f"Failed to remove volume group {vg}", rc=rc, err=err)
                else:
                    module.fail_json(msg=f"Refuse to remove non-empty volume group {vg} without force=true")
        # activate/deactivate existing VG
        elif state == "active":
            changed = activate_vg(module=module, vg=vg, active=True)
        elif state == "inactive":
            changed = activate_vg(module=module, vg=vg, active=False)

        # reset VG uuid
        if reset_vg_uuid:
            changed = reset_uuid_vg(module=module, vg=vg) or changed

        # resize VG
        if dev_list:
            current_devs = [os.path.realpath(pv["name"]) for pv in pvs if pv["vg_name"] == vg]
            devs_to_remove = list(set(current_devs) - set(dev_list))
            devs_to_add = list(set(dev_list) - set(current_devs))

            if not remove_extra_pvs:
                devs_to_remove = []

            if current_devs:
                if present_state:
                    for device in current_devs:
                        if pvresize:
                            changed = resize_pv(module=module, device=device) or changed
                        if reset_pv_uuid:
                            changed = reset_uuid_pv(module=module, device=device) or changed

            if devs_to_add or devs_to_remove:
                if module.check_mode:
                    changed = True
                else:
                    if devs_to_add:
                        # create PV
                        pvcreate_cmd = module.get_bin_path("pvcreate", True)
                        for current_dev in devs_to_add:
                            rc, dummy, err = module.run_command([pvcreate_cmd] + pvoptions + ["-f", str(current_dev)])
                            if rc == 0:
                                changed = True
                            else:
                                module.fail_json(msg=f"Creating physical volume '{current_dev}' failed", rc=rc, err=err)
                        # add PV to our VG
                        vgextend_cmd = module.get_bin_path("vgextend", True)
                        rc, dummy, err = module.run_command([vgextend_cmd, vg] + devs_to_add)
                        if rc == 0:
                            changed = True
                        else:
                            module.fail_json(msg=f"Unable to extend {vg} by {' '.join(devs_to_add)}.", rc=rc, err=err)

                    # remove some PV from our VG
                    if devs_to_remove:
                        vgreduce_cmd = module.get_bin_path("vgreduce", True)
                        rc, dummy, err = module.run_command([vgreduce_cmd, "--force", vg] + devs_to_remove)
                        if rc == 0:
                            changed = True
                        else:
                            module.fail_json(
                                msg=f"Unable to reduce {vg} by {' '.join(devs_to_remove)}.", rc=rc, err=err
                            )

    module.exit_json(changed=changed)


if __name__ == "__main__":
    main()