#!/usr/bin/env python3
"""Filter an already-secret-resolved keepalived configuration without printing it."""

import argparse
from pathlib import Path

MAP = {
    "front": ("chk_front_vip", "FRONT_VIP"),
    "mqtt": ("chk_mqtt_vip", "MQTT_VIP"),
    "db-proxy": ("chk_db_proxy_vip", "DB_PROXY_VIP"),
}


def blocks(lines):
    result = []
    index = 0
    while index < len(lines):
        stripped = lines[index].strip()
        if not stripped or stripped.startswith("#"):
            result.append(("raw", "", [lines[index]]))
            index += 1
            continue
        if "{" not in stripped:
            result.append(("raw", "", [lines[index]]))
            index += 1
            continue
        head = stripped.split("{", 1)[0].strip().split()
        kind = head[0]
        name = head[1] if len(head) > 1 else ""
        depth = 0
        body = []
        while index < len(lines):
            line = lines[index]
            body.append(line)
            depth += line.count("{") - line.count("}")
            index += 1
            if depth == 0:
                break
        if depth != 0:
            raise SystemExit("unbalanced keepalived configuration")
        result.append((kind, name, body))
    return result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--enabled", required=True)
    args = parser.parse_args()

    enabled = [item for item in args.enabled.split(",") if item]
    if not enabled or any(item not in MAP for item in enabled):
        raise SystemExit("invalid enabled VIP set")
    wanted_scripts = {MAP[item][0] for item in enabled}
    wanted_instances = {MAP[item][1] for item in enabled}

    source = Path(args.input)
    target = Path(args.output)
    rendered = []
    seen_global = False
    seen_scripts = set()
    seen_instances = set()
    for kind, name, body in blocks(source.read_text(encoding="utf-8").splitlines(True)):
        keep = kind in {"raw", "global_defs"}
        if kind == "global_defs":
            seen_global = True
        if kind == "vrrp_script" and name in wanted_scripts:
            keep = True
            seen_scripts.add(name)
        if kind == "vrrp_instance" and name in wanted_instances:
            keep = True
            seen_instances.add(name)
        if keep:
            rendered.extend(body)

    if not seen_global or seen_scripts != wanted_scripts or seen_instances != wanted_instances:
        raise SystemExit("required keepalived blocks are absent")
    target.write_text("".join(rendered), encoding="utf-8")


if __name__ == "__main__":
    main()