#!/usr/bin/python
# -*- coding: utf-8 -*-
# GNU General Public License v3.0+
from __future__ import absolute_import, division, print_function


__metaclass__ = type

DOCUMENTATION = r"""
module: vyos_ha
short_description: Manage VRRP and load balancer configuration on VyOS via REST API
description:
- Manages VRRP groups, global VRRP parameters, sync-groups, virtual servers, and LVS real servers on VyOS devices via the
  REST API.
- Uses REST API (C(connection=httpapi)) instead of CLI.
- Targets VyOS 1.4+.
version_added: 1.0.0
author:
- Evgeny Molotkov (@omnom62)
options:
  config:
    description: High-availability configuration.
    type: dict
    suboptions:
      disable:
        description: Disable all high-availability configuration.
        type: bool
        default: false
      virtual_servers:
        description: List of load balancer virtual server definitions.
        type: list
        elements: dict
        suboptions:
          name:
            type: str
            required: true
            description: Name.
          address:
            type: str
            description: Address.
          algorithm:
            type: str
            description: Algorithm.
          delay_loop:
            type: int
            description: Delay loop.
          forward_method:
            type: str
            choices:
            - direct
            - nat
            description: Forward method.
          fwmark:
            type: int
            description: Fwmark.
          persistence_timeout:
            type: int
            description: Persistence timeout.
          port:
            type: int
            description: Port.
          protocol:
            type: str
            choices:
            - tcp
            - udp
            description: Protocol.
          real_server:
            type: list
            elements: dict
            suboptions:
              address:
                type: str
                required: true
                description: Address.
              port:
                type: int
                description: Port.
              connection_timeout:
                type: int
                description: Connection timeout.
              health_check_script:
                type: str
                description: Health check script.
            description: Real server.
      vrrp:
        description: VRRP configuration.
        type: dict
        suboptions:
          global_parameters:
            type: dict
            suboptions:
              garp:
                type: dict
                suboptions:
                  interval:
                    type: int
                    description: Interval.
                  master_delay:
                    type: int
                    description: Master delay.
                  master_refresh:
                    type: int
                    description: Master refresh.
                  master_refresh_repeat:
                    type: int
                    description: Master refresh repeat.
                  master_repeat:
                    type: int
                    description: Master repeat.
                description: Garp.
              startup_delay:
                type: int
                description: Startup delay.
              version:
                type: str
                description: Version.
            description: Global parameters.
          groups:
            type: list
            elements: dict
            suboptions:
              name:
                type: str
                required: true
                description: Name.
              address:
                type: list
                elements: str
                description: Address.
              advertise_interval:
                type: int
                description: Advertise interval.
              authentication:
                type: dict
                suboptions:
                  password:
                    type: str
                    description: Password.
                  type:
                    type: str
                    description: Type.
                description: Authentication.
              description:
                type: str
                description: Description.
              disable:
                type: bool
                default: false
                description: Disable.
              excluded_address:
                type: list
                elements: str
                description: Excluded address.
              garp:
                type: dict
                suboptions:
                  interval:
                    type: int
                    description: Interval.
                  master_delay:
                    type: int
                    description: Master delay.
                  master_refresh:
                    type: int
                    description: Master refresh.
                  master_refresh_repeat:
                    type: int
                    description: Master refresh repeat.
                  master_repeat:
                    type: int
                    description: Master repeat.
                description: Garp.
              health_check:
                type: dict
                suboptions:
                  failure_count:
                    type: int
                    description: Failure count.
                  interval:
                    type: int
                    description: Interval.
                  ping:
                    type: str
                    description: Ping.
                  script:
                    type: str
                    description: Script.
                description: Health check.
              hello_source_address:
                type: str
                description: Hello source address.
              interface:
                type: str
                description: Interface.
              no_preempt:
                type: bool
                default: false
                description: No preempt.
              peer_address:
                type: str
                description: Peer address.
              preempt_delay:
                type: int
                description: Preempt delay.
              priority:
                type: int
                description: Priority.
              rfc3768_compatibility:
                type: bool
                default: false
                description: Rfc3768 compatibility.
              track:
                type: dict
                suboptions:
                  exclude_vrrp_interface:
                    type: bool
                    description: Exclude vrrp interface.
                  interface:
                    type: list
                    elements: str
                    description: Interface.
                description: Track.
              transition_script:
                type: dict
                suboptions:
                  backup:
                    type: str
                    description: Backup.
                  fault:
                    type: str
                    description: Fault.
                  master:
                    type: str
                    description: Master.
                  stop:
                    type: str
                    description: Stop.
                description: Transition script.
              vrid:
                type: int
                description: Vrid.
            description: Groups.
          snmp:
            type: str
            choices:
            - enabled
            - disabled
            description: Snmp.
          sync_groups:
            type: list
            elements: dict
            suboptions:
              name:
                type: str
                required: true
                description: Name.
              health_check:
                type: dict
                suboptions:
                  failure_count:
                    type: int
                    description: Failure count.
                  interval:
                    type: int
                    description: Interval.
                  ping:
                    type: str
                    description: Ping.
                  script:
                    type: str
                    description: Script.
                description: Health check.
              member:
                type: list
                elements: str
                description: Member.
              transition_script:
                type: dict
                suboptions:
                  backup:
                    type: str
                    description: Backup.
                  fault:
                    type: str
                    description: Fault.
                  master:
                    type: str
                    description: Master.
                  stop:
                    type: str
                    description: Stop.
                description: Transition script.
            description: Sync groups.
  state:
    description: Desired end state of the configuration.
    type: str
    choices:
    - merged
    - replaced
    - overridden
    - deleted
    - gathered
    default: merged

"""

