#!/usr/bin/env python3

import re
import sys


def _quote_yaml_single(s: str) -> str:
    # YAML single-quoted scalars escape a single quote by doubling it.
    return "'" + s.replace("'", "''") + "'"


def patch_realtime_deployment_yaml(content: str) -> tuple[str, bool]:
    changed = False

    # Case A: kompose sometimes emits a single-quoted curl command as a single argv entry.
    # Wrap it into `sh -c '<curl ...>'`.
    def _wrap_single_curl(m: re.Match[str]) -> str:
        nonlocal changed
        changed = True
        cmd_hdr = m.group(1)
        indent = m.group(2)
        curl_cmd = m.group(3)
        return (
            cmd_hdr
            + indent
            + "sh\n"
            + indent
            + "-c\n"
            + indent
            + _quote_yaml_single(curl_cmd)
            + "\n"
        )

    content = re.sub(
        r"(\s+command:\n)(\s+- )'(curl[^\n]+)'\n",
        _wrap_single_curl,
        content,
        count=1,
    )

    # Case B: unquoted scalars containing `Authorization: ...` can be parsed as YAML mappings.
    # Quote them to guarantee Kubernetes receives a string.
    out_lines: list[str] = []
    for line in content.splitlines(True):
        m = re.match(r"^(\s*-\s+)(curl\b.*)$", line)
        if m:
            prefix = m.group(1)
            cmd = m.group(2).rstrip("\n")
            if "Authorization:" in cmd and not cmd.lstrip().startswith(("'", '"')):
                out_lines.append(prefix + _quote_yaml_single(cmd) + "\n")
                changed = True
                continue
        out_lines.append(line)

    return "".join(out_lines), changed


def main() -> int:
    if len(sys.argv) != 2:
        print("Usage: patch_realtime_probe.py <path-to-realtime-deployment.yaml>", file=sys.stderr)
        return 2

    path = sys.argv[1]
    with open(path, "r", encoding="utf-8") as f:
        content = f.read()

    patched, changed = patch_realtime_deployment_yaml(content)
    if changed:
        with open(path, "w", encoding="utf-8") as f:
            f.write(patched)
        print("  [OK] realtime liveness probe YAML sanitized")
    else:
        print("  [OK] realtime liveness probe already sane")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())