#!/usr/bin/env python3
# Companion to scripts/sync-host-keys.sh. Applies a set of additive edits
# to .sops.yaml via targeted line insertion -- not a full YAML parse and
# re-serialize -- so every untouched byte of the file is guaranteed to stay
# exactly as it was (a full round-trip through a generic YAML library was
# tested and silently reformatted the whole file's indentation style).
#
# Reads a JSON edit plan from stdin:
#   {
#     "add_keys": [{"host": "lxc-server", "age_key": "age1..."}],
#     "add_aliases": [{"host": "lxc-server", "basename": "common.yaml"}],
#     "remove_keys": ["lxc-server"],
#     "remove_aliases_for_hosts": ["lxc-server"]
#   }
# Idempotent: an anchor or alias that's already present/absent is left
# alone -- adding what's already there, or removing what's already gone,
# is a no-op rather than an error.
# Prints a JSON summary of what actually changed to stdout, so the caller
# knows which secrets/*.yaml files need `sops updatekeys` and which don't.
import json
import re
import sys

KEY_LINE_RE = re.compile(r"^  - &\S+ age1")


def add_keys(lines, add_keys):
    existing_anchors = set()
    for line in lines:
        m = re.match(r"^  - &(\S+) age1", line)
        if m:
            existing_anchors.add(m.group(1))

    new = [e for e in add_keys if e["host"] not in existing_anchors]
    if not new:
        return lines, []

    last_idx = None
    for i, line in enumerate(lines):
        if KEY_LINE_RE.match(line):
            last_idx = i
    if last_idx is None:
        print("ERROR: no '  - &<name> age1...' line found under keys: in .sops.yaml", file=sys.stderr)
        sys.exit(1)

    insert = [f"  - &{e['host']} {e['age_key']}\n" for e in new]
    lines = lines[: last_idx + 1] + insert + lines[last_idx + 1 :]
    return lines, [e["host"] for e in new]


def add_aliases(lines, add_aliases):
    changed_files = []
    for e in add_aliases:
        host = e["host"]
        basename = e["basename"]
        target = "path_regex: secrets/" + basename.replace(".", r"\.") + "$"

        rule_start = None
        for i, line in enumerate(lines):
            if target in line:
                rule_start = i
                break
        if rule_start is None:
            print(
                f"WARNING: no creation_rule for secrets/{basename} in .sops.yaml "
                f"(needed by {host}) -- add one manually, then re-run.",
                file=sys.stderr,
            )
            continue

        rule_end = len(lines)
        for i in range(rule_start + 1, len(lines)):
            if re.match(r"^  - path_regex:", lines[i]):
                rule_end = i
                break

        block = lines[rule_start:rule_end]
        if any(re.search(r"\*" + re.escape(host) + r"\s*$", line) for line in block):
            continue  # already present

        last_alias_idx = None
        last_alias_indent = None
        for i in range(rule_start, rule_end):
            m = re.match(r"^(\s*)- \*\S+\s*$", lines[i])
            if m:
                last_alias_idx = i
                last_alias_indent = m.group(1)
        if last_alias_idx is None:
            print(
                f"WARNING: creation_rule for secrets/{basename} has no existing "
                f"'- *alias' line to anchor the insertion point -- add {host} manually.",
                file=sys.stderr,
            )
            continue

        new_line = f"{last_alias_indent}- *{host}\n"
        lines = lines[: last_alias_idx + 1] + [new_line] + lines[last_alias_idx + 1 :]
        changed_files.append(basename)
    return lines, changed_files


def remove_keys(lines, hosts):
    hosts = set(hosts)
    removed = []
    kept = []
    for line in lines:
        m = re.match(r"^  - &(\S+) age1", line)
        if m and m.group(1) in hosts:
            removed.append(m.group(1))
            continue
        kept.append(line)
    return kept, removed


def remove_aliases_for_hosts(lines, hosts):
    hosts = set(hosts)
    if not hosts:
        return lines, []

    changed_files = []
    result = []
    current_basename = None
    current_block_changed = False

    def flush():
        if current_block_changed and current_basename:
            changed_files.append(current_basename)

    for line in lines:
        if re.match(r"^  - path_regex:", line):
            flush()
            current_block_changed = False
            m = re.search(r"path_regex: secrets/(.+)\$", line)
            current_basename = m.group(1).replace(r"\.", ".") if m else None
            result.append(line)
            continue

        if current_basename is not None and any(
            re.search(r"\*" + re.escape(h) + r"\s*$", line) for h in hosts
        ):
            current_block_changed = True
            continue  # drop this alias line

        result.append(line)

    flush()
    return result, changed_files


def main():
    sops_path = sys.argv[1]
    edits = json.load(sys.stdin)

    with open(sops_path) as f:
        lines = f.readlines()

    lines, added_keys = add_keys(lines, edits.get("add_keys", []))
    lines, added_alias_files = add_aliases(lines, edits.get("add_aliases", []))
    lines, removed_keys = remove_keys(lines, edits.get("remove_keys", []))
    lines, removed_alias_files = remove_aliases_for_hosts(
        lines, edits.get("remove_aliases_for_hosts", [])
    )

    with open(sops_path, "w") as f:
        f.writelines(lines)

    changed_files = sorted(set(added_alias_files) | set(removed_alias_files))
    json.dump(
        {
            "added_keys": added_keys,
            "removed_keys": removed_keys,
            "changed_secrets_files": changed_files,
        },
        sys.stdout,
    )
    print()


if __name__ == "__main__":
    main()