EXAMPLES = r"""
- name: Merge VRRP configuration
  vyos.rest.vyos_ha:
    config:
      vrrp:
        global_parameters:
          startup_delay: 30
        groups:
          - name: g1
            interface: eth0
            vrid: 20
            priority: 100
            address:
              - 192.168.1.100/24
        sync_groups:
          - name: sg1
            member: [g1]
        snmp: enabled
    state: merged

- name: Delete all HA configuration
  vyos.rest.vyos_ha:
    state: deleted

- name: Gather current HA configuration
  vyos.rest.vyos_ha:
    state: gathered
"""

RETURN = r"""
before:
  description: HA configuration before this module ran.
  returned: always
  type: dict
after:
  description: HA configuration after this module ran.
  returned: when changed
  type: dict
commands:
  description: List of API commands sent to the device.
  returned: always
  type: list
gathered:
  description: Current HA configuration as structured data.
  returned: when state is gathered
  type: dict
saved:
  description: Whether the config was saved after changes.
  returned: when changed
  type: bool
"""

from ansible.module_utils.basic import AnsibleModule
from ansible_collections.vyos.rest.plugins.module_utils.vyos import (
    VyOSModule,
    dict_op,
)


_BASE = ["high-availability"]

ARGUMENT_SPEC = dict(
    config=dict(
        type="dict",
        options=dict(
            disable=dict(type="bool", default=False),
            virtual_servers=dict(
                type="list",
                elements="dict",
                options=dict(
                    name=dict(type="str", required=True),
                    address=dict(type="str"),
                    algorithm=dict(type="str"),
                    delay_loop=dict(type="int"),
                    forward_method=dict(type="str", choices=["direct", "nat"]),
                    fwmark=dict(type="int"),
                    persistence_timeout=dict(type="int"),
                    port=dict(type="int"),
                    protocol=dict(type="str", choices=["tcp", "udp"]),
                    real_server=dict(
                        type="list",
                        elements="dict",
                        options=dict(
                            address=dict(type="str", required=True),
                            port=dict(type="int"),
                            connection_timeout=dict(type="int"),
                            health_check_script=dict(type="str"),
                        ),
                    ),
                ),
            ),
            vrrp=dict(
                type="dict",
                options=dict(
                    global_parameters=dict(
                        type="dict",
                        options=dict(
                            garp=dict(
                                type="dict",
                                options=dict(
                                    interval=dict(type="int"),
                                    master_delay=dict(type="int"),
                                    master_refresh=dict(type="int"),
                                    master_refresh_repeat=dict(type="int"),
                                    master_repeat=dict(type="int"),
                                ),
                            ),
                            startup_delay=dict(type="int"),
                            version=dict(type="str"),
                        ),
                    ),
                    groups=dict(
                        type="list",
                        elements="dict",
                        options=dict(
                            name=dict(type="str", required=True),
                            address=dict(type="list", elements="str"),
                            advertise_interval=dict(type="int"),
                            authentication=dict(
                                type="dict",
                                options=dict(
                                    password=dict(type="str", no_log=True),
                                    type=dict(type="str"),
                                ),
                            ),
                            description=dict(type="str"),
                            disable=dict(type="bool", default=False),
                            excluded_address=dict(type="list", elements="str"),
                            garp=dict(
                                type="dict",
                                options=dict(
                                    interval=dict(type="int"),
                                    master_delay=dict(type="int"),
                                    master_refresh=dict(type="int"),
                                    master_refresh_repeat=dict(type="int"),
                                    master_repeat=dict(type="int"),
                                ),
                            ),
                            health_check=dict(
                                type="dict",
                                options=dict(
                                    failure_count=dict(type="int"),
                                    interval=dict(type="int"),
                                    ping=dict(type="str"),
                                    script=dict(type="str"),
                                ),
                            ),
                            hello_source_address=dict(type="str"),
                            interface=dict(type="str"),
                            no_preempt=dict(type="bool", default=False),
                            peer_address=dict(type="str"),
                            preempt_delay=dict(type="int"),
                            priority=dict(type="int"),
                            rfc3768_compatibility=dict(type="bool", default=False),
                            track=dict(
                                type="dict",
                                options=dict(
                                    exclude_vrrp_interface=dict(type="bool"),
                                    interface=dict(type="list", elements="str"),
                                ),
                            ),
                            transition_script=dict(
                                type="dict",
                                options=dict(
                                    backup=dict(type="str"),
                                    fault=dict(type="str"),
                                    master=dict(type="str"),
                                    stop=dict(type="str"),
                                ),
                            ),
                            vrid=dict(type="int"),
                        ),
                    ),
                    snmp=dict(type="str", choices=["enabled", "disabled"]),
                    sync_groups=dict(
                        type="list",
                        elements="dict",
                        options=dict(
                            name=dict(type="str", required=True),
                            health_check=dict(
                                type="dict",
                                options=dict(
                                    failure_count=dict(type="int"),
                                    interval=dict(type="int"),
                                    ping=dict(type="str"),
                                    script=dict(type="str"),
                                ),
                            ),
                            member=dict(type="list", elements="str"),
                            transition_script=dict(
                                type="dict",
                                options=dict(
                                    backup=dict(type="str"),
                                    fault=dict(type="str"),
                                    master=dict(type="str"),
                                    stop=dict(type="str"),
                                ),
                            ),
                        ),
                    ),
                ),
            ),
        ),
    ),
    state=dict(
        default="merged",
        choices=["merged", "replaced", "overridden", "deleted", "gathered"],
    ),
)


