from __future__ import annotations

import argparse
from pathlib import Path


def _replace_tag(text: str, name: str, tag: str) -> str:
    marker = f'name: {name}'
    parts = text.split(marker)
    if len(parts) == 1:
        raise ValueError(f"image marker not found for {name}")
    rebuilt = [parts[0]]
    for chunk in parts[1:]:
        lines = chunk.splitlines()
        replaced = False
        for index, line in enumerate(lines):
            if line.strip().startswith("newTag:"):
                prefix = line.split("newTag:", 1)[0]
                lines[index] = f"{prefix}newTag: {tag}"
                replaced = True
                break
        if not replaced:
            raise ValueError(f"newTag not found for {name}")
        rebuilt.append(marker + "\n".join(lines))
    return "".join(rebuilt)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--iac-root", required=True)
    parser.add_argument("--tag", required=True)
    args = parser.parse_args()

    kustomization = (
        Path(args.iac_root)
        / "services"
        / "uforia"
        / "prod"
        / "kustomization.yaml"
    )
    text = kustomization.read_text()
    for image_name in (
        "europe-north1-docker.pkg.dev/polynomial-uforia/uforia/web",
        "europe-north1-docker.pkg.dev/polynomial-uforia/uforia/api",
        "europe-north1-docker.pkg.dev/polynomial-uforia/uforia/python",
    ):
        text = _replace_tag(text, image_name, args.tag)
    kustomization.write_text(text)


if __name__ == "__main__":
    main()