# ---------------------------------------------------------------------------
# Generic converters — zero field-name knowledge
# ---------------------------------------------------------------------------


def _try_int(v):
    """Try casting v to int; return original on failure."""
    try:
        return int(v)
    except (ValueError, TypeError):
        return v


def _to_dict(value):
    """Coerce VyOS tag-node value (str/list/dict/None) to {key: {}} dict."""
    if not value:
        return {}
    if isinstance(value, dict):
        return value
    if isinstance(value, str):
        return {value: {}}
    if isinstance(value, list):
        return {str(v): {} for v in value}
    return {}


def _to_device_generic(d, tag_node_fields=None):
    """Convert argspec snake_case dict to device kebab-case shape.

    Rules (no field-name knowledge needed):
    - snake_case key -> kebab-case key
    - bool True -> presence node {}; bool False -> skip
    - int/str -> str(v)
    - None -> skip
    - dict -> recurse
    - list whose argspec key is in tag_node_fields -> {v: {}} tag-node dict
    - list of dicts -> caller handles (named-object lists)
    """
    tag_node_fields = tag_node_fields or set()
    result = {}
    for k, v in (d or {}).items():
        if v is None:
            continue
        device_key = k.replace("_", "-")
        if isinstance(v, bool):
            if v:
                result[device_key] = {}
        elif isinstance(v, (int, str)):
            result[device_key] = str(v)
        elif isinstance(v, dict):
            sub = _to_device_generic(v, tag_node_fields)
            if sub:
                result[device_key] = sub
        elif isinstance(v, list) and k in tag_node_fields:
            result[device_key] = {str(item): {} for item in v}
    return result


def _from_device_generic(raw):
    """Convert device kebab-case dict to argspec snake_case shape.

    Rules (no field-name knowledge needed):
    - kebab-case key -> snake_case key
    - Empty dict {} -> True (presence node = bool True)
    - Numeric strings -> int via _try_int
    - Nested non-empty dicts -> recurse
    - Everything else -> as-is
    """
    if not raw or not isinstance(raw, dict):
        return {}
    result = {}
    for k, v in raw.items():
        arg_key = k.replace("-", "_")
        if isinstance(v, dict) and not v:
            result[arg_key] = True
        elif isinstance(v, dict):
            result[arg_key] = _from_device_generic(v)
        elif isinstance(v, str):
            result[arg_key] = _try_int(v)
        else:
            result[arg_key] = v
    return result


# ---------------------------------------------------------------------------
# Structural adapters
#
# The ONLY reasons these exist — not field names, but structural mismatches:
# 1. Named-object lists (groups, sync_groups, virtual_servers, real_server)
#    argspec uses [{name: "x", ...}], device uses {"x": {...}}
# 2. Tag-node lists (address, excluded_address, track.interface, member)
#    argspec uses ["a", "b"], device uses {"a": {}, "b": {}}
# 3. snmp: argspec "enabled"/"disabled" string <-> device presence node / absent
# 4. health_check_script: argspec flat field <-> device nested under health-check.script
# ---------------------------------------------------------------------------


def _want_to_device(config):
    """Convert argspec config to device shape.

    Only handles the four structural exceptions above.
    All field conversion delegated to _to_device_generic.
    """
    if not config:
        return {}
    want = {}

    # Top-level scalars and bools via generic converter
    top = _to_device_generic(
        {k: v for k, v in config.items() if k not in ("virtual_servers", "vrrp")},
    )
    want.update(top)

    # virtual-server: named-object list -> dict keyed by name
    for vs in config.get("virtual_servers") or []:
        vs_dev = _to_device_generic(
            {k: v for k, v in vs.items() if k not in ("name", "real_server")},
        )
        for rs in vs.get("real_server") or []:
            rs_dev = _to_device_generic(
                {k: v for k, v in rs.items() if k not in ("address", "health_check_script")},
            )
            # health_check_script -> nested under health-check.script
            if rs.get("health_check_script"):
                rs_dev["health-check"] = {"script": rs["health_check_script"]}
            vs_dev.setdefault("real-server", {})[rs["address"]] = rs_dev
        want.setdefault("virtual-server", {})[vs["name"]] = vs_dev

    # vrrp subtree
    vrrp = config.get("vrrp") or {}
    if vrrp:
        vrrp_dev = _to_device_generic(
            {k: v for k, v in vrrp.items() if k not in ("groups", "sync_groups", "snmp")},
        )

        # snmp: "enabled" -> presence node; "disabled" -> handled via delete in main()
        if vrrp.get("snmp") == "enabled":
            vrrp_dev["snmp"] = {}

        # group: named-object list -> dict keyed by name
        # address, excluded_address, track.interface are tag-node lists
        for grp in vrrp.get("groups") or []:
            grp_dev = _to_device_generic(
                {
                    k: v
                    for k, v in grp.items()
                    if k not in ("name", "address", "excluded_address", "track")
                },
                tag_node_fields=set(),
            )
            if grp.get("address"):
                grp_dev["address"] = {a: {} for a in grp["address"]}
            if grp.get("excluded_address"):
                grp_dev["excluded-address"] = {a: {} for a in grp["excluded_address"]}
            if grp.get("track"):
                t = grp["track"]
                track_dev = _to_device_generic(
                    {k: v for k, v in t.items() if k != "interface"},
                )
                if t.get("interface"):
                    track_dev["interface"] = {i: {} for i in t["interface"]}
                if track_dev:
                    grp_dev["track"] = track_dev
            vrrp_dev.setdefault("group", {})[grp["name"]] = grp_dev

        # sync-group: named-object list -> dict keyed by name
        # member is a tag-node list
        for sg in vrrp.get("sync_groups") or []:
            sg_dev = _to_device_generic(
                {k: v for k, v in sg.items() if k not in ("name", "member")},
            )
            if sg.get("member"):
                sg_dev["member"] = {m: {} for m in sg["member"]}
            vrrp_dev.setdefault("sync-group", {})[sg["name"]] = sg_dev

        if vrrp_dev:
            want["vrrp"] = vrrp_dev

    return want


def _device_to_argspec(raw):
    """Convert raw device response to argspec shape.

    Only handles the four structural exceptions above.
    All field conversion delegated to _from_device_generic.
    """
    if not raw:
        return {}

    # Start with generic conversion of everything
    result = _from_device_generic(
        {k: v for k, v in raw.items() if k not in ("virtual-server", "vrrp")},
    )

    # virtual-server: dict keyed by name -> named-object list
    vs_raw = raw.get("virtual-server") or {}
    if vs_raw and isinstance(vs_raw, dict):
        vs_list = []
        for name, vs_data in sorted(vs_raw.items()):
            vs_data = vs_data or {}
            rs_raw = vs_data.get("real-server") or {}
            vs = {"name": name}
            vs.update(
                _from_device_generic(
                    {k: v for k, v in vs_data.items() if k != "real-server"},
                ),
            )
            # real-server: dict keyed by address -> named-object list
            if rs_raw and isinstance(rs_raw, dict):
                rs_list = []
                for addr, rs_data in sorted(rs_raw.items()):
                    rs_data = rs_data or {}
                    hc = rs_data.get("health-check") or {}
                    rs = {"address": addr}
                    rs.update(
                        _from_device_generic(
                            {k: v for k, v in rs_data.items() if k != "health-check"},
                        ),
                    )
                    # health-check.script -> health_check_script flat field
                    if hc.get("script"):
                        rs["health_check_script"] = hc["script"]
                    rs_list.append(rs)
                if rs_list:
                    vs["real_server"] = rs_list
            vs_list.append(vs)
        if vs_list:
            result["virtual_servers"] = vs_list

    # vrrp subtree
    vrrp_raw = raw.get("vrrp") or {}
    if vrrp_raw:
        vrrp_arg = _from_device_generic(
            {k: v for k, v in vrrp_raw.items() if k not in ("group", "sync-group", "snmp")},
        )

        # snmp: presence node -> "enabled" string
        if "snmp" in vrrp_raw:
            vrrp_arg["snmp"] = "enabled"

        # group: dict keyed by name -> named-object list
        # address, excluded-address, track.interface are tag-node dicts -> lists
        grp_raw = vrrp_raw.get("group") or {}
        if grp_raw and isinstance(grp_raw, dict):
            grps = []
            for name, gdata in sorted(grp_raw.items()):
                gdata = gdata or {}
                track_raw = gdata.get("track") or {}
                grp = {"name": name}
                grp.update(
                    _from_device_generic(
                        {
                            k: v
                            for k, v in gdata.items()
                            if k not in ("track", "address", "excluded-address")
                        },
                    ),
                )
                # address and excluded-address: tag-node dict -> sorted list
                if "address" in gdata:
                    grp["address"] = sorted(_to_dict(gdata["address"]).keys())
                if "excluded-address" in gdata:
                    grp["excluded_address"] = sorted(
                        _to_dict(gdata["excluded-address"]).keys(),
                    )
                # track: convert recursively, interface is tag-node dict -> list
                if track_raw:
                    track = _from_device_generic(
                        {k: v for k, v in track_raw.items() if k != "interface"},
                    )
                    if "interface" in track_raw:
                        track["interface"] = sorted(
                            _to_dict(track_raw["interface"]).keys(),
                        )
                    if track:
                        grp["track"] = track
                grps.append(grp)
            if grps:
                vrrp_arg["groups"] = grps

        # sync-group: dict keyed by name -> named-object list
        # member: string or tag-node dict -> list
        sg_raw = vrrp_raw.get("sync-group") or {}
        if sg_raw and isinstance(sg_raw, dict):
            sgs = []
            for name, sdata in sorted(sg_raw.items()):
                sdata = sdata or {}
                member_raw = sdata.get("member")
                sg = {"name": name}
                sg.update(
                    _from_device_generic(
                        {k: v for k, v in sdata.items() if k != "member"},
                    ),
                )
                # member: string (single) or dict (multiple) -> list
                if member_raw is not None:
                    sg["member"] = (
                        sorted(_to_dict(member_raw).keys())
                        if isinstance(member_raw, dict)
                        else [member_raw]
                    )
                sgs.append(sg)
            if sgs:
                vrrp_arg["sync_groups"] = sgs

        if vrrp_arg:
            result["vrrp"] = vrrp_arg

    return result


# ---------------------------------------------------------------------------
# Normalize have for dict_op
# ---------------------------------------------------------------------------


def _normalize_have(raw):
    """Normalize raw device response so dict_op only sees dicts.

    Coerces tag-node values that the device returns as strings (single entry)
    to {key: {}} dicts. Only applied to known tag-node contexts.
    """
    if not raw or not isinstance(raw, dict):
        return {}
    result = {}
    for k, v in raw.items():
        if isinstance(v, dict):
            result[k] = _normalize_have(v)
        elif isinstance(v, str) and k == "member":
            result[k] = {v: {}}
        else:
            result[k] = v
    return result


def _normalize_groups(norm_have):
    """Coerce VRRP group tag-nodes from string to dict in-place."""
    for gdata in (norm_have.get("vrrp", {}).get("group", {}) or {}).values():
        if not isinstance(gdata, dict):
            continue
        for tag_key in ("address", "excluded-address"):
            if isinstance(gdata.get(tag_key), str):
                gdata[tag_key] = {gdata[tag_key]: {}}
        track = gdata.get("track") or {}
        if isinstance(track.get("interface"), str):
            track["interface"] = {track["interface"]: {}}


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main():
    module = AnsibleModule(ARGUMENT_SPEC, supports_check_mode=True)
    vyos = VyOSModule(module)

    state = module.params["state"]
    config = module.params.get("config") or {}

    raw_have = vyos.get_config(_BASE)
    have = _device_to_argspec(raw_have)

    if state == "gathered":
        module.exit_json(changed=False, gathered=have)

    want_device = _want_to_device(config)
    norm_have = _normalize_have(raw_have)
    _normalize_groups(norm_have)

    if state == "deleted":
        commands = [("delete", _BASE)] if raw_have else []
    elif state == "overridden":
        commands = []
        for section in list(norm_have.keys()):
            if section not in want_device:
                commands.append(("delete", _BASE + [section]))
            else:
                commands += dict_op(
                    want_device[section],
                    norm_have[section],
                    _BASE + [section],
                    op="purge",
                )
        commands += dict_op(want_device, norm_have, _BASE, op="set")
    elif state == "replaced":
        commands = []
        for section, section_want in want_device.items():
            section_have = norm_have.get(section, {})
            commands += dict_op(section_want, section_have, _BASE + [section], op="purge")
        commands += dict_op(want_device, norm_have, _BASE, op="set")
    else:
        commands = dict_op(want_device, norm_have, _BASE, op="set")

    # snmp: disabled -> explicit delete (presence node, no "disabled" value on device)
    if (config.get("vrrp") or {}).get("snmp") == "disabled":
        if "snmp" in (raw_have.get("vrrp") or {}):
            commands.append(("delete", _BASE + ["vrrp", "snmp"]))

    if module.check_mode:
        module.exit_json(changed=bool(commands), commands=commands, before=have)

    if commands:
        response = vyos.apply_commands(commands)
        saved = vyos.save_config()
        after = _device_to_argspec(vyos.get_config(_BASE))
        module.exit_json(
            changed=True,
            before=have,
            after=after,
            commands=commands,
            saved=saved,
            response=response,
        )

    module.exit_json(changed=False, before=have, after=have, commands=[])


if __name__ == "__main__":
    main()